feat: Improved the Architecture Modification Agent

Author: lak7

CHANGELOG.md

@@ -1,5 +1,18 @@
 # Changelog
 
+## v0.5.1
+### Added
+- Structured architecture update flow with a user-facing update action
+- Tool-driven agent flow for architecture modification
+- Token-usage based credit accounting for architecture-related AI operations
+
+### Improved
+- Architecture modification dialog and decision logic for clearer tool selection
+- More concise and deterministic agent responses
+- Chat interface now includes per-message inline options for quicker actions
+- Responses include termination metadata
+- Support for partial-result recovery when errors occur
+
 ## v0.5.0
 ### Added
 - Tool-driven chat agent with structured interview flow and clearer multi-stage responses

README.md

@@ -9,7 +9,7 @@
 </p>
 
 <p align="center">
-  <img src="https://img.shields.io/badge/version-0.5.0-blue" alt="Version 0.5.0" />
+  <img src="https://img.shields.io/badge/version-0.5.1-blue" alt="Version 0.5.1" />
    <img src="https://img.shields.io/badge/status-experimental-orange" alt="Status: Experimental" />
    <a href="https://opensource.org/licenses/Apache-2.0">
     <img src="https://img.shields.io/badge/License-Apache%202.0-blue.svg" alt="License: Apache 2.0" />

actions/agentsFlow.ts

@@ -1,17 +1,19 @@
 "use server";
 import { ChatOpenAI } from "@langchain/openai";
-import { ChatPromptTemplate, HumanMessagePromptTemplate, MessagesPlaceholder, PromptTemplate } from "@langchain/core/prompts";
-import { architectureModificationPrompt, BASE_DEVILDEV_AGENT_PROMPT, DevilDevAgentAfterInterviewRole, DevilDevAgentBeforeInterviewRole } from "../prompts/dev/Chatbot";
+import { ChatPromptTemplate, HumanMessagePromptTemplate, MessagesPlaceholder } from "@langchain/core/prompts";
+import { architectureModificationAgentPrompt, BASE_DEVILDEV_AGENT_PROMPT, DevilDevAgentAfterInterviewRole, DevilDevAgentBeforeInterviewRole } from "../prompts/dev/Chatbot";
 import { deductCredits, getCredits } from "./credits";
 import { minSoulsToSendMessage } from "../Limits";
-import { extractTextContent } from "@/lib/ai/extractTextContent";
 import { interviewTool } from "../tools/ptoA-tools/interview";
 import { generalResTool } from "../tools/ptoA-tools/general_res";
 import { AgentExecutor, createToolCallingAgent } from "langchain/agents";
 import { TokenUsageCallbackHandler } from "../common/TokenUsageHandler";
 import { TERMINATING_TOOLS, TerminatingTools } from "../types/pToA/tools";
 import { tier1Tool } from "../tools/ptoA-tools/tier-1";
 import { tier2Tool } from "../tools/ptoA-tools/tier-2";
+import { updateArchitectureTool } from "../tools/ptoA-tools/update_architecture";
+import { DynamicStructuredTool } from "langchain/tools";
+import { z } from "zod";
 
 // Return type for agent flow functions
 export type AgentFlowResult = {
@@ -132,49 +134,96 @@ export async function chatbot(userInput: string, conversationHistory: any[] = []
 
 export async function architectureModificationBot(userInput: string, conversationHistory: any[] = [], architectureData: any, userId: string | null = null) {
     // Check credits before running the agent
-    if (userId) {
+    if (userId) { 
         const creditsResult = await getCredits(userId);
         if (creditsResult.success && creditsResult.credits !== undefined && creditsResult.credits < minSoulsToSendMessage) {
             return { error: 'INSUFFICIENT_CREDITS', remainingCredits: creditsResult.credits };
         }
     }
 
-    const template = architectureModificationPrompt;
-
     // Format conversation history for the prompt
-    const formattedHistory = conversationHistory.map(msg => 
+    const formattedHistory = conversationHistory.map(msg =>
         `${msg.type === 'user' ? 'User' : 'Assistant'}: ${msg.content}`
     ).join('\n');
 
-    const prompt = PromptTemplate.fromTemplate(template);
-    const chain = prompt.pipe(llm);
-    const result = await chain.invoke({
-        userInput: userInput,
-        conversationHistory: formattedHistory,
-        architecture_data: JSON.stringify(architectureData)
+    const get_architecture = new DynamicStructuredTool({
+        name: "get_architecture",
+        description:
+            "Use this when you need to inspect the current architecture. It returns the full architecture JSON as a string. This does NOT close the agent.",
+        schema: z.object({}),
+        func: async (): Promise<string> => {
+            return JSON.stringify(architectureData);
+        },
     });
-    
-    // Extract text content - handle both string and complex content types
-    const textContent = extractTextContent(result.content);
-    
-    // Deduct credits if userId is provided
+
+    const tools = [get_architecture, generalResTool, interviewTool, updateArchitectureTool];
+
+    const prompt = ChatPromptTemplate.fromMessages([
+        ["system", architectureModificationAgentPrompt],
+        HumanMessagePromptTemplate.fromTemplate("{userInput}"),
+        new MessagesPlaceholder("agent_scratchpad"),
+    ]);
+
+    const agent = await createToolCallingAgent({ llm, tools, prompt });
+
+    const agentExecutor = new AgentExecutor({
+        agent,
+        tools,
+        verbose: true,
+        maxIterations: 10,
+    });
+
+    const inputs = {
+        userInput,
+        conversationHistory: formattedHistory,
+    };
+
+    const tokenUsageHandler = new TokenUsageCallbackHandler();
+
+    const stream = (agentExecutor as any)._streamIterator(inputs, { callbacks: [tokenUsageHandler] });
+
+    let result: { output?: string; terminatingTool?: TerminatingTools } | undefined;
+
+    for await (const step of stream) {
+        if (!step) continue;
+
+        if (step.output !== undefined) {
+            result = { output: step.output, terminatingTool: step.terminatingTool };
+            break;
+        }
+
+        const steps = step.intermediateSteps as Array<{ action: { tool: string }; observation: string }> | undefined;
+        if (steps?.length) {
+            const last = steps[steps.length - 1];
+            const toolName = last?.action?.tool?.toLowerCase();
+            if (toolName && TERMINATING_TOOLS.has(toolName)) {
+                result = { output: last.observation, terminatingTool: toolName as TerminatingTools };
+                break;
+            }
+        }
+    }
+
+    console.log("architectureModificationBot result:", result);
+    const tokenUsage = tokenUsageHandler.getUsage();
+
     if (userId) {
-        const inputTokens = result.usage_metadata?.input_tokens ?? 0;
-        const outputTokens = result.usage_metadata?.output_tokens ?? 0;
-        
-        const creditResult = await deductCredits(userId, inputTokens, outputTokens);
+        const inputTokens = tokenUsage.inputTokens;
+        const outputTokens = tokenUsage.outputTokens;
+
+        const creditResult = await deductCredits(userId, inputTokens, outputTokens, true);
         if (!creditResult.success) {
             console.error("Failed to deduct credits:", creditResult.error);
-            return { textContent };
+            return { response: result?.output, terminatingTool: result?.terminatingTool };
         } else {
             console.log(`Credits deducted: ${creditResult.deducted}, remaining: ${creditResult.remaining}`);
             return {
-                textContent,
+                response: result?.output,
+                terminatingTool: result?.terminatingTool,
                 remainingCredits: creditResult.remaining,
                 deducted: creditResult.deducted
             };
         }
     }
-    
-    return { textContent };
+
+    return { response: result?.output, terminatingTool: result?.terminatingTool };
 }

actions/architecture.ts

@@ -18,7 +18,7 @@ import { minSoulsToGenArch } from "../Limits";
 import { TokenUsageCallbackHandler } from "../common/TokenUsageHandler";
 import { db } from "@/lib/db";
 import { createParallelWebSearchTool } from "../tools/parallel";
-import { ARCHITECTURE_GENERATION_PROMPT } from "../prompts/dev/architecture";
+import { ARCHITECTURE_GENERATION_PROMPT, ARCHITECTURE_UPDATE_PROMPT } from "../prompts/dev/architecture";
 const { inngest } = await import('../src/inngest/client');
 
 
@@ -490,3 +490,85 @@ export async function generateMainArchitecture({
     return { success: false, error: error instanceof Error ? error.message : "Unknown error" };
   }
 }
+
+export async function updateArchitecture({
+  changeRequirement,
+  chatId,
+  stackId,
+  architecture_data,
+  userId,
+}: {
+  changeRequirement: string;
+  chatId: string;
+  stackId?: string;
+  architecture_data: any;
+  userId: string;
+}) {
+  const creditsResult = await getCredits(userId);
+  if (creditsResult.success && creditsResult.credits !== undefined && creditsResult.credits < minSoulsToGenArch) {
+    return { success: false, error: 'INSUFFICIENT_CREDITS', remainingCredits: creditsResult.credits };
+  }
+
+  const tokenUsageHandler = new TokenUsageCallbackHandler();
+
+  try {
+    const llm = new ChatOpenAI({ openAIApiKey: process.env.OPENAI_API_KEY, model: "gpt-5-mini-2025-08-07" });
+
+    const prompt = PromptTemplate.fromTemplate(ARCHITECTURE_UPDATE_PROMPT);
+
+    const structuredLlm = llm.withStructuredOutput(architectureJsonSchemaWithPrd);
+    const chain = prompt.pipe(structuredLlm);
+
+    const updatedArchitecture = await chain.invoke(
+      {
+        changeRequirement,
+        currentArchitecture: JSON.stringify(architecture_data),
+      },
+      { callbacks: [tokenUsageHandler] }
+    );
+
+    const dbOps: any[] = [
+      db.architecture.create({
+        data: {
+          chatId,
+          stackId,
+          requirement: changeRequirement,
+          architectureRationale: updatedArchitecture.architectureRationale,
+          components: updatedArchitecture.components,
+          connectionLabels: updatedArchitecture.connectionLabels || {},
+          componentPositions: {},
+        },
+      }),
+    ];
+
+    if (stackId) {
+      dbOps.push(
+        db.stack.update({
+          where: { id: stackId },
+          data: { prd: updatedArchitecture.prd },
+        })
+      );
+    }
+
+    const [newArchitecture] = await db.$transaction(dbOps);
+
+    const usage = tokenUsageHandler.getUsage();
+    const creditResult = await deductCredits(userId, usage.inputTokens, usage.outputTokens);
+    if (!creditResult.success) {
+      console.error("Failed to deduct credits:", creditResult.error);
+    } else {
+      console.log(`Credits deducted: ${creditResult.deducted}, remaining: ${creditResult.remaining}`);
+    }
+
+    return {
+      success: true,
+      architecture: newArchitecture,
+      prd: updatedArchitecture.prd,
+      creditsRemaining: creditResult.remaining,
+    };
+
+  } catch (error) {
+    console.error("Error in updateArchitecture:", error);
+    return { success: false, error: error instanceof Error ? error.message : "Unknown error" };
+  }
+}

actions/chat.ts

@@ -15,6 +15,7 @@ export interface ChatMessage {
   interviewAnswers?: InterviewAnswer[]; 
   prompt?: string;
   tier2Context?: string;
+  changeRequirement?: string;
 }
 
 // Create a new chat with a specific ID (for localStorage flow)