From ab8b1b05a379a6528d01e2cb6269b36b85fa12c7 Mon Sep 17 00:00:00 2001 From: tawsifkamal Date: Wed, 5 Mar 2025 17:31:06 -0800 Subject: [PATCH 1/6] lol --- .../swebench_agent_run/local_run.ipynb | 86 ++++++++++++++++--- 1 file changed, 73 insertions(+), 13 deletions(-) diff --git a/codegen-examples/examples/swebench_agent_run/local_run.ipynb b/codegen-examples/examples/swebench_agent_run/local_run.ipynb index 1f27f470c..1d44789de 100644 --- a/codegen-examples/examples/swebench_agent_run/local_run.ipynb +++ b/codegen-examples/examples/swebench_agent_run/local_run.ipynb @@ -6,8 +6,19 @@ "metadata": {}, "outputs": [], "source": [ + "from codegen.configs.models.codebase import CodebaseConfig\n", + "from codegen import CodeAgent\n", + "from codegen import Codebase\n", + "\n", "%load_ext autoreload\n", - "%autoreload 2" + "%autoreload 2\n", + "\n", + "from dotenv import load_dotenv # type: ignore\n", + "\n", + "load_dotenv()\n", + "\n", + "from codegen.extensions.swebench.utils import SWEBenchDataset, get_swe_bench_examples # noqa: E402\n", + "from run_eval import run_eval # noqa: E402" ] }, { @@ -16,9 +27,7 @@ "metadata": {}, "outputs": [], "source": [ - "from codegen.sdk.core.codebase import Codebase\n", - "from codegen.extensions.swebench.utils import SWEBenchDataset, get_swe_bench_examples\n", - "from run_eval import run_eval" + "examples = get_swe_bench_examples(dataset=SWEBenchDataset.LITE, split=\"test\", offset=0, length=10)" ] }, { @@ -27,18 +36,27 @@ "metadata": {}, "outputs": [], "source": [ - "examples = get_swe_bench_examples(dataset=SWEBenchDataset.LITE, split=\"test\", offset=0, length=1)" + "await run_eval(use_existing_preds=None, dataset=\"lite\", length=10, repo=\"django/django\", local=False, num_workers=5)" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ - "codebase = Codebase.from_repo(examples[0].repo, commit=examples[0].base_commit, tmp_dir=f\"/tmp/{examples[0].instance_id}\")\n", - "# this will allow us to reuse the codebase for multiple examples\n", - "codebases = {examples[0].instance_id: codebase}" + "config = CodebaseConfig(\n", + " disable_file_parse=True, # Disable the graph AND disable file parsing (file.edit only)\n", + ")\n", + "\n", + "codebase = Codebase(\"./\")" ] }, { @@ -47,8 +65,14 @@ "metadata": {}, "outputs": [], "source": [ - "await run_eval(use_existing_preds=None, dataset=\"lite\", length=None, instance_id=examples[0].instance_id, local=True, codebases=codebases)\n", - "codebases[examples[0].instance_id].reset()" + "codegen = CodeAgent(\n", + " codebase=codebase,\n", + " model=\"claude-3-5-sonnet-20240620\",\n", + " temperature=0.0,\n", + " max_tokens=4096,\n", + " max_retries=3,\n", + " timeout=1000,\n", + ")" ] }, { @@ -56,14 +80,50 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": [] + "source": [ + "codegen.run(\"Hello!\")" + ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], - "source": [] + "source": [ + "from langchain_anthropic import ChatAnthropic" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "llm = ChatAnthropic(model=\"claude-3-5-sonnet-20240620\", temperature=0.0, max_tokens=4096, max_retries=3, timeout=1000)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "llm.invoke(\"Hello!\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import anthropic\n", + "import openai\n", + "\n", + "anthropic.RateLimitError\n", + "\n", + "openai.RateLimitError" + ] } ], "metadata": { @@ -82,7 +142,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.1" + "version": "3.13.0" } }, "nbformat": 4, From 18a13ac42e3a78b346f1822253b328017fb136f7 Mon Sep 17 00:00:00 2001 From: tawsifkamal Date: Fri, 7 Mar 2025 13:30:53 -0800 Subject: [PATCH 2/6] done --- .../examples/langchain_agent/run.py | 4 +- .../swebench_agent_run/local_run.ipynb | 92 +-------------- pyproject.toml | 1 + src/codegen/agents/chat_agent.py | 4 +- src/codegen/agents/code_agent.py | 12 +- src/codegen/extensions/langchain/agent.py | 17 +-- src/codegen/extensions/langchain/graph.py | 105 ++++++++++++++++++ src/codegen/extensions/langchain/llm.py | 4 +- src/codegen/extensions/langchain/tools.py | 33 +++++- src/codegen/extensions/tools/search.py | 2 +- 10 files changed, 158 insertions(+), 116 deletions(-) create mode 100644 src/codegen/extensions/langchain/graph.py diff --git a/codegen-examples/examples/langchain_agent/run.py b/codegen-examples/examples/langchain_agent/run.py index 0d4d4f837..5c6891889 100644 --- a/codegen-examples/examples/langchain_agent/run.py +++ b/codegen-examples/examples/langchain_agent/run.py @@ -20,7 +20,7 @@ from langgraph.checkpoint.memory import MemorySaver from langgraph.graph.graph import CompiledGraph -from langgraph.prebuilt import create_react_agent +from codegen.extensions.langchain.graph import create_react_agent from langchain_core.messages import SystemMessage @@ -70,7 +70,7 @@ def create_codebase_agent( memory = MemorySaver() if memory else None - return create_react_agent(model=llm, tools=tools, prompt=system_message, checkpointer=memory, debug=debug) + return create_react_agent(model=llm, tools=tools, system_message=system_message, checkpointer=memory, debug=debug) if __name__ == "__main__": diff --git a/codegen-examples/examples/swebench_agent_run/local_run.ipynb b/codegen-examples/examples/swebench_agent_run/local_run.ipynb index 1d44789de..0b212fa40 100644 --- a/codegen-examples/examples/swebench_agent_run/local_run.ipynb +++ b/codegen-examples/examples/swebench_agent_run/local_run.ipynb @@ -6,10 +6,6 @@ "metadata": {}, "outputs": [], "source": [ - "from codegen.configs.models.codebase import CodebaseConfig\n", - "from codegen import CodeAgent\n", - "from codegen import Codebase\n", - "\n", "%load_ext autoreload\n", "%autoreload 2\n", "\n", @@ -36,93 +32,7 @@ "metadata": {}, "outputs": [], "source": [ - "await run_eval(use_existing_preds=None, dataset=\"lite\", length=10, repo=\"django/django\", local=False, num_workers=5)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "config = CodebaseConfig(\n", - " disable_file_parse=True, # Disable the graph AND disable file parsing (file.edit only)\n", - ")\n", - "\n", - "codebase = Codebase(\"./\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "codegen = CodeAgent(\n", - " codebase=codebase,\n", - " model=\"claude-3-5-sonnet-20240620\",\n", - " temperature=0.0,\n", - " max_tokens=4096,\n", - " max_retries=3,\n", - " timeout=1000,\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "codegen.run(\"Hello!\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from langchain_anthropic import ChatAnthropic" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "llm = ChatAnthropic(model=\"claude-3-5-sonnet-20240620\", temperature=0.0, max_tokens=4096, max_retries=3, timeout=1000)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "llm.invoke(\"Hello!\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import anthropic\n", - "import openai\n", - "\n", - "anthropic.RateLimitError\n", - "\n", - "openai.RateLimitError" + "await run_eval(use_existing_preds=None, dataset=\"lite\", length=20, repo=\"django/django\", num_workers=10, model=\"claude-3-7-sonnet-latest\")" ] } ], diff --git a/pyproject.toml b/pyproject.toml index ea7fa4ae1..e7e0fef4f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ dependencies = [ "tiktoken<1.0.0,>=0.5.1", "tabulate>=0.9.0,<1.0.0", "codeowners<1.0.0,>=0.6.0", + "anthropic", "dataclasses-json<1.0.0,>=0.6.4", "dicttoxml<2.0.0,>=1.7.16", "xmltodict<1.0.0,>=0.13.0", diff --git a/src/codegen/agents/chat_agent.py b/src/codegen/agents/chat_agent.py index 2ecae36fb..24ecada26 100644 --- a/src/codegen/agents/chat_agent.py +++ b/src/codegen/agents/chat_agent.py @@ -44,7 +44,7 @@ def run(self, prompt: str, thread_id: Optional[str] = None) -> str: if thread_id is None: thread_id = str(uuid4()) - input = {"messages": [("user", prompt)]} + input = {"query": prompt} stream = self.agent.stream(input, config={"configurable": {"thread_id": thread_id}}, stream_mode="values") for s in stream: @@ -57,7 +57,7 @@ def run(self, prompt: str, thread_id: Optional[str] = None) -> str: else: message.pretty_print() - return s["messages"][-1].content + return s["final_answer"] def chat(self, prompt: str, thread_id: Optional[str] = None) -> tuple[str, str]: """Chat with the agent, maintaining conversation history. diff --git a/src/codegen/agents/code_agent.py b/src/codegen/agents/code_agent.py index a7c500c90..efa84bfde 100644 --- a/src/codegen/agents/code_agent.py +++ b/src/codegen/agents/code_agent.py @@ -3,7 +3,7 @@ from uuid import uuid4 from langchain.tools import BaseTool -from langchain_core.messages import AIMessage +from langchain_core.messages import AIMessage, HumanMessage from langchain_core.runnables.config import RunnableConfig from langsmith import Client @@ -67,7 +67,7 @@ def run(self, prompt: str, thread_id: Optional[str] = None) -> str: # this message has a reducer which appends the current message to the existing history # see more https://langchain-ai.github.io/langgraph/concepts/low_level/#reducers - input = {"messages": [("user", prompt)]} + input = {"query": prompt} metadata = {"project": self.project_name} tags = [] # Add SWEBench run ID and instance ID to the metadata and tags for filtering @@ -87,7 +87,11 @@ def run(self, prompt: str, thread_id: Optional[str] = None) -> str: run_ids = [] for s in stream: - message = s["messages"][-1] + if len(s["messages"]) == 0: + message = HumanMessage(content=prompt) + else: + message = s["messages"][-1] + if isinstance(message, tuple): print(message) else: @@ -101,7 +105,7 @@ def run(self, prompt: str, thread_id: Optional[str] = None) -> str: run_ids.append(message.additional_kwargs["run_id"]) # Get the last message content - result = s["messages"][-1].content + result = s["final_answer"] # Try to find run IDs in the LangSmith client's recent runs try: diff --git a/src/codegen/extensions/langchain/agent.py b/src/codegen/extensions/langchain/agent.py index aabc57847..fe44594b1 100644 --- a/src/codegen/extensions/langchain/agent.py +++ b/src/codegen/extensions/langchain/agent.py @@ -6,11 +6,10 @@ from langchain_core.messages import SystemMessage from langgraph.checkpoint.memory import MemorySaver from langgraph.graph.graph import CompiledGraph -from langgraph.prebuilt import create_react_agent -from .llm import LLM -from .prompts import REASONER_SYSTEM_MESSAGE -from .tools import ( +from codegen.extensions.langchain.llm import LLM +from codegen.extensions.langchain.prompts import REASONER_SYSTEM_MESSAGE +from codegen.extensions.langchain.tools import ( CreateFileTool, DeleteFileTool, ListDirectoryTool, @@ -25,6 +24,8 @@ ViewFileTool, ) +from .graph import create_react_agent + if TYPE_CHECKING: from codegen import Codebase @@ -88,7 +89,7 @@ def create_codebase_agent( memory = MemorySaver() if memory else None - return create_react_agent(model=llm, tools=tools, prompt=system_message, checkpointer=memory, debug=debug) + return create_react_agent(model=llm, tools=tools, system_message=system_message, checkpointer=memory, debug=debug) def create_chat_agent( @@ -137,7 +138,7 @@ def create_chat_agent( memory = MemorySaver() if memory else None - return create_react_agent(model=llm, tools=tools, prompt=system_message, checkpointer=memory, debug=debug) + return create_react_agent(model=llm, tools=tools, system_message=system_message, checkpointer=memory, debug=debug) def create_codebase_inspector_agent( @@ -174,7 +175,7 @@ def create_codebase_inspector_agent( ] memory = MemorySaver() if memory else None - return create_react_agent(model=llm, tools=tools, prompt=system_message, checkpointer=memory, debug=debug) + return create_react_agent(model=llm, tools=tools, system_message=system_message, checkpointer=memory, debug=debug) def create_agent_with_tools( @@ -208,4 +209,4 @@ def create_agent_with_tools( memory = MemorySaver() if memory else None - return create_react_agent(model=llm, tools=tools, prompt=system_message, checkpointer=memory, debug=debug) + return create_react_agent(model=llm, tools=tools, system_message=system_message, checkpointer=memory, debug=debug) diff --git a/src/codegen/extensions/langchain/graph.py b/src/codegen/extensions/langchain/graph.py new file mode 100644 index 000000000..88ef8f8dd --- /dev/null +++ b/src/codegen/extensions/langchain/graph.py @@ -0,0 +1,105 @@ +"""Demo implementation of an agent with Codegen tools.""" + +from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional + +import anthropic +import openai +from langchain.tools import BaseTool +from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, SystemMessage, ToolMessage +from langgraph.checkpoint.memory import MemorySaver +from langgraph.graph import END, START +from langgraph.graph.message import add_messages +from langgraph.graph.state import CompiledGraph, StateGraph +from langgraph.pregel import RetryPolicy + + +class GraphState(dict[str, Any]): + """State of the graph.""" + + query: str + final_answer: str + messages: Annotated[list[AnyMessage], add_messages] + + +class AgentGraph: + """Main graph class for the agent.""" + + def __init__(self, model: "LLM", tools: list[BaseTool], system_message: SystemMessage): + self.model = model.bind_tools(tools) + self.tools = tools + self.tools_by_name = {tool.name: tool for tool in self.tools} + self.system_message = system_message + + # =================================== NODES ==================================== + + # Reasoner node + def reasoner(self, state: GraphState) -> dict[str, Any]: + new_turn = len(state["messages"]) == 0 or isinstance(state["messages"][-1], AIMessage) + messages = state["messages"] + if new_turn: + query = state["query"] + messages.append(HumanMessage(content=query)) + + result = self.model.invoke([self.system_message, *messages]) + + if isinstance(result, AIMessage): + return {"messages": [*messages, result], "final_answer": result.content} + + return {"messages": [*messages, result]} + + # Tool node + def tool_node(self, state: GraphState) -> dict[str, Any]: + """Performs the tool call""" + result = [] + for tool_call in state["messages"][-1].tool_calls: + tool = self.tools_by_name[tool_call["name"]] + observation = tool.invoke(tool_call["args"]) + result.append(ToolMessage(content=observation, tool_call_id=tool_call["id"])) + return {"messages": result} + + # =================================== EDGE CONDITIONS ==================================== + def should_continue(self, state: GraphState) -> Literal["tools", END]: + messages = state["messages"] + last_message = messages[-1] + if hasattr(last_message, "tool_calls") and last_message.tool_calls: + return "tools" + return END + + # =================================== COMPILE GRAPH ==================================== + def create(self, checkpointer: Optional[MemorySaver] = None, debug: bool = False) -> CompiledGraph: + """Create and compile the graph.""" + builder = StateGraph(GraphState) + + # the retry policy has an initial interval, a backoff factor, and a max interval of controlling the + # amount of time between retries + retry_policy = RetryPolicy(retry_on=[anthropic.RateLimitError, openai.RateLimitError], max_attempts=10) + + # Add nodes + builder.add_node("reasoner", self.reasoner, retry=retry_policy) + builder.add_node("tools", self.tool_node, retry=retry_policy) + + # Add edges + builder.add_edge(START, "reasoner") + builder.add_edge("tools", "reasoner") + builder.add_conditional_edges( + "reasoner", + self.should_continue, + ) + + return builder.compile(checkpointer=checkpointer, debug=debug) + + +def create_react_agent( + model: "LLM", + tools: list[BaseTool], + system_message: SystemMessage, + checkpointer: Optional[MemorySaver] = None, + debug: bool = False, +) -> CompiledGraph: + """Create a reactive agent graph.""" + graph = AgentGraph(model, tools, system_message) + return graph.create(checkpointer=checkpointer, debug=debug) + + +if TYPE_CHECKING: + from codegen.extensions.langchain.llm import LLM diff --git a/src/codegen/extensions/langchain/llm.py b/src/codegen/extensions/langchain/llm.py index 1aafce31c..00bef68c9 100644 --- a/src/codegen/extensions/langchain/llm.py +++ b/src/codegen/extensions/langchain/llm.py @@ -85,13 +85,13 @@ def _get_model(self) -> BaseChatModel: if not os.getenv("ANTHROPIC_API_KEY"): msg = "ANTHROPIC_API_KEY not found in environment. Please set it in your .env file or environment variables." raise ValueError(msg) - return ChatAnthropic(**self._get_model_kwargs()) + return ChatAnthropic(**self._get_model_kwargs(), max_retries=10, timeout=1000) elif self.model_provider == "openai": if not os.getenv("OPENAI_API_KEY"): msg = "OPENAI_API_KEY not found in environment. Please set it in your .env file or environment variables." raise ValueError(msg) - return ChatOpenAI(**self._get_model_kwargs()) + return ChatOpenAI(**self._get_model_kwargs(), max_retries=10, timeout=1000) msg = f"Unknown model provider: {self.model_provider}. Must be one of: anthropic, openai" raise ValueError(msg) diff --git a/src/codegen/extensions/langchain/tools.py b/src/codegen/extensions/langchain/tools.py index acfbdaf59..68cde9674 100644 --- a/src/codegen/extensions/langchain/tools.py +++ b/src/codegen/extensions/langchain/tools.py @@ -202,6 +202,7 @@ class SearchTool(BaseTool): If no exact matches are found, automatically falls back to semantic search to find relevant code even without exact text matches. + Features: - Plain text or regex pattern matching - Directory and file type filtering @@ -867,25 +868,45 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]: class ReplacementEditInput(BaseModel): - filepath: str = Field(..., description="Path to the file to edit relative to the workspace root. The file must exist and be a text file.") + """Input for replacement editing.""" + + filepath: str = Field( + ..., + description=("Path to the file to edit relative to the workspace root. The file must exist and be a text file."), + ) pattern: str = Field( ..., - description="Regular expression pattern to match text that should be replaced. Supports all Python regex syntax including capture groups (\1, \2, etc). The pattern is compiled with re.MULTILINE flag by default.", + description=( + "Regular expression pattern to match text that should be replaced. " + "Supports all Python regex syntax including capture groups (\\1, \\2, etc). " + "The pattern is compiled with re.MULTILINE flag by default." + ), ) replacement: str = Field( ..., - description="Text to replace matched patterns with. Can reference regex capture groups using \1, \2, etc. If using regex groups in pattern, make sure to preserve them in replacement if needed.", + description=( + "Text to replace matched patterns with. Can reference regex capture groups using \\1, \\2, etc. If using regex groups in pattern, make sure to preserve them in replacement if needed." + ), ) start: int = Field( - default=1, description="Starting line number (1-indexed, inclusive) to begin replacements from. Use this with 'end' to limit changes to a specific region. Default is 1 (start of file)." + default=1, + description=("Starting line number (1-indexed, inclusive) to begin replacements from. Use this with 'end' to limit changes to a specific region. Default is 1 (start of file)."), ) end: int = Field( default=-1, - description="Ending line number (1-indexed, inclusive) to stop replacements at. Use -1 to indicate end of file. Use this with 'start' to limit changes to a specific region. Default is -1 (end of file).", + description=( + "Ending line number (1-indexed, inclusive) to stop replacements at. " + "Use -1 to indicate end of file. Use this with 'start' to limit changes to a specific region. " + "Default is -1 (end of file)." + ), ) count: Optional[int] = Field( default=None, - description="Maximum number of replacements to make. Use None to replace all occurrences (default), or specify a number to limit replacements. Useful when you only want to replace the first N occurrences.", + description=( + "Maximum number of replacements to make. " + "Use None to replace all occurrences (default), or specify a number to limit replacements. " + "Useful when you only want to replace the first N occurrences." + ), ) diff --git a/src/codegen/extensions/tools/search.py b/src/codegen/extensions/tools/search.py index 8083e7db8..2a7b090e4 100644 --- a/src/codegen/extensions/tools/search.py +++ b/src/codegen/extensions/tools/search.py @@ -149,7 +149,7 @@ def _search_with_ripgrep( pass # Add the query and path - cmd.append(query) + cmd.append(f"{query}") cmd.append(search_path) # Run ripgrep From f2505f64ad5cf8add02a52b4ff16495b621f3148 Mon Sep 17 00:00:00 2001 From: tawsifkamal Date: Mon, 10 Mar 2025 10:31:54 -0700 Subject: [PATCH 3/6] done --- src/codegen/extensions/langchain/graph.py | 15 ++------- src/codegen/extensions/langchain/tools.py | 38 ++++++++++++++++------- 2 files changed, 29 insertions(+), 24 deletions(-) diff --git a/src/codegen/extensions/langchain/graph.py b/src/codegen/extensions/langchain/graph.py index 88ef8f8dd..d5ea786c4 100644 --- a/src/codegen/extensions/langchain/graph.py +++ b/src/codegen/extensions/langchain/graph.py @@ -5,11 +5,12 @@ import anthropic import openai from langchain.tools import BaseTool -from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, SystemMessage, ToolMessage +from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, SystemMessage from langgraph.checkpoint.memory import MemorySaver from langgraph.graph import END, START from langgraph.graph.message import add_messages from langgraph.graph.state import CompiledGraph, StateGraph +from langgraph.prebuilt import ToolNode from langgraph.pregel import RetryPolicy @@ -47,16 +48,6 @@ def reasoner(self, state: GraphState) -> dict[str, Any]: return {"messages": [*messages, result]} - # Tool node - def tool_node(self, state: GraphState) -> dict[str, Any]: - """Performs the tool call""" - result = [] - for tool_call in state["messages"][-1].tool_calls: - tool = self.tools_by_name[tool_call["name"]] - observation = tool.invoke(tool_call["args"]) - result.append(ToolMessage(content=observation, tool_call_id=tool_call["id"])) - return {"messages": result} - # =================================== EDGE CONDITIONS ==================================== def should_continue(self, state: GraphState) -> Literal["tools", END]: messages = state["messages"] @@ -76,7 +67,7 @@ def create(self, checkpointer: Optional[MemorySaver] = None, debug: bool = False # Add nodes builder.add_node("reasoner", self.reasoner, retry=retry_policy) - builder.add_node("tools", self.tool_node, retry=retry_policy) + builder.add_node("tools", ToolNode(self.tools), retry=retry_policy) # Add edges builder.add_edge(START, "reasoner") diff --git a/src/codegen/extensions/langchain/tools.py b/src/codegen/extensions/langchain/tools.py index 96469ae89..2ab5b0179 100644 --- a/src/codegen/extensions/langchain/tools.py +++ b/src/codegen/extensions/langchain/tools.py @@ -117,7 +117,7 @@ class SearchInput(BaseModel): query: str = Field( ..., description="""The search query to find in the codebase. When ripgrep is available, this will be passed as a ripgrep pattern. For regex searches, set use_regex=True. -Ripgrep is the preferred method.""", + Ripgrep is the preferred method.""", ) target_directories: Optional[list[str]] = Field(default=None, description="Optional list of directories to search in") file_extensions: Optional[list[str]] = Field(default=None, description="Optional list of file extensions to search (e.g. ['.py', '.ts'])") @@ -850,31 +850,45 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]: class ReplacementEditInput(BaseModel): - filepath: str = Field(..., description="Path to the file to edit relative to the workspace root. The file must exist and be a text file.") + """Input for replacement editing.""" + + filepath: str = Field( + ..., + description=("Path to the file to edit relative to the workspace root. The file must exist and be a text file."), + ) pattern: str = Field( ..., - description="""Regular expression pattern to match text that should be replaced. -Supports all Python regex syntax including capture groups (\1, \2, etc). The pattern is compiled with re.MULTILINE flag by default.""", + description=( + "Regular expression pattern to match text that should be replaced. " + "Supports all Python regex syntax including capture groups (\\1, \\2, etc). " + "The pattern is compiled with re.MULTILINE flag by default." + ), ) replacement: str = Field( ..., - description="""Text to replace matched patterns with. -Can reference regex capture groups using \1, \2, etc. If using regex groups in pattern, make sure to preserve them in replacement if needed.""", + description=( + "Text to replace matched patterns with. Can reference regex capture groups using \\1, \\2, etc. If using regex groups in pattern, make sure to preserve them in replacement if needed." + ), ) start: int = Field( default=1, - description="""Starting line number (1-indexed, inclusive) to begin replacements from. -Use this with 'end' to limit changes to a specific region. Default is 1 (start of file).""", + description=("Starting line number (1-indexed, inclusive) to begin replacements from. Use this with 'end' to limit changes to a specific region. Default is 1 (start of file)."), ) end: int = Field( default=-1, - description="""Ending line number (1-indexed, inclusive) to stop replacements at. -Use -1 to indicate end of file. Use this with 'start' to limit changes to a specific region. Default is -1 (end of file).""", + description=( + "Ending line number (1-indexed, inclusive) to stop replacements at. " + "Use -1 to indicate end of file. Use this with 'start' to limit changes to a specific region. " + "Default is -1 (end of file)." + ), ) count: Optional[int] = Field( default=None, - description="""Maximum number of replacements to make. Use None to replace all occurrences (default), or specify a number to limit replacements. -Useful when you only want to replace the first N occurrences.""", + description=( + "Maximum number of replacements to make. " + "Use None to replace all occurrences (default), or specify a number to limit replacements. " + "Useful when you only want to replace the first N occurrences." + ), ) From 583fe42c6008c4430d1ca7308ff4ddaf1c9b6391 Mon Sep 17 00:00:00 2001 From: tawsifkamal Date: Mon, 10 Mar 2025 10:33:34 -0700 Subject: [PATCH 4/6] retry stuff --- src/codegen/extensions/langchain/llm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/codegen/extensions/langchain/llm.py b/src/codegen/extensions/langchain/llm.py index 54b9a91a2..d2cec02a3 100644 --- a/src/codegen/extensions/langchain/llm.py +++ b/src/codegen/extensions/langchain/llm.py @@ -89,13 +89,13 @@ def _get_model(self) -> BaseChatModel: if not os.getenv("ANTHROPIC_API_KEY"): msg = "ANTHROPIC_API_KEY not found in environment. Please set it in your .env file or environment variables." raise ValueError(msg) - return ChatAnthropic(**self._get_model_kwargs(), max_retries=10, timeout=1000) + return ChatAnthropic(**self._get_model_kwargs()) elif self.model_provider == "openai": if not os.getenv("OPENAI_API_KEY"): msg = "OPENAI_API_KEY not found in environment. Please set it in your .env file or environment variables." raise ValueError(msg) - return ChatOpenAI(**self._get_model_kwargs(), max_retries=10, timeout=1000) + return ChatOpenAI(**self._get_model_kwargs()) elif self.model_provider == "xai": if not os.getenv("XAI_API_KEY"): From cf30d0f627d070c9ec21d544033129fe7ea6cc61 Mon Sep 17 00:00:00 2001 From: tawsifkamal Date: Mon, 10 Mar 2025 11:14:55 -0700 Subject: [PATCH 5/6] done --- src/codegen/extensions/langchain/graph.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/codegen/extensions/langchain/graph.py b/src/codegen/extensions/langchain/graph.py index d5ea786c4..654bba69c 100644 --- a/src/codegen/extensions/langchain/graph.py +++ b/src/codegen/extensions/langchain/graph.py @@ -28,7 +28,6 @@ class AgentGraph: def __init__(self, model: "LLM", tools: list[BaseTool], system_message: SystemMessage): self.model = model.bind_tools(tools) self.tools = tools - self.tools_by_name = {tool.name: tool for tool in self.tools} self.system_message = system_message # =================================== NODES ==================================== @@ -63,7 +62,14 @@ def create(self, checkpointer: Optional[MemorySaver] = None, debug: bool = False # the retry policy has an initial interval, a backoff factor, and a max interval of controlling the # amount of time between retries - retry_policy = RetryPolicy(retry_on=[anthropic.RateLimitError, openai.RateLimitError], max_attempts=10) + retry_policy = RetryPolicy( + retry_on=[anthropic.RateLimitError, openai.RateLimitError], + max_attempts=10, + initial_interval=30.0, # Start with 10 second wait + backoff_factor=0.5, # Double the wait time each retry + max_interval=60.0, # Cap at 1 minute max wait + jitter=True, + ) # Add nodes builder.add_node("reasoner", self.reasoner, retry=retry_policy) From d0227977176744b9c9ec4e6db713997f0d599256 Mon Sep 17 00:00:00 2001 From: tawsifkamal Date: Mon, 10 Mar 2025 11:30:01 -0700 Subject: [PATCH 6/6] done --- src/codegen/extensions/langchain/graph.py | 6 +++--- src/codegen/extensions/langchain/llm.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/codegen/extensions/langchain/graph.py b/src/codegen/extensions/langchain/graph.py index 654bba69c..3685ea322 100644 --- a/src/codegen/extensions/langchain/graph.py +++ b/src/codegen/extensions/langchain/graph.py @@ -65,9 +65,9 @@ def create(self, checkpointer: Optional[MemorySaver] = None, debug: bool = False retry_policy = RetryPolicy( retry_on=[anthropic.RateLimitError, openai.RateLimitError], max_attempts=10, - initial_interval=30.0, # Start with 10 second wait - backoff_factor=0.5, # Double the wait time each retry - max_interval=60.0, # Cap at 1 minute max wait + initial_interval=30.0, # Start with 30 second wait + backoff_factor=2, # Double the wait time each retry + max_interval=1000.0, # Cap at 1000 second max wait jitter=True, ) diff --git a/src/codegen/extensions/langchain/llm.py b/src/codegen/extensions/langchain/llm.py index d2cec02a3..54b9a91a2 100644 --- a/src/codegen/extensions/langchain/llm.py +++ b/src/codegen/extensions/langchain/llm.py @@ -89,13 +89,13 @@ def _get_model(self) -> BaseChatModel: if not os.getenv("ANTHROPIC_API_KEY"): msg = "ANTHROPIC_API_KEY not found in environment. Please set it in your .env file or environment variables." raise ValueError(msg) - return ChatAnthropic(**self._get_model_kwargs()) + return ChatAnthropic(**self._get_model_kwargs(), max_retries=10, timeout=1000) elif self.model_provider == "openai": if not os.getenv("OPENAI_API_KEY"): msg = "OPENAI_API_KEY not found in environment. Please set it in your .env file or environment variables." raise ValueError(msg) - return ChatOpenAI(**self._get_model_kwargs()) + return ChatOpenAI(**self._get_model_kwargs(), max_retries=10, timeout=1000) elif self.model_provider == "xai": if not os.getenv("XAI_API_KEY"):