feat: top-level code agent (#564)

Author: jayhack

codegen-examples/examples/codegen-mcp-server/server.py

@@ -1,8 +1,9 @@
 import asyncio
 from dataclasses import dataclass, field
-from typing import Annotated, Optional, Dict, Any, List
-from mcp.server.fastmcp import FastMCP
+from typing import Annotated, Any, Dict, List, Optional
+
 from codegen import Codebase
+from mcp.server.fastmcp import FastMCP
 
 
 @dataclass

codegen-examples/examples/sqlalchemy_soft_delete/run.py

@@ -92,8 +92,8 @@ def process_soft_deletes(codebase):
 
 
 if __name__ == "__main__":
+    from codegen import Codebase
     from codegen.sdk.codebase.config import CodebaseConfig
-    from codegen.sdk.core.codebase import Codebase
 
     repo_path = Path("/tmp/core")
     repo_url = "https://github.com/hasgeek/funnel.git"

codegen-examples/examples/sqlalchemy_type_annotations/run.py

@@ -1,11 +1,10 @@
-import codegen
-
+import os
+import shutil
+import subprocess
 
+import codegen
 from codegen import Codebase
 from codegen.sdk.core.detached_symbols.function_call import FunctionCall
-import subprocess
-import shutil
-import os
 
 
 def init_git_repo(repo_path: str) -> None:

src/codegen/__init__.py

@@ -1,3 +1,4 @@
+from codegen.agents.code_agent import CodeAgent
 from codegen.cli.sdk.decorator import function
 from codegen.cli.sdk.functions import Function
 
@@ -6,4 +7,4 @@
 from codegen.sdk.core.codebase import Codebase
 from codegen.shared.enums.programming_language import ProgrammingLanguage
 
-__all__ = ["Codebase", "Function", "ProgrammingLanguage", "function"]
+__all__ = ["CodeAgent", "Codebase", "Function", "ProgrammingLanguage", "function"]

src/codegen/agents/__init__.py

@@ -0,0 +1,21 @@
+from typing import Optional
+from uuid import uuid4
+
+from codegen.extensions.langchain.agent import create_codebase_agent
+from codegen.sdk.core.codebase import Codebase
+
+
+class CodeAgent:
+    """Agent for interacting with a codebase."""
+
+    def __init__(self, codebase: Codebase):
+        self.codebase = codebase
+        self.agent = create_codebase_agent(self.codebase)
+
+    def run(self, prompt: str, session_id: Optional[str] = None) -> str:
+        if session_id is None:
+            session_id = str(uuid4())
+        return self.agent.invoke(
+            {"input": prompt},
+            config={"configurable": {"session_id": session_id}},
+        )

src/codegen/agents/code_agent.py

@@ -3,7 +3,7 @@
 
 def convert_to_cli(input: str, language: str, name: str) -> str:
     return f"""import codegen
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 
 @codegen.function('{name}')

src/codegen/cli/codemod/convert.py

@@ -7,7 +7,6 @@
 from rich.markdown import Markdown
 from rich.prompt import Prompt
 
-from codegen import Codebase
 from codegen.extensions.langchain.agent import create_agent_with_tools
 from codegen.extensions.langchain.tools import (
     CreateFileTool,
@@ -20,6 +19,7 @@
     SearchTool,
     ViewFileTool,
 )
+from codegen.sdk.core.codebase import Codebase
 
 # Suppress specific warnings
 warnings.filterwarnings("ignore", message=".*Helicone.*")

src/codegen/cli/commands/agent/main.py

@@ -4,9 +4,9 @@
 from rich.panel import Panel
 from rich.status import Status
 
-from codegen import Codebase
 from codegen.cli.auth.session import CodegenSession
 from codegen.cli.utils.function_finder import DecoratedFunction
+from codegen.sdk.core.codebase import Codebase
 
 
 def parse_codebase(repo_root: Path) -> Codebase:

src/codegen/cli/commands/run/run_local.py

@@ -1,5 +1,5 @@
 DEFAULT_CODEMOD = '''import codegen
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 
 @codegen.function("{name}")

src/codegen/cli/utils/default_code.py

@@ -5,7 +5,7 @@
 DEFAULT_CELLS = [
     {
         "cell_type": "code",
-        "source": """from codegen import Codebase
+        "source": """from codegen.sdk.core.codebase import Codebase
 
 # Initialize codebase
 codebase = Codebase('../../')
@@ -33,7 +33,7 @@
     },
     {
         "cell_type": "code",
-        "source": """from codegen import Codebase
+        "source": """from codegen.sdk.core.codebase import Codebase
 
 # Initialize FastAPI codebase
 print('Cloning and parsing FastAPI to /tmp/codegen/fastapi...')

src/codegen/cli/utils/notebooks.py

@@ -1,6 +1,6 @@
-from codegen import Codebase
 from codegen.extensions.graph.create_graph import create_codebase_graph
 from codegen.extensions.graph.neo4j_exporter import Neo4jExporter
+from codegen.sdk.core.codebase import Codebase
 
 
 def visualize_codebase(codebase, neo4j_uri: str, username: str, password: str):

src/codegen/extensions/graph/main.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 T = TypeVar("T")  # Type of the items being indexed (e.g., File, Symbol)
 

src/codegen/extensions/index/code_index.py

@@ -9,8 +9,8 @@
 from openai import OpenAI
 from tqdm import tqdm
 
-from codegen import Codebase
 from codegen.extensions.index.code_index import CodeIndex
+from codegen.sdk.core.codebase import Codebase
 from codegen.sdk.core.file import File
 
 logger = logging.getLogger(__name__)

src/codegen/extensions/index/file_index.py

@@ -8,8 +8,8 @@
 from openai import OpenAI
 from tqdm import tqdm
 
-from codegen import Codebase
 from codegen.extensions.index.code_index import CodeIndex
+from codegen.sdk.core.codebase import Codebase
 from codegen.sdk.core.symbol import Symbol
 
 logger = logging.getLogger(__name__)

src/codegen/extensions/index/symbol_index.py

@@ -2,7 +2,7 @@
 
 from langchain_core.tools.base import BaseTool
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from .tools import (
     CommitTool,

src/codegen/extensions/langchain/__init__.py

@@ -11,7 +11,7 @@
 from langchain_core.runnables.history import RunnableWithMessageHistory
 from langchain_openai import ChatOpenAI
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from .tools import (
     CreateFileTool,

src/codegen/extensions/langchain/agent.py

@@ -6,7 +6,6 @@
 from langchain_core.tools.base import BaseTool
 from pydantic import BaseModel, Field
 
-from codegen import Codebase
 from codegen.extensions.linear.linear_client import LinearClient
 from codegen.extensions.tools.bash import run_bash_command
 from codegen.extensions.tools.linear.linear import (
@@ -20,10 +19,10 @@
 from codegen.extensions.tools.link_annotation import add_links_to_message
 from codegen.extensions.tools.replacement_edit import replacement_edit
 from codegen.extensions.tools.reveal_symbol import reveal_symbol
-from codegen.extensions.tools.run_codemod import run_codemod
 from codegen.extensions.tools.search import search
 from codegen.extensions.tools.semantic_edit import semantic_edit
 from codegen.extensions.tools.semantic_search import semantic_search
+from codegen.sdk.core.codebase import Codebase
 
 from ..tools import (
     commit,
@@ -712,8 +711,7 @@ def get_workspace_tools(codebase: Codebase) -> list["BaseTool"]:
         RenameFileTool(codebase),
         ReplacementEditTool(codebase),
         RevealSymbolTool(codebase),
-        RunBashCommandTool(),
-        RunCodemodTool(codebase),
+        RunBashCommandTool(),  # Note: This tool doesn't need the codebase
         SearchTool(codebase),
         SemanticEditTool(codebase),
         SemanticSearchTool(codebase),
@@ -774,38 +772,3 @@ def _run(
             count=count,
         )
         return json.dumps(result, indent=2)
-
-
-class RunCodemodInput(BaseModel):
-    """Input for running a codemod."""
-
-    codemod_source: str = Field(
-        ...,
-        description="""Source code of the codemod function. Must define a 'run(codebase: Codebase)' function that makes the desired changes.
-Example:
-```python
-def run(codebase: Codebase):
-    for file in codebase.files:
-        if file.filepath.endswith('.py'):
-            content = file.content
-            # Make changes to content
-            file.edit(new_content)
-```
-""",
-    )
-
-
-class RunCodemodTool(BaseTool):
-    """Tool for running custom codemod functions."""
-
-    name: ClassVar[str] = "run_codemod"
-    description: ClassVar[str] = "Run a custom codemod function to make systematic changes across the codebase"
-    args_schema: ClassVar[type[BaseModel]] = RunCodemodInput
-    codebase: Codebase = Field(exclude=True)
-
-    def __init__(self, codebase: Codebase) -> None:
-        super().__init__(codebase=codebase)
-
-    def _run(self, codemod_source: str) -> str:
-        result = run_codemod(self.codebase, codemod_source)
-        return json.dumps(result, indent=2)

src/codegen/extensions/langchain/tools.py

@@ -4,7 +4,7 @@
 
 from pydantic import Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from .observation import Observation
 

src/codegen/extensions/tools/commit.py

@@ -4,7 +4,7 @@
 
 from pydantic import Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from .observation import Observation
 from .view_file import ViewFileObservation, view_file

src/codegen/extensions/tools/create_file.py

@@ -4,7 +4,7 @@
 
 from pydantic import Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from .observation import Observation
 

src/codegen/extensions/tools/delete_file.py

@@ -4,7 +4,7 @@
 
 from pydantic import Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from .observation import Observation
 from .view_file import ViewFileObservation, view_file

src/codegen/extensions/tools/edit_file.py

@@ -6,7 +6,7 @@
 from github import GithubException
 from pydantic import Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from ..observation import Observation
 

src/codegen/extensions/tools/github/create_pr.py

@@ -4,7 +4,7 @@
 
 from pydantic import Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from ..observation import Observation
 

src/codegen/extensions/tools/github/create_pr_comment.py

@@ -4,7 +4,7 @@
 
 from pydantic import Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from ..observation import Observation
 

src/codegen/extensions/tools/github/create_pr_review_comment.py

@@ -4,7 +4,7 @@
 
 from pydantic import Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from ..observation import Observation
 

src/codegen/extensions/tools/github/view_pr.py

@@ -4,7 +4,7 @@
 from enum import StrEnum
 from typing import Callable
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 
 class MessageChannel(StrEnum):

src/codegen/extensions/tools/link_annotation.py

@@ -4,7 +4,7 @@
 
 from pydantic import BaseModel, Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 from codegen.sdk.core.directory import Directory
 
 from .observation import Observation

src/codegen/extensions/tools/list_directory.py

@@ -4,7 +4,7 @@
 
 from pydantic import Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from .observation import Observation
 from .view_file import ViewFileObservation, view_file

src/codegen/extensions/tools/move_symbol.py

@@ -4,7 +4,7 @@
 
 from pydantic import Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from .observation import Observation
 from .view_file import ViewFileObservation, view_file

src/codegen/extensions/tools/rename_file.py

@@ -6,7 +6,7 @@
 
 from pydantic import Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from .observation import Observation
 from .view_file import add_line_numbers

src/codegen/extensions/tools/replacement_edit.py

@@ -5,8 +5,8 @@
 import tiktoken
 from pydantic import Field
 
-from codegen import Codebase
 from codegen.sdk.ai.utils import count_tokens
+from codegen.sdk.core.codebase import Codebase
 from codegen.sdk.core.external_module import ExternalModule
 from codegen.sdk.core.import_resolution import Import
 from codegen.sdk.core.symbol import Symbol