diff --git a/codegen-examples/examples/codegen-mcp-server/server.py b/codegen-examples/examples/codegen-mcp-server/server.py
index 7f312f9f9..91e21e3b3 100644
--- a/codegen-examples/examples/codegen-mcp-server/server.py
+++ b/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
diff --git a/codegen-examples/examples/sqlalchemy_soft_delete/run.py b/codegen-examples/examples/sqlalchemy_soft_delete/run.py
index 22a4d68df..f25c64e47 100644
--- a/codegen-examples/examples/sqlalchemy_soft_delete/run.py
+++ b/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"
diff --git a/codegen-examples/examples/sqlalchemy_type_annotations/run.py b/codegen-examples/examples/sqlalchemy_type_annotations/run.py
index 0bcc6173e..96574152d 100644
--- a/codegen-examples/examples/sqlalchemy_type_annotations/run.py
+++ b/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:
diff --git a/src/codegen/__init__.py b/src/codegen/__init__.py
index 7349dde28..7f37ec124 100644
--- a/src/codegen/__init__.py
+++ b/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"]
diff --git a/src/codegen/agents/__init__.py b/src/codegen/agents/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/codegen/agents/code_agent.py b/src/codegen/agents/code_agent.py
new file mode 100644
index 000000000..7858d2609
--- /dev/null
+++ b/src/codegen/agents/code_agent.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}},
+        )
diff --git a/src/codegen/cli/codemod/convert.py b/src/codegen/cli/codemod/convert.py
index 15c80f4d5..b86b58d14 100644
--- a/src/codegen/cli/codemod/convert.py
+++ b/src/codegen/cli/codemod/convert.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}')
diff --git a/src/codegen/cli/commands/agent/main.py b/src/codegen/cli/commands/agent/main.py
index acf3b4a84..e3ba521a6 100644
--- a/src/codegen/cli/commands/agent/main.py
+++ b/src/codegen/cli/commands/agent/main.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.*")
diff --git a/src/codegen/cli/commands/run/run_local.py b/src/codegen/cli/commands/run/run_local.py
index 71c3a07f7..bfe8c7f1a 100644
--- a/src/codegen/cli/commands/run/run_local.py
+++ b/src/codegen/cli/commands/run/run_local.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:
diff --git a/src/codegen/cli/utils/default_code.py b/src/codegen/cli/utils/default_code.py
index 2bed33f80..5d42596e0 100644
--- a/src/codegen/cli/utils/default_code.py
+++ b/src/codegen/cli/utils/default_code.py
@@ -1,5 +1,5 @@
 DEFAULT_CODEMOD = '''import codegen
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 
 @codegen.function("{name}")
diff --git a/src/codegen/cli/utils/notebooks.py b/src/codegen/cli/utils/notebooks.py
index 2311cabdf..f859adc0f 100644
--- a/src/codegen/cli/utils/notebooks.py
+++ b/src/codegen/cli/utils/notebooks.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...')
diff --git a/src/codegen/extensions/graph/main.py b/src/codegen/extensions/graph/main.py
index 851d572e9..c2a655b2e 100644
--- a/src/codegen/extensions/graph/main.py
+++ b/src/codegen/extensions/graph/main.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):
diff --git a/src/codegen/extensions/index/code_index.py b/src/codegen/extensions/index/code_index.py
index 4cc50ab52..4cf8a5de3 100644
--- a/src/codegen/extensions/index/code_index.py
+++ b/src/codegen/extensions/index/code_index.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)
 
diff --git a/src/codegen/extensions/index/file_index.py b/src/codegen/extensions/index/file_index.py
index 81b0174d9..76759246f 100644
--- a/src/codegen/extensions/index/file_index.py
+++ b/src/codegen/extensions/index/file_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__)
diff --git a/src/codegen/extensions/index/symbol_index.py b/src/codegen/extensions/index/symbol_index.py
index 8874dffc2..42e39e4d5 100644
--- a/src/codegen/extensions/index/symbol_index.py
+++ b/src/codegen/extensions/index/symbol_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__)
diff --git a/src/codegen/extensions/langchain/__init__.py b/src/codegen/extensions/langchain/__init__.py
index 957e7913d..0df13e62b 100644
--- a/src/codegen/extensions/langchain/__init__.py
+++ b/src/codegen/extensions/langchain/__init__.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,
diff --git a/src/codegen/extensions/langchain/agent.py b/src/codegen/extensions/langchain/agent.py
index 9d778072b..f8e4cbf26 100644
--- a/src/codegen/extensions/langchain/agent.py
+++ b/src/codegen/extensions/langchain/agent.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,
diff --git a/src/codegen/extensions/langchain/tools.py b/src/codegen/extensions/langchain/tools.py
index 6c17f4087..86f3cc3dc 100644
--- a/src/codegen/extensions/langchain/tools.py
+++ b/src/codegen/extensions/langchain/tools.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)
diff --git a/src/codegen/extensions/tools/commit.py b/src/codegen/extensions/tools/commit.py
index 3bd931756..d01634142 100644
--- a/src/codegen/extensions/tools/commit.py
+++ b/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
 
diff --git a/src/codegen/extensions/tools/create_file.py b/src/codegen/extensions/tools/create_file.py
index 903cd11bf..d40cdef4c 100644
--- a/src/codegen/extensions/tools/create_file.py
+++ b/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
 from .view_file import ViewFileObservation, view_file
diff --git a/src/codegen/extensions/tools/delete_file.py b/src/codegen/extensions/tools/delete_file.py
index 1f23ef265..8ac3e4c28 100644
--- a/src/codegen/extensions/tools/delete_file.py
+++ b/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
 
diff --git a/src/codegen/extensions/tools/edit_file.py b/src/codegen/extensions/tools/edit_file.py
index 50e85b73d..784697441 100644
--- a/src/codegen/extensions/tools/edit_file.py
+++ b/src/codegen/extensions/tools/edit_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
diff --git a/src/codegen/extensions/tools/github/create_pr.py b/src/codegen/extensions/tools/github/create_pr.py
index 0e10b850b..70da33b3d 100644
--- a/src/codegen/extensions/tools/github/create_pr.py
+++ b/src/codegen/extensions/tools/github/create_pr.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
 
diff --git a/src/codegen/extensions/tools/github/create_pr_comment.py b/src/codegen/extensions/tools/github/create_pr_comment.py
index 3f538413c..aaff7179d 100644
--- a/src/codegen/extensions/tools/github/create_pr_comment.py
+++ b/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
 
diff --git a/src/codegen/extensions/tools/github/create_pr_review_comment.py b/src/codegen/extensions/tools/github/create_pr_review_comment.py
index ac6ea6b86..a9a924e85 100644
--- a/src/codegen/extensions/tools/github/create_pr_review_comment.py
+++ b/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
 
diff --git a/src/codegen/extensions/tools/github/view_pr.py b/src/codegen/extensions/tools/github/view_pr.py
index a698d80d2..1210df107 100644
--- a/src/codegen/extensions/tools/github/view_pr.py
+++ b/src/codegen/extensions/tools/github/view_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
 
diff --git a/src/codegen/extensions/tools/link_annotation.py b/src/codegen/extensions/tools/link_annotation.py
index bcb128a16..543f0fc3a 100644
--- a/src/codegen/extensions/tools/link_annotation.py
+++ b/src/codegen/extensions/tools/link_annotation.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):
diff --git a/src/codegen/extensions/tools/list_directory.py b/src/codegen/extensions/tools/list_directory.py
index ed76a953b..d1f247c35 100644
--- a/src/codegen/extensions/tools/list_directory.py
+++ b/src/codegen/extensions/tools/list_directory.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
diff --git a/src/codegen/extensions/tools/move_symbol.py b/src/codegen/extensions/tools/move_symbol.py
index 6a86be4e4..9ef3a5a85 100644
--- a/src/codegen/extensions/tools/move_symbol.py
+++ b/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
diff --git a/src/codegen/extensions/tools/rename_file.py b/src/codegen/extensions/tools/rename_file.py
index ce4111865..832588d8b 100644
--- a/src/codegen/extensions/tools/rename_file.py
+++ b/src/codegen/extensions/tools/rename_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
diff --git a/src/codegen/extensions/tools/replacement_edit.py b/src/codegen/extensions/tools/replacement_edit.py
index 3bac90e17..0c61c8f96 100644
--- a/src/codegen/extensions/tools/replacement_edit.py
+++ b/src/codegen/extensions/tools/replacement_edit.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
diff --git a/src/codegen/extensions/tools/reveal_symbol.py b/src/codegen/extensions/tools/reveal_symbol.py
index 6279eaeba..c91b0a111 100644
--- a/src/codegen/extensions/tools/reveal_symbol.py
+++ b/src/codegen/extensions/tools/reveal_symbol.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
diff --git a/src/codegen/extensions/tools/run_codemod.py b/src/codegen/extensions/tools/run_codemod.py
index de9bd991c..f73d13aff 100644
--- a/src/codegen/extensions/tools/run_codemod.py
+++ b/src/codegen/extensions/tools/run_codemod.py
@@ -6,7 +6,7 @@
 from tempfile import NamedTemporaryFile
 from typing import Any
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 
 def run_codemod(codebase: Codebase, codemod_source: str) -> dict[str, Any]:
@@ -32,7 +32,7 @@ def run(codebase: Codebase):
     # Create a temporary module to run the codemod
     with NamedTemporaryFile(suffix=".py", mode="w", delete=False) as temp_file:
         # Add imports and write the codemod source
-        temp_file.write("from codegen import Codebase\n\n")
+        temp_file.write("from codegen.sdk.core.codebase import Codebase\n\n")
         temp_file.write(codemod_source)
         temp_file.flush()
 
diff --git a/src/codegen/extensions/tools/search.py b/src/codegen/extensions/tools/search.py
index 0923f6837..1b3a21d0b 100644
--- a/src/codegen/extensions/tools/search.py
+++ b/src/codegen/extensions/tools/search.py
@@ -10,7 +10,7 @@
 
 from pydantic import Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from .observation import Observation
 
diff --git a/src/codegen/extensions/tools/semantic_edit.py b/src/codegen/extensions/tools/semantic_edit.py
index 642a245f4..fd30efa93 100644
--- a/src/codegen/extensions/tools/semantic_edit.py
+++ b/src/codegen/extensions/tools/semantic_edit.py
@@ -8,7 +8,7 @@
 from langchain_core.prompts import ChatPromptTemplate
 from pydantic import Field
 
-from codegen import Codebase
+from codegen.sdk.core.codebase import Codebase
 
 from .observation import Observation
 from .semantic_edit_prompts import _HUMAN_PROMPT_DRAFT_EDITOR, COMMANDER_SYSTEM_PROMPT
diff --git a/src/codegen/extensions/tools/semantic_search.py b/src/codegen/extensions/tools/semantic_search.py
index f2876ca3e..93cf05212 100644
--- a/src/codegen/extensions/tools/semantic_search.py
+++ b/src/codegen/extensions/tools/semantic_search.py
@@ -4,8 +4,8 @@
 
 from pydantic import Field
 
-from codegen import Codebase
 from codegen.extensions.index.file_index import FileIndex
+from codegen.sdk.core.codebase import Codebase
 
 from .observation import Observation
 
diff --git a/src/codegen/extensions/tools/view_file.py b/src/codegen/extensions/tools/view_file.py
index 41d1276c5..68300885b 100644
--- a/src/codegen/extensions/tools/view_file.py
+++ b/src/codegen/extensions/tools/view_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
 
diff --git a/tests/integration/codegen/test_imports.py b/tests/integration/codegen/test_imports.py
index 6004c3ffc..b7b3fb039 100644
--- a/tests/integration/codegen/test_imports.py
+++ b/tests/integration/codegen/test_imports.py
@@ -1,8 +1,8 @@
 import os
 
 import codegen
-from codegen import Codebase
 from codegen.sdk.code_generation.current_code_codebase import get_graphsitter_repo_path
+from codegen.sdk.core.codebase import Codebase
 
 
 def test_codegen_imports():
diff --git a/tests/integration/extension/test_github.py b/tests/integration/extension/test_github.py
index a7ced0a3e..60e0a32e6 100644
--- a/tests/integration/extension/test_github.py
+++ b/tests/integration/extension/test_github.py
@@ -4,9 +4,9 @@
 
 import pytest
 
-from codegen import Codebase
 from codegen.extensions.linear.linear_client import LinearClient
 from codegen.extensions.tools.github import view_pr
+from codegen.sdk.core.codebase import Codebase
 
 
 @pytest.fixture