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}")