From 10f43d9e707db71e6ff2142cb811710dac32944a Mon Sep 17 00:00:00 2001
From: codegen-bot <team+codegenbot@codegen.sh>
Date: Mon, 27 Jan 2025 13:10:09 -0800
Subject: [PATCH 1/7] .

---
 docs/introduction/overview.mdx | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/docs/introduction/overview.mdx b/docs/introduction/overview.mdx
index 24a10d0a6..217005511 100644
--- a/docs/introduction/overview.mdx
+++ b/docs/introduction/overview.mdx
@@ -9,10 +9,12 @@ iconType: "solid"
 
 It provides a scriptable interface to a powerful, multi-lingual language server built on top of [Tree-sitter](https://tree-sitter.github.io/tree-sitter/).
 
-export const intoSnippet = `# Codegen builds a complete graph connecting
+export const metaCode = `# Codegen builds a complete graph connecting
 #  functions, classes, imports and their relationships
 from codegen import Codebase
 
+codebase = Codebase('./')
+
 # Work with code without dealing with syntax trees or parsing
 for function in codebase.functions:
     # Comprehensive static analysis for references, dependencies, etc.
@@ -21,13 +23,23 @@ for function in codebase.functions:
         function.remove()
 `
 
+export const code = `def foo():
+  pass
+
+def bar():
+  foo()
+
+def baz():
+  pass
+`
+
 <iframe
   width="100%"
   height="300px"
   scrolling="no"
   src={`https://chadcode.sh/embedded/codemod/?code=${encodeURIComponent(
-    intoSnippet
-  )}`}
+    metaCode
+  )}&input=${encodeURIComponent(code)}`}
   style={{
     backgroundColor: "#15141b",
   }}

From 276ae81f77c0bc7f4ef496457b78a20729146810 Mon Sep 17 00:00:00 2001
From: codegen-bot <team+codegenbot@codegen.sh>
Date: Mon, 27 Jan 2025 13:40:31 -0800
Subject: [PATCH 2/7] .

---
 docs/building-with-codegen/dot-codegen.mdx    | 29 ++++++++---
 src/codegen/cli/commands/create/main.py       | 16 +-----
 src/codegen/cli/commands/init/render.py       |  7 +--
 src/codegen/cli/commands/notebook/main.py     | 46 +++--------------
 src/codegen/cli/commands/run/main.py          | 11 ++++
 src/codegen/cli/utils/default_code.py         | 18 +++++++
 src/codegen/cli/utils/notebooks.py            | 44 ++++++++++++++++
 .../cli/workspace/initialize_workspace.py     | 50 ++++---------------
 src/codegen/cli/workspace/venv_manager.py     | 36 +++++++++----
 9 files changed, 144 insertions(+), 113 deletions(-)
 create mode 100644 src/codegen/cli/utils/default_code.py
 create mode 100644 src/codegen/cli/utils/notebooks.py

diff --git a/docs/building-with-codegen/dot-codegen.mdx b/docs/building-with-codegen/dot-codegen.mdx
index c4a4ca7a5..9e2488037 100644
--- a/docs/building-with-codegen/dot-codegen.mdx
+++ b/docs/building-with-codegen/dot-codegen.mdx
@@ -11,12 +11,11 @@ The `.codegen` directory contains your project's Codegen configuration, codemods
 
 ```bash
 .codegen/
-├── config.toml         # Project configuration
-├── codemods/          # Your codemod implementations
-├── jupyter/           # Jupyter notebooks for exploration
-├── docs/             # API documentation
-├── examples/         # Example code
-└── prompts/          # AI system prompts
+├── .venv/            # Python virtual environment (gitignored)
+├── config.toml       # Project configuration
+├── codemods/         # Your codemod implementations
+├── jupyter/          # Jupyter notebooks for exploration
+└── codegen-system-prompt.txt  # AI system prompt
 ```
 
 ## Initialization
@@ -31,6 +30,18 @@ codegen init [--fetch-docs] [--repo-name NAME] [--organization-name ORG]
 The `--fetch-docs` flag downloads API documentation and examples specific to your project's programming language.
 </Note>
 
+## Virtual Environment
+
+Codegen maintains its own virtual environment in `.codegen/.venv/` to ensure consistent package versions and isolation from your project's dependencies. This environment is:
+
+- Created using `uv` for fast, reliable package management
+- Initialized with Python 3.13
+- Automatically managed by Codegen commands
+- Used for running codemods and Jupyter notebooks
+- Gitignored to avoid committing environment-specific files
+
+The environment is created during `codegen init` and used by commands like `codegen run` and `codegen notebook`.
+
 ### Configuration
 
 The `config.toml` file stores your project settings:
@@ -52,10 +63,14 @@ Codegen automatically adds appropriate entries to your `.gitignore`:
 .codegen/prompts/
 .codegen/docs/
 .codegen/examples/
+.codegen/.venv/
+.codegen/jupyter/
+.codegen/codegen-system-prompt.txt
 ```
 
 <Info>
-While prompts, docs, and examples are ignored, your codemods in `.codegen/codemods/` are tracked in Git.
+- While most directories are ignored, your codemods in `.codegen/codemods/` are tracked in Git
+- The virtual environment and Jupyter notebooks are gitignored to avoid environment-specific issues
 </Info>
 
 ## Working with Codemods
diff --git a/src/codegen/cli/commands/create/main.py b/src/codegen/cli/commands/create/main.py
index 3fafe819e..14e34510c 100644
--- a/src/codegen/cli/commands/create/main.py
+++ b/src/codegen/cli/commands/create/main.py
@@ -12,6 +12,7 @@
 from codegen.cli.rich.pretty_print import pretty_print_error
 from codegen.cli.rich.spinners import create_spinner
 from codegen.cli.utils.constants import ProgrammingLanguage
+from codegen.cli.utils.default_code import DEFAULT_CODEMOD
 from codegen.cli.workspace.decorators import requires_init
 
 
@@ -62,19 +63,6 @@ def make_relative(path: Path) -> str:
         return f"./{path.name}"
 
 
-def get_default_code(name: str) -> str:
-    """Get the default function code without using the API."""
-    return f'''import codegen
-from codegen import Codebase
-
-@codegen.function("{name}")
-def run(codebase: Codebase):
-    """Add a description of what this codemod does."""
-    # Add your code here
-    pass
-'''
-
-
 @click.command(name="create")
 @requires_init
 @click.argument("name", type=str)
@@ -111,7 +99,7 @@ def create_command(session: CodegenSession, name: str, path: Path, description:
                     prompt_path.write_text(response.context)
         else:
             # Use default implementation
-            code = get_default_code(name)
+            code = DEFAULT_CODEMOD.format(name=name)
 
         # Create the target directory if needed
         target_path.parent.mkdir(parents=True, exist_ok=True)
diff --git a/src/codegen/cli/commands/init/render.py b/src/codegen/cli/commands/init/render.py
index 665db8246..d8244ac47 100644
--- a/src/codegen/cli/commands/init/render.py
+++ b/src/codegen/cli/commands/init/render.py
@@ -4,6 +4,7 @@
 def get_success_message(codegen_dir: Path, docs_dir: Path, examples_dir: Path) -> str:
     """Get the success message to display after initialization."""
     return """📁 .codegen configuration folder created:
-   [dim]config.toml[/dim]  Project configuration
-   [dim]codemods/[/dim]    Your codemod implementations
-   [dim]codegen-system-prompt.txt[/dim]     AI system prompt (gitignored)"""
+   [dim]config.toml[/dim]                Project configuration
+   [dim]codemods/[/dim]                  Your codemod implementations
+   [dim].venv/[/dim]                     Python virtual environment (gitignored)
+   [dim]codegen-system-prompt.txt[/dim]  AI system prompt (gitignored)"""
diff --git a/src/codegen/cli/commands/notebook/main.py b/src/codegen/cli/commands/notebook/main.py
index 29bf92375..81be1c370 100644
--- a/src/codegen/cli/commands/notebook/main.py
+++ b/src/codegen/cli/commands/notebook/main.py
@@ -8,6 +8,7 @@
 from codegen.cli.auth.decorators import requires_auth
 from codegen.cli.auth.session import CodegenSession
 from codegen.cli.rich.spinners import create_spinner
+from codegen.cli.utils.notebooks import create_notebook
 from codegen.cli.workspace.decorators import requires_init
 from codegen.cli.workspace.venv_manager import VenvManager
 
@@ -19,57 +20,24 @@ def create_jupyter_dir() -> Path:
     return jupyter_dir
 
 
-def create_notebook(jupyter_dir: Path) -> Path:
-    """Create a new Jupyter notebook if it doesn't exist."""
-    notebook_path = jupyter_dir / "tmp.ipynb"
-    if not notebook_path.exists():
-        notebook_content = {
-            "cells": [
-                {
-                    "cell_type": "code",
-                    "execution_count": None,
-                    "metadata": {},
-                    "outputs": [],
-                    "source": ["from codegen import Codebase\n", "\n", "# Initialize codebase\n", "codebase = Codebase('../../')\n"],
-                }
-            ],
-            "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}},
-            "nbformat": 4,
-            "nbformat_minor": 4,
-        }
-        import json
-
-        notebook_path.write_text(json.dumps(notebook_content, indent=2))
-    return notebook_path
-
-
 @click.command(name="notebook")
-@click.option("--background", is_flag=True, help="Run Jupyter Lab in the background")
 @requires_auth
 @requires_init
-def notebook_command(session: CodegenSession, background: bool = False):
+def notebook_command(session: CodegenSession):
     """Open a Jupyter notebook with the current codebase loaded."""
     with create_spinner("Setting up Jupyter environment...") as status:
         venv = VenvManager()
 
-        if not venv.is_initialized():
-            status.update("Creating virtual environment...")
-            venv.create_venv()
-
-        status.update("Installing required packages...")
-        venv.install_packages("codegen", "jupyterlab")
+        status.update("Checking Jupyter installation...")
+        venv.ensure_jupyter()
 
         jupyter_dir = create_jupyter_dir()
         notebook_path = create_notebook(jupyter_dir)
 
-        status.update("Starting Jupyter Lab...")
+        status.update("Running Jupyter Lab...")
 
         # Prepare the environment with the virtual environment activated
         env = {**os.environ, "VIRTUAL_ENV": str(venv.venv_dir), "PATH": f"{venv.venv_dir}/bin:{os.environ['PATH']}"}
 
-        # Start Jupyter Lab
-        if background:
-            subprocess.Popen(["jupyter", "lab", str(notebook_path)], env=env, start_new_session=True)
-        else:
-            # Run in foreground
-            subprocess.run(["jupyter", "lab", str(notebook_path)], env=env, check=True)
+        # Run Jupyter Lab
+        subprocess.run(["jupyter", "lab", str(notebook_path)], env=env, check=True)
diff --git a/src/codegen/cli/commands/run/main.py b/src/codegen/cli/commands/run/main.py
index 357f1f380..161814473 100644
--- a/src/codegen/cli/commands/run/main.py
+++ b/src/codegen/cli/commands/run/main.py
@@ -1,4 +1,5 @@
 import json
+import os
 
 import rich_click as click
 
@@ -6,6 +7,7 @@
 from codegen.cli.utils.codemod_manager import CodemodManager
 from codegen.cli.utils.json_schema import validate_json
 from codegen.cli.workspace.decorators import requires_init
+from codegen.cli.workspace.venv_manager import VenvManager
 
 
 @click.command(name="run")
@@ -22,6 +24,15 @@ def run_command(
     arguments: str | None = None,
 ):
     """Run a codegen function by its label."""
+    # Ensure venv is initialized
+    venv = VenvManager()
+    if not venv.is_initialized():
+        raise click.ClickException("Virtual environment not found. Please run 'codegen init' first.")
+
+    # Set up environment with venv
+    os.environ["VIRTUAL_ENV"] = str(venv.venv_dir)
+    os.environ["PATH"] = f"{venv.venv_dir}/bin:{os.environ['PATH']}"
+
     # Get and validate the codemod
     codemod = CodemodManager.get_codemod(label)
 
diff --git a/src/codegen/cli/utils/default_code.py b/src/codegen/cli/utils/default_code.py
new file mode 100644
index 000000000..804ac1a53
--- /dev/null
+++ b/src/codegen/cli/utils/default_code.py
@@ -0,0 +1,18 @@
+DEFAULT_CODEMOD = '''import codegen
+from codegen import Codebase
+
+@codegen.function("{name}")
+def run(codebase: Codebase):
+    """Add a description of what this codemod does."""
+    # Add your code here
+    print('Total files: ', len(codebase.files))
+    print('Total functions: ', len(codebase.functions))
+    print('Total imports: ', len(codebase.imports))
+
+if __name__ == "__main__":
+    print('Parsing codebase...')
+    codebase = Codebase("./")
+
+    print('Running...')
+    run(codebase)
+'''
diff --git a/src/codegen/cli/utils/notebooks.py b/src/codegen/cli/utils/notebooks.py
new file mode 100644
index 000000000..c4b047ebf
--- /dev/null
+++ b/src/codegen/cli/utils/notebooks.py
@@ -0,0 +1,44 @@
+import json
+from pathlib import Path
+
+DEFAULT_NOTEBOOK_CONTENT = """from codegen import Codebase
+
+# Initialize codebase
+codebase = Codebase('../../')
+
+# Print out stats
+print("🔍 Codebase Analysis")
+print("=" * 50)
+print(f"📚 Total Files: {len(codebase.files)}")
+print(f"⚡ Total Functions: {len(codebase.functions)}")
+print(f"🔄 Total Imports: {len(codebase.imports)}")
+""".strip()
+
+
+def create_notebook(jupyter_dir: Path) -> Path:
+    """Create a new Jupyter notebook if it doesn't exist.
+
+    Args:
+        jupyter_dir: Directory where the notebook should be created
+
+    Returns:
+        Path to the created or existing notebook
+    """
+    notebook_path = jupyter_dir / "tmp.ipynb"
+    if not notebook_path.exists():
+        notebook_content = {
+            "cells": [
+                {
+                    "cell_type": "code",
+                    "execution_count": None,
+                    "metadata": {},
+                    "outputs": [],
+                    "source": DEFAULT_NOTEBOOK_CONTENT,
+                }
+            ],
+            "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}},
+            "nbformat": 4,
+            "nbformat_minor": 4,
+        }
+        notebook_path.write_text(json.dumps(notebook_content, indent=2))
+    return notebook_path
diff --git a/src/codegen/cli/workspace/initialize_workspace.py b/src/codegen/cli/workspace/initialize_workspace.py
index bfcdbcf78..e57d7ee50 100644
--- a/src/codegen/cli/workspace/initialize_workspace.py
+++ b/src/codegen/cli/workspace/initialize_workspace.py
@@ -13,46 +13,10 @@
 from codegen.cli.git.repo import get_git_repo
 from codegen.cli.git.url import get_git_organization_and_repo
 from codegen.cli.rich.spinners import create_spinner
+from codegen.cli.utils.notebooks import create_notebook
 from codegen.cli.workspace.docs_workspace import populate_api_docs
 from codegen.cli.workspace.examples_workspace import populate_examples
-
-DEFAULT_CODE = """
-from codegen import Codebase
-
-# Initialize codebase
-codebase = Codebase('../../')
-
-# Print out stats
-print("🔍 Codebase Analysis")
-print("=" * 50)
-print(f"📚 Total Files: {len(codebase.files)}")
-print(f"⚡ Total Functions: {len(codebase.functions)}")
-print(f"🔄 Total Imports: {len(codebase.imports)}")
-"""
-
-
-def create_notebook(jupyter_dir: Path) -> Path:
-    """Create a new Jupyter notebook if it doesn't exist."""
-    notebook_path = jupyter_dir / "tmp.ipynb"
-    if not notebook_path.exists():
-        notebook_content = {
-            "cells": [
-                {
-                    "cell_type": "code",
-                    "execution_count": None,
-                    "metadata": {},
-                    "outputs": [],
-                    "source": [DEFAULT_CODE],
-                }
-            ],
-            "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}},
-            "nbformat": 4,
-            "nbformat_minor": 4,
-        }
-        import json
-
-        notebook_path.write_text(json.dumps(notebook_content, indent=2))
-    return notebook_path
+from codegen.cli.workspace.venv_manager import VenvManager
 
 
 def initialize_codegen(
@@ -93,6 +57,13 @@ def initialize_codegen(
         JUPYTER_DIR.mkdir(parents=True, exist_ok=True)
         CODEMODS_DIR.mkdir(parents=True, exist_ok=True)
 
+        # Initialize virtual environment
+        status_obj.update(f"   {'Creating' if isinstance(status, str) else 'Checking'} virtual environment...")
+        venv = VenvManager()
+        if not venv.is_initialized():
+            venv.create_venv()
+            venv.install_packages("codegen")
+
         # Download system prompt
         try:
             from codegen.cli.api.endpoints import CODEGEN_SYSTEM_PROMPT_URL
@@ -164,7 +135,8 @@ def modify_gitignore(codegen_folder: Path):
         "examples/",
         "prompts/",
         "jupyter/",
-        "codegen-system-prompt.txt",  # Add system prompt to gitignore
+        ".venv/",  # Add venv to gitignore
+        "codegen-system-prompt.txt",
         "",
         "# Python cache files",
         "__pycache__/",
diff --git a/src/codegen/cli/workspace/venv_manager.py b/src/codegen/cli/workspace/venv_manager.py
index cb137d060..6dc1bd86e 100644
--- a/src/codegen/cli/workspace/venv_manager.py
+++ b/src/codegen/cli/workspace/venv_manager.py
@@ -3,19 +3,21 @@
 from pathlib import Path
 
 from codegen.cli.auth.constants import CODEGEN_DIR
-from codegen.cli.git.repo import get_git_repo
 
 
 class VenvManager:
+    """Manages the virtual environment for codegen."""
+
     def __init__(self):
-        repo = get_git_repo()
-        self.repo_path = Path(repo.workdir)
-        self.codegen_dir = self.repo_path / CODEGEN_DIR
+        self.codegen_dir = Path.cwd() / CODEGEN_DIR
         self.venv_dir = self.codegen_dir / ".venv"
-        self.python_path = self.venv_dir / "bin" / "python"
-        self.pip_path = self.venv_dir / "bin" / "pip"
 
-    def create_venv(self, python_version: str = "3.12") -> None:
+    def is_initialized(self) -> bool:
+        """Check if virtual environment exists."""
+        python_path = self.venv_dir / "bin" / "python"
+        return self.venv_dir.exists() and python_path.exists()
+
+    def create_venv(self, python_version: str = "3.13"):
         """Create a virtual environment using uv."""
         self.codegen_dir.mkdir(parents=True, exist_ok=True)
         subprocess.run(
@@ -23,7 +25,7 @@ def create_venv(self, python_version: str = "3.12") -> None:
             check=True,
         )
 
-    def install_packages(self, *packages: str) -> None:
+    def install_packages(self, *packages: str):
         """Install packages into the virtual environment using uv pip."""
         subprocess.run(
             ["uv", "pip", "install", *packages],
@@ -31,6 +33,11 @@ def install_packages(self, *packages: str) -> None:
             env={**os.environ, "VIRTUAL_ENV": str(self.venv_dir)},
         )
 
+    def run_python(self, script: str, *args: str):
+        """Run a Python script in the virtual environment."""
+        python_path = self.venv_dir / "bin" / "python"
+        subprocess.run([str(python_path), "-c", script, *args], check=True)
+
     def get_activate_command(self) -> str:
         """Get the command to activate the virtual environment."""
         return f"source {self.venv_dir}/bin/activate"
@@ -39,6 +46,13 @@ def is_active(self) -> bool:
         """Check if a virtual environment is active."""
         return "VIRTUAL_ENV" in os.environ
 
-    def is_initialized(self) -> bool:
-        """Check if the virtual environment exists and is properly set up."""
-        return self.venv_dir.exists() and self.python_path.exists()
+    def ensure_jupyter(self):
+        """Ensure Jupyter Lab is installed in the virtual environment."""
+        try:
+            subprocess.run(
+                [str(self.venv_dir / "bin" / "jupyter"), "--version"],
+                check=True,
+                capture_output=True,
+            )
+        except (subprocess.CalledProcessError, FileNotFoundError):
+            self.install_packages("jupyterlab")

From d22a442b7a768632f4d9d695227b786170896b68 Mon Sep 17 00:00:00 2001
From: codegen-bot <team+codegenbot@codegen.sh>
Date: Mon, 27 Jan 2025 13:41:51 -0800
Subject: [PATCH 3/7] .

---
 docs/building-with-codegen/dot-codegen.mdx | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/docs/building-with-codegen/dot-codegen.mdx b/docs/building-with-codegen/dot-codegen.mdx
index 9e2488037..72942b247 100644
--- a/docs/building-with-codegen/dot-codegen.mdx
+++ b/docs/building-with-codegen/dot-codegen.mdx
@@ -42,6 +42,8 @@ Codegen maintains its own virtual environment in `.codegen/.venv/` to ensure con
 
 The environment is created during `codegen init` and used by commands like `codegen run` and `codegen notebook`.
 
+<Note>To debug codemods, you will need to set the python virtual environment in your IDE to `.codegen/.venv`</Note>
+
 ### Configuration
 
 The `config.toml` file stores your project settings:

From a153ccfecefb5e4c1d051d6c1cc78b3a3ce28733 Mon Sep 17 00:00:00 2001
From: codegen-bot <team+codegenbot@codegen.sh>
Date: Mon, 27 Jan 2025 13:51:45 -0800
Subject: [PATCH 4/7] .

---
 docs/building-with-codegen/dot-codegen.mdx | 6 ++----
 src/codegen/cli/utils/default_code.py      | 2 ++
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/docs/building-with-codegen/dot-codegen.mdx b/docs/building-with-codegen/dot-codegen.mdx
index 72942b247..a555e0e60 100644
--- a/docs/building-with-codegen/dot-codegen.mdx
+++ b/docs/building-with-codegen/dot-codegen.mdx
@@ -62,16 +62,14 @@ Codegen automatically adds appropriate entries to your `.gitignore`:
 
 ```gitignore
 # Codegen
-.codegen/prompts/
-.codegen/docs/
-.codegen/examples/
 .codegen/.venv/
+.codegen/docs/
 .codegen/jupyter/
 .codegen/codegen-system-prompt.txt
 ```
 
 <Info>
-- While most directories are ignored, your codemods in `.codegen/codemods/` are tracked in Git
+- While most directories are ignored, your codemods in `.codegen/codemods/` and `config.toml` are tracked in Git
 - The virtual environment and Jupyter notebooks are gitignored to avoid environment-specific issues
 </Info>
 
diff --git a/src/codegen/cli/utils/default_code.py b/src/codegen/cli/utils/default_code.py
index 804ac1a53..2bed33f80 100644
--- a/src/codegen/cli/utils/default_code.py
+++ b/src/codegen/cli/utils/default_code.py
@@ -1,6 +1,7 @@
 DEFAULT_CODEMOD = '''import codegen
 from codegen import Codebase
 
+
 @codegen.function("{name}")
 def run(codebase: Codebase):
     """Add a description of what this codemod does."""
@@ -9,6 +10,7 @@ def run(codebase: Codebase):
     print('Total functions: ', len(codebase.functions))
     print('Total imports: ', len(codebase.imports))
 
+
 if __name__ == "__main__":
     print('Parsing codebase...')
     codebase = Codebase("./")

From 094c6c144d1f2a0704525b5fa9ac60fc21818512 Mon Sep 17 00:00:00 2001
From: codegen-bot <team+codegenbot@codegen.sh>
Date: Mon, 27 Jan 2025 16:16:55 -0800
Subject: [PATCH 5/7] makes CREATE work

---
 src/codegen/cli/api/client.py           |  2 +-
 src/codegen/cli/api/schemas.py          |  7 +++---
 src/codegen/cli/auth/session.py         |  8 ++++++
 src/codegen/cli/codemod/convert.py      | 20 ++++++++-------
 src/codegen/cli/commands/create/main.py | 33 +++++++++++--------------
 5 files changed, 38 insertions(+), 32 deletions(-)

diff --git a/src/codegen/cli/api/client.py b/src/codegen/cli/api/client.py
index bef1fcbde..ff81459ab 100644
--- a/src/codegen/cli/api/client.py
+++ b/src/codegen/cli/api/client.py
@@ -174,7 +174,7 @@ def create(self, name: str, query: str) -> CreateResponse:
         return self._make_request(
             "GET",
             CREATE_ENDPOINT,
-            CreateInput(input=CreateInput.BaseCreateInput(name=name, query=query, repo_full_name=session.repo_name)),
+            CreateInput(input=CreateInput.BaseCreateInput(name=name, query=query, language=session.language)),
             CreateResponse,
         )
 
diff --git a/src/codegen/cli/api/schemas.py b/src/codegen/cli/api/schemas.py
index 34583200e..9fffd418e 100644
--- a/src/codegen/cli/api/schemas.py
+++ b/src/codegen/cli/api/schemas.py
@@ -92,8 +92,8 @@ class DocsResponse(SafeBaseModel):
 class CreateInput(SafeBaseModel):
     class BaseCreateInput(SafeBaseModel):
         name: str
-        query: str | None = None
-        repo_full_name: str | None = None
+        query: str
+        language: ProgrammingLanguage
 
     input: BaseCreateInput
 
@@ -102,8 +102,7 @@ class CreateResponse(SafeBaseModel):
     success: bool
     response: str
     code: str
-    codemod_id: int
-    context: str | None = None
+    context: str
 
 
 ###########################################################################
diff --git a/src/codegen/cli/auth/session.py b/src/codegen/cli/auth/session.py
index bdf92a8f6..07dddb3b1 100644
--- a/src/codegen/cli/auth/session.py
+++ b/src/codegen/cli/auth/session.py
@@ -8,6 +8,7 @@
 from codegen.cli.errors import AuthError, NoTokenError
 from codegen.cli.git.repo import get_git_repo
 from codegen.cli.utils.config import Config, get_config, write_config
+from codegen.sdk.enums import ProgrammingLanguage
 
 
 @dataclass
@@ -109,6 +110,13 @@ def repo_name(self) -> str:
         """Get the current repository name"""
         return self.config.repo_full_name
 
+    @property
+    def language(self) -> ProgrammingLanguage:
+        """Get the current language"""
+        # TODO(jayhack): This is a temporary solution to get the language.
+        # We should eventually get the language on init.
+        return self.config.programming_language or ProgrammingLanguage.PYTHON
+
     @property
     def codegen_dir(self) -> Path:
         """Get the path to the  codegen-sh directory"""
diff --git a/src/codegen/cli/codemod/convert.py b/src/codegen/cli/codemod/convert.py
index 512eeebd4..15c80f4d5 100644
--- a/src/codegen/cli/codemod/convert.py
+++ b/src/codegen/cli/codemod/convert.py
@@ -2,19 +2,21 @@
 
 
 def convert_to_cli(input: str, language: str, name: str) -> str:
-    codebase_type = "PyCodebaseType" if language.lower() == "python" else "TSCodebaseType"
-    return f"""import codegen.cli.sdk.decorator
-# from app.codemod.compilation.models.context import CodemodContext
-#from app.codemod.compilation.models.pr_options import PROptions
+    return f"""import codegen
+from codegen import Codebase
 
-from codegen.sdk import {codebase_type}
 
-context: Any
+@codegen.function('{name}')
+def run(codebase: Codebase):
+{indent(input, "    ")}
 
 
-@codegen.cli.sdk.decorator.function('{name}')
-def run(codebase: {codebase_type}, pr_options: Any):
-{indent(input, "    ")}
+if __name__ == "__main__":
+    print('Parsing codebase...')
+    codebase = Codebase("./")
+
+    print('Running function...')
+    codegen.run(run)
 """
 
 
diff --git a/src/codegen/cli/commands/create/main.py b/src/codegen/cli/commands/create/main.py
index 14e34510c..390276a2d 100644
--- a/src/codegen/cli/commands/create/main.py
+++ b/src/codegen/cli/commands/create/main.py
@@ -11,7 +11,6 @@
 from codegen.cli.rich.codeblocks import format_command, format_path
 from codegen.cli.rich.pretty_print import pretty_print_error
 from codegen.cli.rich.spinners import create_spinner
-from codegen.cli.utils.constants import ProgrammingLanguage
 from codegen.cli.utils.default_code import DEFAULT_CODEMOD
 from codegen.cli.workspace.decorators import requires_init
 
@@ -29,7 +28,7 @@ def get_prompts_dir() -> Path:
     return PROMPTS_DIR
 
 
-def get_target_path(name: str, path: Path) -> Path:
+def get_target_paths(name: str, path: Path) -> tuple[Path, Path]:
     """Get the target path for the new function file.
 
     Creates a directory structure like:
@@ -47,7 +46,9 @@ def get_target_path(name: str, path: Path) -> Path:
     # Create path within .codegen/codemods
     codemods_dir = base_dir / ".codegen" / "codemods"
     function_dir = codemods_dir / name_snake
-    return function_dir / f"{name_snake}.py"
+    codemod_path = function_dir / f"{name_snake}.py"
+    prompt_path = function_dir / f"{name_snake}-system-prompt.txt"
+    return codemod_path, prompt_path
 
 
 def make_relative(path: Path) -> str:
@@ -76,11 +77,11 @@ def create_command(session: CodegenSession, name: str, path: Path, description:
     PATH is where to create the function (default: current directory)
     """
     # Get the target path for the function
-    target_path = get_target_path(name, path)
+    codemod_path, prompt_path = get_target_paths(name, path)
 
     # Check if file exists
-    if target_path.exists() and not overwrite:
-        rel_path = make_relative(target_path)
+    if codemod_path.exists() and not overwrite:
+        rel_path = make_relative(codemod_path)
         pretty_print_error(f"File already exists at {format_path(rel_path)}\n\nTo overwrite the file:\n{format_command(f'codegen create {name} {rel_path} --overwrite')}")
         return
 
@@ -89,34 +90,30 @@ def create_command(session: CodegenSession, name: str, path: Path, description:
     try:
         if description:
             # Use API to generate implementation
-            with create_spinner("Generating function (using LLM, this will take ~30s)") as status:
+            with create_spinner("Generating function (using LLM, this will take ~10s)") as status:
                 response = RestAPI(session.token).create(name=name, query=description)
-                code = convert_to_cli(response.code, session.config.programming_language or ProgrammingLanguage.PYTHON, name)
-
-                # Write the system prompt if provided
-                if response.context:
-                    prompt_path = get_prompts_dir() / f"{name.lower().replace(' ', '-')}-system-prompt.md"
-                    prompt_path.write_text(response.context)
+                code = convert_to_cli(response.code, session.language, name)
         else:
             # Use default implementation
             code = DEFAULT_CODEMOD.format(name=name)
 
         # Create the target directory if needed
-        target_path.parent.mkdir(parents=True, exist_ok=True)
+        codemod_path.parent.mkdir(parents=True, exist_ok=True)
 
         # Write the function code
-        target_path.write_text(code)
+        codemod_path.write_text(code)
+        prompt_path.write_text(response.context)
 
     except (ServerError, ValueError) as e:
         raise click.ClickException(str(e))
 
     # Success message
-    rich.print(f"\n✅ {'Overwrote' if overwrite and target_path.exists() else 'Created'} function '{name}'")
+    rich.print(f"\n✅ {'Overwrote' if overwrite and codemod_path.exists() else 'Created'} function '{name}'")
     rich.print("")
     rich.print("📁 Files Created:")
-    rich.print(f"   [dim]Function:[/dim]  {make_relative(target_path)}")
+    rich.print(f"   [dim]Function:[/dim]  {make_relative(codemod_path)}")
     if description and response.context:
-        rich.print(f"   [dim]Prompt:[/dim]    {make_relative(get_prompts_dir() / f'{name.lower().replace(" ", "-")}-system-prompt.md')}")
+        rich.print(f"   [dim]Prompt:[/dim]    {make_relative(prompt_path)}")
 
     # Next steps
     rich.print("\n[bold]What's next?[/bold]\n")

From bfdd541fc25c3112e39fbb38583271230e78ca98 Mon Sep 17 00:00:00 2001
From: codegen-bot <team+codegenbot@codegen.sh>
Date: Mon, 27 Jan 2025 16:24:39 -0800
Subject: [PATCH 6/7] .

---
 docs/cli/create.mdx   | 15 ++++++++++++---
 docs/cli/notebook.mdx |  2 +-
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/docs/cli/create.mdx b/docs/cli/create.mdx
index 5b7b5dd55..4a6e14b7d 100644
--- a/docs/cli/create.mdx
+++ b/docs/cli/create.mdx
@@ -33,8 +33,8 @@ When you run `codegen create rename-function`, it creates:
 .codegen/
 └── codemods/
     └── rename_function/
-        ├── rename_function.py       # The codemod implementation
-        └── rename_function_prompt.md  # System prompt (if --description used)
+        ├── rename_function.py              # The codemod implementation
+        └── rename_function-system-prompt.txt  # System prompt (if --description used)
 ```
 
 The generated codemod will have this structure:
@@ -47,7 +47,16 @@ from codegen import Codebase
 def run(codebase: Codebase):
     """Add a description of what this codemod does."""
     # Add your code here
-    pass
+    print('Total files: ', len(codebase.files))
+    print('Total functions: ', len(codebase.functions))
+    print('Total imports: ', len(codebase.imports))
+
+if __name__ == "__main__":
+    print('Parsing codebase...')
+    codebase = Codebase("./")
+
+    print('Running...')
+    run(codebase)
 ```
 
 ## Examples
diff --git a/docs/cli/notebook.mdx b/docs/cli/notebook.mdx
index 343856970..4941f1a94 100644
--- a/docs/cli/notebook.mdx
+++ b/docs/cli/notebook.mdx
@@ -1,5 +1,5 @@
 ---
-title: "notebook"
+title: "Notebook Command"
 sidebarTitle: "notebook"
 description: "Open a Jupyter notebook with the current codebase loaded"
 icon: "book"

From d862efcbc60c196c4bdedbc5a1e3161f6434fb2c Mon Sep 17 00:00:00 2001
From: codegen-bot <team+codegenbot@codegen.sh>
Date: Mon, 27 Jan 2025 16:36:16 -0800
Subject: [PATCH 7/7] .

---
 README.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index faae1a232..7fad15c74 100644
--- a/README.md
+++ b/README.md
@@ -16,9 +16,9 @@
 
 <div align="center">
 
-[![PyPI](https://img.shields.io/pypi/v/codegen?style=flat-square&color=blue)](https://pypi.org/project/codegen/) 
-[![Documentation](https://img.shields.io/badge/docs-docs.codegen.com-purple?style=flat-square)](https://docs.codegen.com) 
-[![Slack Community](https://img.shields.io/badge/slack-community-4A154B?logo=slack&style=flat-square)](https://community.codegen.com) 
+[![PyPI](https://img.shields.io/pypi/v/codegen?style=flat-square&color=blue)](https://pypi.org/project/codegen/)
+[![Documentation](https://img.shields.io/badge/docs-docs.codegen.com-purple?style=flat-square)](https://docs.codegen.com)
+[![Slack Community](https://img.shields.io/badge/slack-community-4A154B?logo=slack&style=flat-square)](https://community.codegen.com)
 [![Follow on X](https://img.shields.io/twitter/follow/codegen?style=social)](https://x.com/codegen)
 
 </div>