From 661c68e37c1e131328c091f390e7022b30007b98 Mon Sep 17 00:00:00 2001 From: Ahmed Butt Date: Sun, 12 Apr 2026 22:21:12 +0500 Subject: [PATCH] improvement: centralise model config and deduplicate message-conversion helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Created agent/common/utils.py with: - DEFAULT_MODEL constant — single place to change the Claude model version across the entire agent (was hardcoded as "claude-sonnet-4-20250514" in 5 separate places) - convert_tools_messages_to_ai_and_human() — shared implementation of the tool-message normaliser that was copy-pasted identically in both architect/graph.py and developer/graph.py - Updated architect/graph.py and developer/graph.py to import from utils.py and removed the local duplicates + now-unused json/AIMessage/HumanMessage imports in each file Co-Authored-By: Claude Sonnet 4.6 --- agent/architect/graph.py | 26 +++++------------------- agent/common/utils.py | 44 ++++++++++++++++++++++++++++++++++++++++ agent/developer/graph.py | 28 ++++++------------------- 3 files changed, 55 insertions(+), 43 deletions(-) create mode 100644 agent/common/utils.py diff --git a/agent/architect/graph.py b/agent/architect/graph.py index 1cbdf94..bf7fcdf 100644 --- a/agent/architect/graph.py +++ b/agent/architect/graph.py @@ -1,4 +1,3 @@ -import json from typing import List, TypedDict, Optional from langchain_anthropic import ChatAnthropic @@ -15,6 +14,7 @@ from agent.tools.write import get_files_structure from helpers.prompts import markdown_to_prompt_template from agent.common.entities import ImplementationPlan +from agent.common.utils import DEFAULT_MODEL, convert_tools_messages_to_ai_and_human class ResearchStep(BaseModel): @@ -33,10 +33,10 @@ class ResearchEvaluation(BaseModel): extract_implementation_prompt = markdown_to_prompt_template("agent/architect/prompts/extract_implementation_plan.md") # runnable -plan_next_step_runnable = plan_next_step_prompt | ChatAnthropic(model="claude-sonnet-4-20250514").with_structured_output(ResearchStep) -check_research_runnable = check_research_prompt | ChatAnthropic(model="claude-sonnet-4-20250514").with_structured_output(ResearchEvaluation) -conduct_research_runnable = conduct_research_prompt | ChatAnthropic(model="claude-sonnet-4-20250514").bind_tools(search_tools+codemap_tools) -extract_implementation_runnable = extract_implementation_prompt | ChatAnthropic(model="claude-sonnet-4-20250514") | JsonOutputParser(pydantic_object=ImplementationPlan) +plan_next_step_runnable = plan_next_step_prompt | ChatAnthropic(model=DEFAULT_MODEL).with_structured_output(ResearchStep) +check_research_runnable = check_research_prompt | ChatAnthropic(model=DEFAULT_MODEL).with_structured_output(ResearchEvaluation) +conduct_research_runnable = conduct_research_prompt | ChatAnthropic(model=DEFAULT_MODEL).bind_tools(search_tools+codemap_tools) +extract_implementation_runnable = extract_implementation_prompt | ChatAnthropic(model=DEFAULT_MODEL) | JsonOutputParser(pydantic_object=ImplementationPlan) tool_node = ToolNode(codemap_tools+search_tools, messages_key="implementation_research_scratchpad") @@ -85,22 +85,6 @@ def conduct_research(state: SoftwareArchitectState): }) return {"implementation_research_scratchpad": [response]} -def convert_tools_messages_to_ai_and_human(implementation_research_scratchpad: List[AnyMessage]): - messages = [] - for message in implementation_research_scratchpad: - if message.type == "ai": - if message.tool_calls: - tool_name = message.tool_calls[0]["name"] - tool_args = json.dumps(message.tool_calls[0]["args"]) - messages.append(AIMessage(content=f"I want to call the tool {tool_name} with the following arguments: {tool_args}")) - else: - messages.append(message) - elif message.type == "tool": - messages.append(HumanMessage(content=f"When executing Tool {message.name} \n The result was {message.content} was called")) - else: - messages.append(message) - return messages - def extract_implementation_plan(state: SoftwareArchitectState): """Extract implementation plan from research findings""" response = extract_implementation_runnable.invoke({ diff --git a/agent/common/utils.py b/agent/common/utils.py new file mode 100644 index 0000000..1f3498c --- /dev/null +++ b/agent/common/utils.py @@ -0,0 +1,44 @@ +"""Shared utilities used across architect and developer agents.""" + +import json +from typing import List + +from langchain_core.messages import AnyMessage, AIMessage, HumanMessage + +# Single source of truth for the Claude model used by all runnables. +# Change this one constant to upgrade the model across the entire agent. +DEFAULT_MODEL = "claude-sonnet-4-20250514" + + +def convert_tools_messages_to_ai_and_human(messages: List[AnyMessage]) -> List[AnyMessage]: + """Convert tool call / tool result messages into plain AI and Human messages. + + LangChain tool messages are not always accepted by every runnable directly. + This helper normalises a mixed message list so every message is either an + AIMessage or a HumanMessage, which any Claude runnable can consume. + + Args: + messages: Raw message list that may contain AIMessages with tool_calls + and ToolMessages with tool results. + + Returns: + List of AIMessage / HumanMessage with the same semantic content. + """ + result = [] + for message in messages: + if message.type == "ai": + if message.tool_calls: + tool_name = message.tool_calls[0]["name"] + tool_args = json.dumps(message.tool_calls[0]["args"]) + result.append( + AIMessage(content=f"I want to call the tool {tool_name} with the following arguments: {tool_args}") + ) + else: + result.append(message) + elif message.type == "tool": + result.append( + HumanMessage(content=f"When executing Tool {message.name} \n The result was {message.content} was called") + ) + else: + result.append(message) + return result diff --git a/agent/developer/graph.py b/agent/developer/graph.py index 22e08c8..3f218be 100644 --- a/agent/developer/graph.py +++ b/agent/developer/graph.py @@ -1,10 +1,9 @@ -import json import os import re from typing import List from diff_match_patch import diff_match_patch from langchain_anthropic import ChatAnthropic -from langchain_core.messages import AnyMessage, AIMessage, HumanMessage +from langchain_core.messages import AnyMessage from langchain_core.output_parsers import StrOutputParser, JsonOutputParser from langgraph.constants import END, START from langgraph.graph import StateGraph @@ -14,21 +13,22 @@ from agent.tools.search import search_tools from agent.tools.codemap import codemap_tools from agent.tools.write import get_files_structure +from agent.common.utils import DEFAULT_MODEL, convert_tools_messages_to_ai_and_human # Load the extract diff prompt extract_diffs_tasks_prompt = markdown_to_prompt_template("agent/developer/prompts/create_diff_prompt.md") implement_diffs_prompt = markdown_to_prompt_template("agent/developer/prompts/implement_diff.md") implement_new_file_prompt = markdown_to_prompt_template("agent/developer/prompts/implement_new_file.md") # Create the runnable with the prompt and model -extract_diff_runnable = extract_diffs_tasks_prompt | ChatAnthropic(model="claude-sonnet-4-20250514") | StrOutputParser() -edit_according_to_diff_runnable = implement_diffs_prompt | ChatAnthropic(model="claude-sonnet-4-20250514") | StrOutputParser() -create_new_file_runnable = implement_new_file_prompt | ChatAnthropic(model="claude-sonnet-4-20250514") | StrOutputParser() +extract_diff_runnable = extract_diffs_tasks_prompt | ChatAnthropic(model=DEFAULT_MODEL) | StrOutputParser() +edit_according_to_diff_runnable = implement_diffs_prompt | ChatAnthropic(model=DEFAULT_MODEL) | StrOutputParser() +create_new_file_runnable = implement_new_file_prompt | ChatAnthropic(model=DEFAULT_MODEL) | StrOutputParser() # Load the get clear implementation plan prompt get_clear_implementation_plan_prompt = markdown_to_prompt_template("agent/developer/prompts/get_clear_implementation_plan.md") # Create the runnable with the prompt and model -get_clear_implementation_plan_runnable = get_clear_implementation_plan_prompt | ChatAnthropic(model="claude-sonnet-4-20250514").bind_tools(search_tools+codemap_tools) +get_clear_implementation_plan_runnable = get_clear_implementation_plan_prompt | ChatAnthropic(model=DEFAULT_MODEL).bind_tools(search_tools+codemap_tools) dmp = diff_match_patch() def start_implementing(state: SoftwareDeveloperState): return { @@ -108,22 +108,6 @@ def is_implementation_complete(state: SoftwareDeveloperState): plan = state.implementation_plan return END if current_task_idx >= len(plan.tasks) else "continue" -def convert_tools_messages_to_ai_and_human(implementation_research_scratchpad: List[AnyMessage]): - messages = [] - for message in implementation_research_scratchpad: - if message.type == "ai": - if message.tool_calls: - tool_name = message.tool_calls[0]["name"] - tool_args = json.dumps(message.tool_calls[0]["args"]) - messages.append(AIMessage(content=f"I want to call the tool {tool_name} with the following arguments: {tool_args}")) - else: - messages.append(message) - elif message.type == "tool": - messages.append(HumanMessage(content=f"When executing Tool {message.name} \n The result was {message.content} was called")) - else: - messages.append(message) - return messages - def creating_diffs_for_task(state: SoftwareDeveloperState): # Get current task information current_task = state.implementation_plan.tasks[state.current_task_idx]