refactor: Use @graph_sitter (#49)

- **rename to CliSession**
- **update to use @graph_sitter**

---------

Co-authored-by: Codegen Team (Internal) <[email protected]>

Author: bagel897

.codegen/codemods/test_language/test_language.py

@@ -1,9 +1,9 @@
-import codegen
-from codegen.sdk.core.codebase import Codebase
-from codegen.shared.enums.programming_language import ProgrammingLanguage
+import graph_sitter
+from graph_sitter.core.codebase import Codebase
+from graph_sitter.shared.enums.programming_language import ProgrammingLanguage
 
 
-@codegen.function("test-language", subdirectories=["src/codegen/cli"], language=ProgrammingLanguage.PYTHON)
+@graph_sitter.function("test-language", subdirectories=["src/codegen/cli"], language=ProgrammingLanguage.PYTHON)
 def run(codebase: Codebase):
     file = codebase.get_file("src/codegen/cli/errors.py")
     print(f"File: {file.path}")

.github/workflows/test.yml

@@ -23,6 +23,10 @@ jobs:
   unit-tests:
     needs: access-check
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        group: [1, 2, 3, 4, 5, 6, 7, 8]
+
     steps:
       - uses: actions/checkout@v4
         with:
@@ -38,9 +42,10 @@ jobs:
           uv run pytest \
             -n auto \
             --cov src \
-            --cov-report=json \
+            --splits 8 --group ${{ matrix.group }} \
+            # --cov-report=json \
             --timeout 15 \
-            -o junit_suite_name="${{github.job}}" \
+            -o junit_suite_name="${{github.job}}-${{ matrix.group }}" \
             tests/unit
 
       - uses: ./.github/actions/report

docs/building-with-graph-sitter/function-decorator.mdx

@@ -1,6 +1,6 @@
 ---
 title: Function Decorator
-sidebarTitle: "@codegen.function"
+sidebarTitle: "@graph_sitter.function"
 icon: "at"
 iconType: "solid"
 ---
@@ -11,17 +11,17 @@ The `function` decorator is used to define codegen functions within your applica
 
 ## Usage
 
-To use the `function` decorator, simply annotate your function with `@codegen.function` and provide a name as an argument.
+To use the `function` decorator, simply annotate your function with `@graph_sitter.function` and provide a name as an argument.
 
 ### Example
 
 ```python
-@codegen.function('my-function')
+@graph_sitter.function('my-function')
 def run(codebase):
     pass
 ```
 
-In this example, the function `run` is decorated with `@codegen.function` and given the name `'my-function'`. This name will be used when the function is ran.
+In this example, the function `run` is decorated with `@graph_sitter.function` and given the name `'my-function'`. This name will be used when the function is ran.
 
 ## Parameters
 

docs/building-with-graph-sitter/reusable-codemods.mdx

@@ -5,7 +5,7 @@ icon: "arrows-rotate"
 iconType: "solid"
 ---
 
-Graph-sitter enables you to create reusable code transformations using Python functions decorated with `@codegen.function`. These codemods can be shared, versioned, and run by your team.
+Graph-sitter enables you to create reusable code transformations using Python functions decorated with `@graph_sitter.function`. These codemods can be shared, versioned, and run by your team.
 
 ## Creating Codemods
 
@@ -18,10 +18,10 @@ gs create rename-function .
 This creates a new codemod in your `.codegen/codemods` directory:
 
 ```python
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
-@codegen.function("rename-function")
+@graph_sitter.function("rename-function")
 def run(codebase: Codebase):
     """Add a description of what this codemod does."""
     # Add your code here
@@ -64,15 +64,15 @@ The execution flow:
 
 A codemod consists of three main parts:
 
-1. The `@codegen.function` decorator that names your codemod
+1. The `@graph_sitter.function` decorator that names your codemod
 2. A `run` function that takes a `Codebase` parameter
 3. Your transformation logic using the Codebase API
 
 ```python
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
-@codegen.function("update-imports")
+@graph_sitter.function("update-imports")
 def run(codebase: Codebase):
     """Update import statements to use new package names."""
     for file in codebase.files:
@@ -93,7 +93,7 @@ class RenameArgs(BaseModel):
     old_name: str
     new_name: str
 
-@codegen.function("rename-function")
+@graph_sitter.function("rename-function")
 def run(codebase: Codebase, arguments: RenameArgs):
     """Rename a function across the codebase."""
     old_func = codebase.get_function(arguments.old_name)

docs/cli/create.mdx

@@ -41,10 +41,10 @@ When you run `gs create rename-function .`, it creates:
 The generated codemod will have this structure:
 
 ```python
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
-@codegen.function("rename-function")
+@graph_sitter.function("rename-function")
 def run(codebase: Codebase):
     """Add a description of what this codemod does."""
     # Add your code here

docs/tutorials/training-data.mdx

@@ -31,7 +31,7 @@ First, we will do a "graph expansion" for each function - grab the function's so
 First, let's import the types we need from Codegen:
 
 ```python
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from codegen.sdk.core.external_module import ExternalModule
 from codegen.sdk.core.import_resolution import Import

examples/STRUCTURE.md

@@ -25,7 +25,7 @@ Your `run.py` should follow this structure, demonstrated well in the `generate_t
 1. **Imports at the top**
 
    ```python
-   import codegen
+   import graph_sitter
    from graph_sitter import Codebase
    from codegen.sdk.core import Function
    # ... other imports
@@ -42,7 +42,7 @@ Your `run.py` should follow this structure, demonstrated well in the `generate_t
 1. **Main Graph-sitter function with decorator**
 
    ```python
-   @codegen.function("your-function-name")
+   @graph_sitter.function("your-function-name")
    def run(codebase: Codebase):
        """Clear docstring explaining what the function does.
 
@@ -98,7 +98,7 @@ Choose between these approaches based on:
 
 1. **Function Decorator**
 
-   - Always use `@codegen.function()` with a descriptive name
+   - Always use `@graph_sitter.function()` with a descriptive name
    - Name should match the example's purpose
 
 1. **Utility Functions**
@@ -136,7 +136,7 @@ def get_function_context(function) -> dict:
 
 
 # Main transformation with decorator
-@codegen.function("generate-training-data")
+@graph_sitter.function("generate-training-data")
 def run(codebase: Codebase):
     """Generate training data using a node2vec-like approach...
 

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

@@ -167,7 +167,7 @@ def codegen_system_prompt():
 ```python
 from graph_sitter import Codebase
 
-@codegen.function('{name}')
+@graph_sitter.function('{name}')
 def codemod(codebase: Codebase):
     for function in codebase.functions:
         if not function.usages:
@@ -355,9 +355,9 @@ async def run_codemod(
                 return {"error": "Invalid JSON in arguments parameter"}
 
         # Create a session for the codemod
-        from graph_sitter.cli.auth.session import CodegenSession
+        from graph_sitter.cli.auth.session import CliSession
 
-        session = CodegenSession(state.parsed_codebase.repo_path)
+        session = CliSession(state.parsed_codebase.repo_path)
         session.codebase = state.parsed_codebase
 
         # Capture output

examples/examples/cyclomatic_complexity/run.py

@@ -1,12 +1,12 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.core.statements.for_loop_statement import ForLoopStatement
 from graph_sitter.core.statements.if_block_statement import IfBlockStatement
 from graph_sitter.core.statements.try_catch_statement import TryCatchStatement
 from graph_sitter.core.statements.while_statement import WhileStatement
 
 
-@codegen.function("cyclomatic-complexity")
+@graph_sitter.function("cyclomatic-complexity")
 def run(codebase: Codebase):
     def calculate_cyclomatic_complexity(code_block):
         # Initialize cyclomatic complexity count

examples/examples/delete_dead_code/run.py

@@ -1,8 +1,8 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
 
-@codegen.function("delete-dead-code")
+@graph_sitter.function("delete-dead-code")
 def run(codebase: Codebase):
     removed_functions_count = 0
     removed_variables_count = 0