diff --git a/.codegen/codemods/test_language/test_language.py b/.codegen/codemods/test_language/test_language.py
index 19ae4c0bd..30323d73c 100644
--- a/.codegen/codemods/test_language/test_language.py
+++ b/.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}")
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 6aa953d41..4a9b8f68b 100644
--- a/.github/workflows/test.yml
+++ b/.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
diff --git a/docs/building-with-graph-sitter/function-decorator.mdx b/docs/building-with-graph-sitter/function-decorator.mdx
index 0010afa66..4ded11a7c 100644
--- a/docs/building-with-graph-sitter/function-decorator.mdx
+++ b/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
 
diff --git a/docs/building-with-graph-sitter/reusable-codemods.mdx b/docs/building-with-graph-sitter/reusable-codemods.mdx
index 0e5602a13..c5cf6344f 100644
--- a/docs/building-with-graph-sitter/reusable-codemods.mdx
+++ b/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)
diff --git a/docs/cli/create.mdx b/docs/cli/create.mdx
index eca5b28e0..30d24fb46 100644
--- a/docs/cli/create.mdx
+++ b/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
diff --git a/docs/tutorials/training-data.mdx b/docs/tutorials/training-data.mdx
index 79784136f..ce14a7c1b 100644
--- a/docs/tutorials/training-data.mdx
+++ b/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
diff --git a/examples/STRUCTURE.md b/examples/STRUCTURE.md
index 91dc6ee3c..c27deb1a5 100644
--- a/examples/STRUCTURE.md
+++ b/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...
 
diff --git a/examples/examples/codegen-mcp-server/server.py b/examples/examples/codegen-mcp-server/server.py
index ddc7bb8b6..f990329c5 100644
--- a/examples/examples/codegen-mcp-server/server.py
+++ b/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
diff --git a/examples/examples/cyclomatic_complexity/run.py b/examples/examples/cyclomatic_complexity/run.py
index 04f1e4b5f..d396dfafb 100644
--- a/examples/examples/cyclomatic_complexity/run.py
+++ b/examples/examples/cyclomatic_complexity/run.py
@@ -1,4 +1,4 @@
-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
@@ -6,7 +6,7 @@
 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
diff --git a/examples/examples/delete_dead_code/run.py b/examples/examples/delete_dead_code/run.py
index 95c6e8c8f..c0b1f410d 100644
--- a/examples/examples/delete_dead_code/run.py
+++ b/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
diff --git a/examples/examples/dict_to_schema/run.py b/examples/examples/dict_to_schema/run.py
index ba9f86f7c..06640885d 100644
--- a/examples/examples/dict_to_schema/run.py
+++ b/examples/examples/dict_to_schema/run.py
@@ -1,8 +1,8 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
 
-@codegen.function("dict-to-pydantic-schema")
+@graph_sitter.function("dict-to-pydantic-schema")
 def run(codebase: Codebase):
     """Convert dictionary literals to Pydantic models in a Python codebase.
 
diff --git a/examples/examples/document_functions/run.py b/examples/examples/document_functions/run.py
index f2ad132f9..53c2e46e6 100644
--- a/examples/examples/document_functions/run.py
+++ b/examples/examples/document_functions/run.py
@@ -1,4 +1,4 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.core.external_module import ExternalModule
 from graph_sitter.core.import_resolution import Import
@@ -54,7 +54,7 @@ def get_extended_context(symbol: Symbol, degree: int) -> tuple[set[Symbol], set[
     return dependencies, usages
 
 
-@codegen.function("document-functions")
+@graph_sitter.function("document-functions")
 def run(codebase: Codebase):
     # Define the maximum degree of dependencies and usages to consider for context
     N_DEGREE = 2
diff --git a/examples/examples/fragment_to_shorthand/run.py b/examples/examples/fragment_to_shorthand/run.py
index 861d9f5a7..311864f20 100644
--- a/examples/examples/fragment_to_shorthand/run.py
+++ b/examples/examples/fragment_to_shorthand/run.py
@@ -1,8 +1,8 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
 
-@codegen.function("fragment_to_shorthand")
+@graph_sitter.function("fragment_to_shorthand")
 def run(codebase: Codebase):
     print("🔍 Starting Fragment syntax conversion...")
 
diff --git a/examples/examples/freezegun_to_timemachine_migration/run.py b/examples/examples/freezegun_to_timemachine_migration/run.py
index 5182fd801..118cf2fe3 100644
--- a/examples/examples/freezegun_to_timemachine_migration/run.py
+++ b/examples/examples/freezegun_to_timemachine_migration/run.py
@@ -1,8 +1,8 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
 
-@codegen.function("freezegun-to-timemachine")
+@graph_sitter.function("freezegun-to-timemachine")
 def run(codebase: Codebase):
     """Convert FreezeGun usage to TimeMachine in test files.
 
diff --git a/examples/examples/generate_training_data/run.py b/examples/examples/generate_training_data/run.py
index 57ac60980..df13f8045 100644
--- a/examples/examples/generate_training_data/run.py
+++ b/examples/examples/generate_training_data/run.py
@@ -1,6 +1,6 @@
 import json
 
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.core.external_module import ExternalModule
 from graph_sitter.core.import_resolution import Import
@@ -42,7 +42,7 @@ def get_function_context(function) -> dict:
     return context
 
 
-@codegen.function("generate-training-data")
+@graph_sitter.function("generate-training-data")
 def run(codebase: Codebase):
     """Generate training data using a node2vec-like approach for code embeddings.
 
diff --git a/examples/examples/modules_dependencies/run.py b/examples/examples/modules_dependencies/run.py
index 094d8f761..284cceeb9 100644
--- a/examples/examples/modules_dependencies/run.py
+++ b/examples/examples/modules_dependencies/run.py
@@ -1,9 +1,9 @@
-import codegen
+import graph_sitter
 import networkx as nx
 from graph_sitter import Codebase
 
 
-@codegen.function("visualize-modules-dependencies")
+@graph_sitter.function("visualize-modules-dependencies")
 def run(codebase: Codebase):
     # Create a directed graph
     G = nx.DiGraph()
diff --git a/examples/examples/openapi_decorators/run.py b/examples/examples/openapi_decorators/run.py
index 2cf92c55c..608408de3 100644
--- a/examples/examples/openapi_decorators/run.py
+++ b/examples/examples/openapi_decorators/run.py
@@ -1,4 +1,4 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
 
@@ -169,7 +169,7 @@ def analyze_method_params(method) -> dict:
     return schema
 
 
-@codegen.function("add-openapi-decorators")
+@graph_sitter.function("add-openapi-decorators")
 def run(codebase: Codebase):
     """Add OpenAPI decorators (@response and @expect) to API endpoints."""
     analytics = {}
diff --git a/examples/examples/promises_to_async_await/run.py b/examples/examples/promises_to_async_await/run.py
index edc78918c..bbf8eb1d9 100644
--- a/examples/examples/promises_to_async_await/run.py
+++ b/examples/examples/promises_to_async_await/run.py
@@ -1,9 +1,9 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.core.statements.statement import StatementType
 
 
-@codegen.function("promises-to-async-await")
+@graph_sitter.function("promises-to-async-await")
 def run(codebase: Codebase):
     """Convert a repeated use of a promise chain to async await in the official twilio js client library twilio/twilio-node repository.
 
diff --git a/examples/examples/python2_to_python3/run.py b/examples/examples/python2_to_python3/run.py
index 904f77629..b7a2a3318 100644
--- a/examples/examples/python2_to_python3/run.py
+++ b/examples/examples/python2_to_python3/run.py
@@ -1,4 +1,4 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
 # Initialize codebase
@@ -115,7 +115,7 @@ def update_iterators(file):
                     stmt.edit(new_stmt)
 
 
-@codegen.function("python2-to-python3")
+@graph_sitter.function("python2-to-python3")
 def run():
     """Main function to run the Python 2 to 3 conversion"""
     print("🚀 Starting Python 2 to 3 conversion...\n")
diff --git a/examples/examples/reexport_management/run.py b/examples/examples/reexport_management/run.py
index d8eb7e512..4f4565562 100644
--- a/examples/examples/reexport_management/run.py
+++ b/examples/examples/reexport_management/run.py
@@ -1,11 +1,11 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.typescript.file import TSImport
 
 processed_imports = set()
 
 
-@codegen.function("reexport_management")
+@graph_sitter.function("reexport_management")
 def run(codebase: Codebase):
     print("🚀 Starting reexport analysis...")
     for file in codebase.files:
diff --git a/examples/examples/remove_default_exports/run.py b/examples/examples/remove_default_exports/run.py
index 6bf3581bb..50d7f3f9e 100644
--- a/examples/examples/remove_default_exports/run.py
+++ b/examples/examples/remove_default_exports/run.py
@@ -1,9 +1,9 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.typescript.file import TSFile
 
 
-@codegen.function("remove-default-exports")
+@graph_sitter.function("remove-default-exports")
 def run(codebase: Codebase):
     """Convert default exports to named exports in TypeScript files.
 
diff --git a/examples/examples/sqlalchemy_1.6_to_2.0/run.py b/examples/examples/sqlalchemy_1.6_to_2.0/run.py
index 34a3bcb3d..4b18ad8f6 100644
--- a/examples/examples/sqlalchemy_1.6_to_2.0/run.py
+++ b/examples/examples/sqlalchemy_1.6_to_2.0/run.py
@@ -1,10 +1,10 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.core.detached_symbols.function_call import FunctionCall
 from graph_sitter.core.expressions.chained_attribute import ChainedAttribute
 
 
-@codegen.function("sqlalchemy-1.4-to-2.0")
+@graph_sitter.function("sqlalchemy-1.4-to-2.0")
 def run(codebase: Codebase):
     """
     Convert SQLAlchemy 1.4 codebases to 2.0.
diff --git a/examples/examples/sqlalchemy_soft_delete/run.py b/examples/examples/sqlalchemy_soft_delete/run.py
index 3daa2fa24..12c33f870 100644
--- a/examples/examples/sqlalchemy_soft_delete/run.py
+++ b/examples/examples/sqlalchemy_soft_delete/run.py
@@ -2,7 +2,7 @@
 import subprocess
 from pathlib import Path
 
-import codegen
+import graph_sitter
 from graph_sitter.core.detached_symbols.function_call import FunctionCall
 
 
@@ -61,7 +61,7 @@ def clone_repo(repo_url: str, repo_path: Path) -> None:
     subprocess.run(["git", "clone", repo_url, str(repo_path)], check=True)
 
 
-@codegen.function("sqlalchemy-soft-delete")
+@graph_sitter.function("sqlalchemy-soft-delete")
 def process_soft_deletes(codebase):
     """Process soft delete conditions for join methods in the codebase."""
     soft_delete_models = {
diff --git a/examples/examples/sqlalchemy_type_annotations/run.py b/examples/examples/sqlalchemy_type_annotations/run.py
index 24f80d15e..bbbeedfc9 100644
--- a/examples/examples/sqlalchemy_type_annotations/run.py
+++ b/examples/examples/sqlalchemy_type_annotations/run.py
@@ -2,7 +2,7 @@
 import shutil
 import subprocess
 
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.core.detached_symbols.function_call import FunctionCall
 
@@ -21,7 +21,7 @@ def cleanup_git_repo(repo_path: str) -> None:
         shutil.rmtree(git_dir)
 
 
-@codegen.function("sqlalchemy-type-annotations")
+@graph_sitter.function("sqlalchemy-type-annotations")
 def run(codebase: Codebase):
     """Add Mapped types to SQLAlchemy models in a codebase.
 
diff --git a/examples/examples/unittest_to_pytest/run.py b/examples/examples/unittest_to_pytest/run.py
index f17b7a8ad..21f2312c3 100644
--- a/examples/examples/unittest_to_pytest/run.py
+++ b/examples/examples/unittest_to_pytest/run.py
@@ -1,4 +1,4 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 
 # Initialize codebase
@@ -53,7 +53,7 @@ def {fixture_name}():
             print(f"🗑️ Removed setUp method from class {cls.name}")
 
 
-@codegen.function("unittest-to-pytest")
+@graph_sitter.function("unittest-to-pytest")
 def run(codebase: Codebase):
     """Main function to run the unittest to pytest conversion"""
     print("🚀 Starting unittest to pytest conversion...")
diff --git a/examples/examples/usesuspensequery_to_usesuspensequeries/run.py b/examples/examples/usesuspensequery_to_usesuspensequeries/run.py
index 2da970eab..1e6ee2187 100644
--- a/examples/examples/usesuspensequery_to_usesuspensequeries/run.py
+++ b/examples/examples/usesuspensequery_to_usesuspensequeries/run.py
@@ -1,9 +1,9 @@
-import codegen
+import graph_sitter
 from graph_sitter import Codebase
 from graph_sitter.core.detached_symbols.function_call import FunctionCall
 
 
-@codegen.function("useSuspenseQuery-to-useSuspenseQueries")
+@graph_sitter.function("useSuspenseQuery-to-useSuspenseQueries")
 def run(codebase: Codebase):
     """Convert useSuspenseQuery calls to useSuspenseQueries in a React codebase.
 
diff --git a/examples/examples/visualize_codebases/blast_radius.py b/examples/examples/visualize_codebases/blast_radius.py
index 445c5716a..804418a4e 100644
--- a/examples/examples/visualize_codebases/blast_radius.py
+++ b/examples/examples/visualize_codebases/blast_radius.py
@@ -1,4 +1,4 @@
-import codegen
+import graph_sitter
 import networkx as nx
 from graph_sitter import Codebase
 from graph_sitter.core.dataclasses.usage import Usage
@@ -83,7 +83,7 @@ def create_blast_radius_visualization(symbol: PySymbol, depth: int = 0):
         create_blast_radius_visualization(usage_symbol, depth + 1)
 
 
-@codegen.function("visualize-function-blast-radius")
+@graph_sitter.function("visualize-function-blast-radius")
 def run(codebase: Codebase):
     """
     Generate a visualization showing the blast radius of changes to a function.
diff --git a/examples/examples/visualize_codebases/call_trace.py b/examples/examples/visualize_codebases/call_trace.py
index 5dc85a979..123795d83 100644
--- a/examples/examples/visualize_codebases/call_trace.py
+++ b/examples/examples/visualize_codebases/call_trace.py
@@ -1,4 +1,4 @@
-import codegen
+import graph_sitter
 import networkx as nx
 from graph_sitter import Codebase
 from graph_sitter.core.class_definition import Class
@@ -86,7 +86,7 @@ def create_downstream_call_trace(src_func: Function, depth: int = 0):
             create_downstream_call_trace(func, depth + 1)
 
 
-@codegen.function("visualize-function-call-relationships")
+@graph_sitter.function("visualize-function-call-relationships")
 def run(codebase: Codebase):
     """Generate a visualization of function call relationships in a codebase.
 
diff --git a/examples/examples/visualize_codebases/dependency_trace.py b/examples/examples/visualize_codebases/dependency_trace.py
index 64168839f..43831fbc6 100644
--- a/examples/examples/visualize_codebases/dependency_trace.py
+++ b/examples/examples/visualize_codebases/dependency_trace.py
@@ -1,4 +1,4 @@
-import codegen
+import graph_sitter
 import networkx as nx
 from graph_sitter import Codebase
 from graph_sitter.core.class_definition import Class
@@ -52,7 +52,7 @@ def create_dependencies_visualization(symbol: Symbol, depth: int = 0):
                 create_dependencies_visualization(dep_symbol, depth + 1)
 
 
-@codegen.function("visualize-symbol-dependencies")
+@graph_sitter.function("visualize-symbol-dependencies")
 def run(codebase: Codebase):
     """Generate a visualization of symbol dependencies in a codebase.
 
diff --git a/examples/examples/visualize_codebases/method_relationships.py b/examples/examples/visualize_codebases/method_relationships.py
index 1e1a5f837..bdfc8b068 100644
--- a/examples/examples/visualize_codebases/method_relationships.py
+++ b/examples/examples/visualize_codebases/method_relationships.py
@@ -1,4 +1,4 @@
-import codegen
+import graph_sitter
 import networkx as nx
 from graph_sitter import Codebase
 from graph_sitter.core.class_definition import Class
@@ -77,7 +77,7 @@ def create_downstream_call_trace(src_func: Function, depth: int = 0):
             create_downstream_call_trace(func, depth + 1)
 
 
-@codegen.function("visualize-class-method-relationships")
+@graph_sitter.function("visualize-class-method-relationships")
 def run(codebase: Codebase):
     """Generate a visualization of method call relationships within a class.
 
diff --git a/pyproject.toml b/pyproject.toml
index df14b5c5c..2f0dd094f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -160,6 +160,7 @@ dev-dependencies = [
   "modal>=0.73.25",
   "pytest-lsp>=1.0.0b1",
   "cython>=3.0.11",
+  "pytest-split>=0.10.0",
 ]
 
 [tool.uv.workspace]
diff --git a/src/graph_sitter/cli/auth/session.py b/src/graph_sitter/cli/auth/session.py
index 354e0916c..f9d138172 100644
--- a/src/graph_sitter/cli/auth/session.py
+++ b/src/graph_sitter/cli/auth/session.py
@@ -13,7 +13,7 @@
 from graph_sitter.git.repo_operator.local_git_repo import LocalGitRepo
 
 
-class CodegenSession:
+class CliSession:
     """Represents an authenticated codegen session with user and repository context"""
 
     repo_path: Path
@@ -38,7 +38,7 @@ def __init__(self, repo_path: Path, git_token: str | None = None) -> None:
         session_manager.set_active_session(repo_path)
 
     @classmethod
-    def from_active_session(cls) -> "CodegenSession | None":
+    def from_active_session(cls) -> "CliSession | None":
         active_session = session_manager.get_active_session()
         if not active_session:
             return None
@@ -84,4 +84,4 @@ def _validate(self) -> None:
             raise click.Abort()
 
     def __str__(self) -> str:
-        return f"CodegenSession(user={self.config.repository.user_name}, repo={self.config.repository.repo_name})"
+        return f"CliSession(user={self.config.repository.user_name}, repo={self.config.repository.repo_name})"
diff --git a/src/graph_sitter/cli/codemod/convert.py b/src/graph_sitter/cli/codemod/convert.py
index 8c99420f4..240812418 100644
--- a/src/graph_sitter/cli/codemod/convert.py
+++ b/src/graph_sitter/cli/codemod/convert.py
@@ -6,11 +6,11 @@ def convert_to_cli(input: str, language: str, name: str) -> str:
 # Run this codemod using `gs run {name}` OR the `run_codemod` MCP tool.
 # Important: if you run this as a regular python file, you MUST run it such that
 #  the base directory './' is the base of your codebase, otherwise it will not work.
-import codegen
+import graph_sitter
 from graph_sitter.core.codebase import Codebase
 
 
-@codegen.function('{name}')
+@graph_sitter.function('{name}')
 def run(codebase: Codebase):
 {indent(input, "    ")}
 
diff --git a/src/graph_sitter/cli/commands/create/main.py b/src/graph_sitter/cli/commands/create/main.py
index 80d8fab06..a1750ec17 100644
--- a/src/graph_sitter/cli/commands/create/main.py
+++ b/src/graph_sitter/cli/commands/create/main.py
@@ -3,7 +3,7 @@
 import rich
 import rich_click as click
 
-from graph_sitter.cli.auth.session import CodegenSession
+from graph_sitter.cli.auth.session import CliSession
 from graph_sitter.cli.errors import ServerError
 from graph_sitter.cli.rich.codeblocks import format_command, format_path
 from graph_sitter.cli.rich.pretty_print import pretty_print_error
@@ -52,7 +52,7 @@ def make_relative(path: Path) -> str:
 @click.argument("name", type=str)
 @click.argument("path", type=click.Path(path_type=Path), default=None)
 @click.option("--overwrite", is_flag=True, help="Overwrites function if it already exists.")
-def create_command(session: CodegenSession, name: str, path: Path | None, overwrite: bool = False):
+def create_command(session: CliSession, name: str, path: Path | None, overwrite: bool = False):
     """Create a new codegen function.
 
     NAME is the name/label for the function
diff --git a/src/graph_sitter/cli/commands/init/main.py b/src/graph_sitter/cli/commands/init/main.py
index 1271e743d..7a37bea5d 100644
--- a/src/graph_sitter/cli/commands/init/main.py
+++ b/src/graph_sitter/cli/commands/init/main.py
@@ -4,7 +4,7 @@
 import rich
 import rich_click as click
 
-from graph_sitter.cli.auth.session import CodegenSession
+from graph_sitter.cli.auth.session import CliSession
 from graph_sitter.cli.commands.init.render import get_success_message
 from graph_sitter.cli.rich.codeblocks import format_command
 from graph_sitter.cli.workspace.initialize_workspace import initialize_codegen
@@ -30,7 +30,7 @@ def init_command(path: str | None = None, token: str | None = None, language: st
         rich.print(format_command("gs init"))
         sys.exit(1)
 
-    session = CodegenSession(repo_path=repo_path, git_token=token)
+    session = CliSession(repo_path=repo_path, git_token=token)
     if language:
         session.config.repository.language = language.upper()
         session.config.save()
diff --git a/src/graph_sitter/cli/commands/list/main.py b/src/graph_sitter/cli/commands/list/main.py
index cec12c91e..37a3b5198 100644
--- a/src/graph_sitter/cli/commands/list/main.py
+++ b/src/graph_sitter/cli/commands/list/main.py
@@ -35,5 +35,5 @@ def list_command():
         rich.print(format_command("gs run <label>"))
     else:
         rich.print("[yellow]No codegen functions found in current directory.[/yellow]")
-        rich.print("\nAdd a function with @codegen.function decorator:")
-        rich.print(format_codeblock("@codegen.function('label')"))
+        rich.print("\nAdd a function with @graph_sitter.function decorator:")
+        rich.print(format_codeblock("@graph_sitter.function('label')"))
diff --git a/src/graph_sitter/cli/commands/notebook/main.py b/src/graph_sitter/cli/commands/notebook/main.py
index 3a3934034..33ff1cf74 100644
--- a/src/graph_sitter/cli/commands/notebook/main.py
+++ b/src/graph_sitter/cli/commands/notebook/main.py
@@ -4,7 +4,7 @@
 
 import rich_click as click
 
-from graph_sitter.cli.auth.session import CodegenSession
+from graph_sitter.cli.auth.session import CliSession
 from graph_sitter.cli.rich.spinners import create_spinner
 from graph_sitter.cli.utils.notebooks import create_notebook
 from graph_sitter.cli.workspace.decorators import requires_init
@@ -22,7 +22,7 @@ def create_jupyter_dir(codegen_dir: Path) -> Path:
 @click.option("--background", is_flag=True, help="Run Jupyter Lab in the background")
 @click.option("--demo", is_flag=True, help="Create a demo notebook with FastAPI example code")
 @requires_init
-def notebook_command(session: CodegenSession, background: bool, demo: bool):
+def notebook_command(session: CliSession, background: bool, demo: bool):
     """Launch Jupyter Lab with a pre-configured notebook for exploring your codebase."""
     with create_spinner("Setting up Jupyter environment...") as status:
         venv = VenvManager(codegen_dir=session.codegen_dir)
diff --git a/src/graph_sitter/cli/commands/run/main.py b/src/graph_sitter/cli/commands/run/main.py
index 9ded2d614..018479a3d 100644
--- a/src/graph_sitter/cli/commands/run/main.py
+++ b/src/graph_sitter/cli/commands/run/main.py
@@ -3,7 +3,7 @@
 
 import rich_click as click
 
-from graph_sitter.cli.auth.session import CodegenSession
+from graph_sitter.cli.auth.session import CliSession
 from graph_sitter.cli.utils.codemod_manager import CodemodManager
 from graph_sitter.cli.utils.json_schema import validate_json
 from graph_sitter.cli.workspace.decorators import requires_init
@@ -17,7 +17,7 @@
 @click.option("--diff-preview", type=int, help="Show a preview of the first N lines of the diff")
 @click.option("--arguments", type=str, help="Arguments as a json string to pass as the function's 'arguments' parameter")
 def run_command(
-    session: CodegenSession,
+    session: CliSession,
     label: str,
     daemon: bool = False,
     diff_preview: int | None = None,
diff --git a/src/graph_sitter/cli/commands/run/run_daemon.py b/src/graph_sitter/cli/commands/run/run_daemon.py
index ca95d2edb..a7bb7c446 100644
--- a/src/graph_sitter/cli/commands/run/run_daemon.py
+++ b/src/graph_sitter/cli/commands/run/run_daemon.py
@@ -2,7 +2,7 @@
 import rich_click as click
 from rich.panel import Panel
 
-from graph_sitter.cli.auth.session import CodegenSession
+from graph_sitter.cli.auth.session import CliSession
 from graph_sitter.cli.commands.start.docker_container import DockerContainer
 from graph_sitter.cli.errors import ServerError
 from graph_sitter.cli.rich.codeblocks import format_command
@@ -11,7 +11,7 @@
 from graph_sitter.runner.enums.warmup_state import WarmupState
 
 
-def run_daemon(session: CodegenSession, function, diff_preview: int | None = None):
+def run_daemon(session: CliSession, function, diff_preview: int | None = None):
     """Run a function on the cloud service.
 
     Args:
@@ -65,7 +65,7 @@ def run_daemon(session: CodegenSession, function, diff_preview: int | None = Non
             raise click.ClickException(str(e))
 
 
-def _get_docker_client(session: CodegenSession) -> DockerClient:
+def _get_docker_client(session: CliSession) -> DockerClient:
     repo_name = session.config.repository.name
     if (container := DockerContainer.get(repo_name)) is None:
         msg = f"graph_sitter.runner does not exist for {repo_name}. Please run 'codegen start' from {session.config.repository.path}."
diff --git a/src/graph_sitter/cli/commands/run/run_local.py b/src/graph_sitter/cli/commands/run/run_local.py
index f3d7fd73c..15c1a2264 100644
--- a/src/graph_sitter/cli/commands/run/run_local.py
+++ b/src/graph_sitter/cli/commands/run/run_local.py
@@ -5,7 +5,7 @@
 from rich.panel import Panel
 from rich.status import Status
 
-from graph_sitter.cli.auth.session import CodegenSession
+from graph_sitter.cli.auth.session import CliSession
 from graph_sitter.cli.utils.function_finder import DecoratedFunction
 from graph_sitter.codebase.config import ProjectConfig
 from graph_sitter.codebase.progress.progress import Progress
@@ -80,7 +80,7 @@ def parse_codebase(
 
 
 def run_local(
-    session: CodegenSession,
+    session: CliSession,
     function: DecoratedFunction,
     diff_preview: int | None = None,
 ) -> None:
@@ -91,7 +91,7 @@ def run_local(
         function: The function to run
         diff_preview: Number of lines of diff to preview (None for all)
     """
-    rich.print("Parsing codebase at {session.repo_path} with subdirectories {function.subdirectories or 'ALL'} and language {function.language or 'AUTO'} ...")
+    rich.print(f"Parsing codebase at {session.repo_path} with subdirectories {function.subdirectories or 'ALL'} and language {function.language or 'AUTO'} ...")
     # Parse codebase and run
     codebase = parse_codebase(repo_path=session.repo_path, subdirectories=function.subdirectories, language=function.language)
     with Status("[bold]Running codemod...", spinner="dots") as status:
diff --git a/src/graph_sitter/cli/mcp/resources/system_prompt.py b/src/graph_sitter/cli/mcp/resources/system_prompt.py
index 46e3af32e..210e71580 100644
--- a/src/graph_sitter/cli/mcp/resources/system_prompt.py
+++ b/src/graph_sitter/cli/mcp/resources/system_prompt.py
@@ -1358,7 +1358,7 @@ def baz():
 To fetch and parse a repository directly from GitHub, use the `from_repo` function.
 
 ```python
-import codegen
+import graph_sitter
 from graph_sitter.shared.enums.programming_language import ProgrammingLanguage
 
 # Fetch and parse a repository (defaults to /tmp/codegen/{repo_name})
@@ -1454,7 +1454,7 @@ def baz():
 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
 
@@ -1467,10 +1467,10 @@ def baz():
 This creates a new codemod in your `.codegen/codemods` directory:
 
 ```python
-import codegen
+import graph_sitter
 from graph_sitter.core.codebase 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
@@ -1513,15 +1513,15 @@ def run(codebase: Codebase):
 
 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.core.codebase 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:
@@ -1542,7 +1542,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)
@@ -1701,7 +1701,7 @@ def run(codebase: Codebase, arguments: RenameArgs):
 
 ---
 title: Function Decorator
-sidebarTitle: "@codegen.function"
+sidebarTitle: "@graph_sitter.function"
 icon: "at"
 iconType: "solid"
 ---
@@ -1712,17 +1712,17 @@ def run(codebase: Codebase, arguments: RenameArgs):
 
 ## 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
 
@@ -6768,7 +6768,7 @@ def create_blast_radius_visualization(symbol: PySymbol, depth: int = 0):
 First, let's import the types we need from Codegen:
 
 ```python
-import codegen
+import graph_sitter
 from graph_sitter.core.codebase import Codebase
 from graph_sitter.core.external_module import ExternalModule
 from graph_sitter.core.import_resolution import Import
diff --git a/src/graph_sitter/cli/sdk/decorator.py b/src/graph_sitter/cli/sdk/decorator.py
index 047d7e548..36f22aa36 100644
--- a/src/graph_sitter/cli/sdk/decorator.py
+++ b/src/graph_sitter/cli/sdk/decorator.py
@@ -55,7 +55,7 @@ def function(name: str, subdirectories: list[str] | None = None, language: Progr
         name: The name of the function to be used when deployed
 
     Example:
-        @codegen.function('my-function')
+        @graph_sitter.function('my-function')
         def run(codebase):
             pass
 
diff --git a/src/graph_sitter/cli/utils/codemod_manager.py b/src/graph_sitter/cli/utils/codemod_manager.py
index 01421b03d..74eabdb26 100644
--- a/src/graph_sitter/cli/utils/codemod_manager.py
+++ b/src/graph_sitter/cli/utils/codemod_manager.py
@@ -15,7 +15,7 @@ def _might_have_decorators(file_path: Path) -> bool:
     try:
         # Read in binary mode and check for b'@codegen' to handle any encoding
         with open(file_path, "rb") as f:
-            return b"@codegen" in f.read()
+            return b"@graph_sitter" in f.read()
     except Exception:
         return False
 
diff --git a/src/graph_sitter/cli/utils/default_code.py b/src/graph_sitter/cli/utils/default_code.py
index eef6c78d1..518bb2663 100644
--- a/src/graph_sitter/cli/utils/default_code.py
+++ b/src/graph_sitter/cli/utils/default_code.py
@@ -1,8 +1,8 @@
-DEFAULT_CODEMOD = '''import codegen
+DEFAULT_CODEMOD = '''import graph_sitter
 from graph_sitter.core.codebase import Codebase
 
 
-@codegen.function("{name}")
+@graph_sitter.function("{name}")
 def run(codebase: Codebase):
     """Add a description of what this codemod does."""
     # Add your code here
diff --git a/src/graph_sitter/cli/utils/function_finder.py b/src/graph_sitter/cli/utils/function_finder.py
index 5e255bff8..0116eaff3 100644
--- a/src/graph_sitter/cli/utils/function_finder.py
+++ b/src/graph_sitter/cli/utils/function_finder.py
@@ -207,7 +207,7 @@ def visit_FunctionDef(self, node):
                 and (len(decorator.args) > 0 or len(decorator.keywords) > 0)
                 and (
                     # Check if it's a direct codegen.X call
-                    (isinstance(decorator.func, ast.Attribute) and isinstance(decorator.func.value, ast.Name) and decorator.func.value.id == "codegen")
+                    (isinstance(decorator.func, ast.Attribute) and isinstance(decorator.func.value, ast.Name) and decorator.func.value.id in ("graph_sitter", "codegen"))
                     or
                     # Check if it starts with codegen.anything.anything...
                     (isinstance(decorator.func, ast.Attribute) and isinstance(decorator.func.value, ast.Attribute) and self._has_codegen_root(decorator.func.value))
diff --git a/src/graph_sitter/cli/workspace/decorators.py b/src/graph_sitter/cli/workspace/decorators.py
index b3bf1a192..04c3b00cf 100644
--- a/src/graph_sitter/cli/workspace/decorators.py
+++ b/src/graph_sitter/cli/workspace/decorators.py
@@ -2,7 +2,7 @@
 import sys
 from collections.abc import Callable
 
-from graph_sitter.cli.auth.session import CodegenSession
+from graph_sitter.cli.auth.session import CliSession
 from graph_sitter.cli.rich.pretty_print import pretty_print_error
 
 
@@ -12,7 +12,7 @@ def requires_init(f: Callable) -> Callable:
     @functools.wraps(f)
     def wrapper(*args, **kwargs):
         # Create a session if one wasn't provided
-        session = kwargs.get("session") or CodegenSession.from_active_session()
+        session = kwargs.get("session") or CliSession.from_active_session()
         if session is None:
             pretty_print_error("Graph-sitter not initialized. Please run `gs init` from a git repo workspace.")
             sys.exit(1)
diff --git a/src/graph_sitter/cli/workspace/initialize_workspace.py b/src/graph_sitter/cli/workspace/initialize_workspace.py
index e19047cc3..556f2cba9 100644
--- a/src/graph_sitter/cli/workspace/initialize_workspace.py
+++ b/src/graph_sitter/cli/workspace/initialize_workspace.py
@@ -4,18 +4,18 @@
 from rich.status import Status
 
 from graph_sitter.cli.auth.constants import CODEGEN_DIR, DOCS_DIR, EXAMPLES_DIR, PROMPTS_DIR
-from graph_sitter.cli.auth.session import CodegenSession
+from graph_sitter.cli.auth.session import CliSession
 from graph_sitter.cli.rich.spinners import create_spinner
 from graph_sitter.cli.utils.notebooks import create_notebook
 from graph_sitter.cli.workspace.venv_manager import VenvManager
 
 
-def initialize_codegen(session: CodegenSession, status: Status | str = "Initializing") -> CodegenSession:
+def initialize_codegen(session: CliSession, status: Status | str = "Initializing") -> CliSession:
     """Initialize or update the codegen directory structure and content.
 
     Args:
         status: Either a Status object to update, or a string action being performed ("Initializing" or "Updating")
-        session: Optional CodegenSession for fetching docs and examples
+        session: Optional CliSession for fetching docs and examples
         fetch_docs: Whether to fetch docs and examples (requires auth)
 
     Returns:
diff --git a/src/graph_sitter/codebase/codebase_context.py b/src/graph_sitter/codebase/codebase_context.py
index f18b6438f..a52e7ac94 100644
--- a/src/graph_sitter/codebase/codebase_context.py
+++ b/src/graph_sitter/codebase/codebase_context.py
@@ -619,11 +619,13 @@ def get_nodes(self, node_type: NodeType | None = None, exclude_type: NodeType |
     def get_edges(self) -> list[tuple[NodeId, NodeId, EdgeType, Usage | None]]:
         return [(x[0], x[1], x[2].type, x[2].usage) for x in self._graph.weighted_edge_list()]
 
-    def get_file(self, file_path: os.PathLike, ignore_case: bool = False) -> SourceFile | None:
+    def get_file(self, file_path: os.PathLike, ignore_case: bool = False, relative_only: bool = False) -> SourceFile | None:
         # Performance hack: just use the relative path
         node_id = self.filepath_idx.get(str(file_path), None)
         if node_id is not None:
             return self.get_node(node_id)
+        if relative_only:
+            return None
         # If not part of repo path, return None
         absolute_path = self.to_absolute(file_path)
         if not self.is_subdir(absolute_path) and not self.config.allow_external:
diff --git a/src/graph_sitter/codebase/config.py b/src/graph_sitter/codebase/config.py
index f5870011e..b6250585b 100644
--- a/src/graph_sitter/codebase/config.py
+++ b/src/graph_sitter/codebase/config.py
@@ -24,7 +24,7 @@ class SessionOptions(BaseModel):
     max_ai_requests: int = Field(default=150, le=HARD_MAX_AI_LIMIT)
 
 
-TestFlags = DefaultCodebaseConfig.model_copy(update=dict(debug=True, track_graph=True, verify_graph=True, full_range_index=True, sync_enabled=True))
+TestFlags = DefaultCodebaseConfig.model_copy(update=dict(debug=True, track_graph=True, verify_graph=True, full_range_index=True, sync_enabled=True, conditional_type_resolution=True))
 LintFlags = DefaultCodebaseConfig.model_copy(update=dict(method_usages=False, sync_enabled=True))
 ParseTestFlags = DefaultCodebaseConfig.model_copy(update=dict(debug=False, track_graph=False, sync_enabled=True))
 
@@ -43,7 +43,7 @@ class ProjectConfig(BaseModel):
     @classmethod
     def from_path(cls, path: str, programming_language: ProgrammingLanguage | None = None) -> Self:
         # Split repo_path into (git_root, base_path)
-        repo_path = os.path.abspath(path)
+        repo_path = os.path.realpath(os.path.abspath(path))
         git_root, base_path = split_git_path(repo_path)
         subdirectories = [base_path] if base_path else None
         programming_language = programming_language or determine_project_language(repo_path)
diff --git a/src/graph_sitter/configs/models/codebase.py b/src/graph_sitter/configs/models/codebase.py
index 06dd81843..84c109739 100644
--- a/src/graph_sitter/configs/models/codebase.py
+++ b/src/graph_sitter/configs/models/codebase.py
@@ -37,6 +37,7 @@ def __init__(self, prefix: str = "CODEBASE", *args, **kwargs) -> None:
     ts_language_engine: bool = False
     v8_ts_engine: bool = False
     unpacking_assignment_partial_removal: bool = True
+    conditional_type_resolution: bool = False
     use_pink: PinkMode = PinkMode.OFF
 
 
diff --git a/src/graph_sitter/core/expressions/name.py b/src/graph_sitter/core/expressions/name.py
index 1c45e8f8d..959d88bac 100644
--- a/src/graph_sitter/core/expressions/name.py
+++ b/src/graph_sitter/core/expressions/name.py
@@ -114,7 +114,7 @@ def resolve_name(self, name: str, start_byte: int | None = None, strict: bool =
         else:
             return
 
-        if hasattr(resolved_name, "parent") and (conditional_parent := resolved_name.parent_of_type(ConditionalBlock)):
+        if self.ctx.config.conditional_type_resolution and hasattr(resolved_name, "parent") and (conditional_parent := resolved_name.parent_of_type(ConditionalBlock)):
             if self.parent_of_type(ConditionalBlock) == conditional_parent:
                 # Use in the same block, should only depend on the inside of the block
                 return
diff --git a/src/graph_sitter/core/interfaces/editable.py b/src/graph_sitter/core/interfaces/editable.py
index 47e929357..cbff625e9 100644
--- a/src/graph_sitter/core/interfaces/editable.py
+++ b/src/graph_sitter/core/interfaces/editable.py
@@ -116,7 +116,6 @@ class Editable(JSONable, Generic[Parent]):
     ctx: CodebaseContext
     parent: Parent
     node_type: NodeType
-    _file: File | None = None
     _hash: int | None = None
 
     def __init__(self, ts_node: TSNode, file_node_id: NodeId, ctx: CodebaseContext, parent: Parent) -> None:
@@ -379,9 +378,7 @@ def file(self) -> SourceFile:
         Returns:
             File: The File object containing this Editable instance.
         """
-        if self._file is None:
-            self._file = self.ctx.get_node(self.file_node_id)
-        return self._file  # type: ignore
+        return self.ctx.get_node(self.file_node_id)
 
     @property
     def filepath(self) -> str:
diff --git a/src/graph_sitter/extensions/attribution/cli.py b/src/graph_sitter/extensions/attribution/cli.py
index eb204f919..d04f85a5e 100644
--- a/src/graph_sitter/extensions/attribution/cli.py
+++ b/src/graph_sitter/extensions/attribution/cli.py
@@ -3,7 +3,7 @@
 
 import pygit2
 
-import codegen
+import graph_sitter
 from graph_sitter.core.codebase import Codebase
 from graph_sitter.extensions.attribution.main import add_attribution_to_symbols, analyze_ai_impact
 
@@ -58,7 +58,7 @@ def diagnose_repository(codebase):
         print(f"⚠️ Error during repository diagnosis: {e}")
 
 
-@codegen.function("analyze-ai-impact")
+@graph_sitter.function("analyze-ai-impact")
 def run(codebase: Codebase):
     """Analyze the impact of AI on the codebase.
 
diff --git a/src/graph_sitter/python/import_resolution.py b/src/graph_sitter/python/import_resolution.py
index 127a726be..87b2ed274 100644
--- a/src/graph_sitter/python/import_resolution.py
+++ b/src/graph_sitter/python/import_resolution.py
@@ -116,11 +116,11 @@ def resolve_import(self, base_path: str | None = None, *, add_module_name: str |
                     return ImportResolution(from_file=file, symbol=None, imports_file=True)
 
             # =====[ Default path ]=====
-            if file := self.ctx.get_file(filepath):
+            if file := self.ctx.get_file(filepath, relative_only=True):
                 return ImportResolution(from_file=file, symbol=None, imports_file=True)
 
             filepath = filepath.replace(".py", "/__init__.py")
-            if file := self.ctx.get_file(filepath):
+            if file := self.ctx.get_file(filepath, relative_only=True):
                 # TODO - I think this is another edge case, due to `dao/__init__.py` etc.
                 # You can't do `from a.b.c import foo` => `foo.utils.x` right now since `foo` is just a file...
                 return ImportResolution(from_file=file, symbol=None, imports_file=True)
@@ -136,7 +136,7 @@ def resolve_import(self, base_path: str | None = None, *, add_module_name: str |
 
             # =====[ Check if `module.py` file exists in the graph ]=====
             filepath = os.path.join(base_path, filepath)
-            if file := self.ctx.get_file(filepath):
+            if file := self.ctx.get_file(filepath, relative_only=True):
                 symbol = file.get_node_by_name(symbol_name)
                 if symbol is None:
                     if file.get_node_from_wildcard_chain(symbol_name):
@@ -165,7 +165,7 @@ def resolve_import(self, base_path: str | None = None, *, add_module_name: str |
                         return ImportResolution(from_file=from_file, symbol=symbol)
 
             # =====[ Check if `module/__init__.py` file exists in the graph ]=====
-            if from_file := self.ctx.get_file(filepath):
+            if from_file := self.ctx.get_file(filepath, relative_only=True):
                 symbol = from_file.get_node_by_name(symbol_name)
                 if symbol is None:
                     if from_file.get_node_from_wildcard_chain(symbol_name):
@@ -179,11 +179,13 @@ def resolve_import(self, base_path: str | None = None, *, add_module_name: str |
 
             # =====[ Case: Can't resolve the import ]=====
             if base_path == "":
-                # Try to resolve with "src" as the base path
-                return self.resolve_import(base_path="src", add_module_name=add_module_name)
+                if self.ctx.get_directory("src"):
+                    # Try to resolve with "src" as the base path
+                    return self.resolve_import(base_path="src", add_module_name=add_module_name)
             if base_path == "src":
-                # Try "test" next
-                return self.resolve_import(base_path="test", add_module_name=add_module_name)
+                if self.ctx.get_directory("test"):
+                    # Try "test" next
+                    return self.resolve_import(base_path="test", add_module_name=add_module_name)
 
             # if not G_override:
             #     for resolver in ctx.import_resolvers:
diff --git a/src/graph_sitter/system-prompt.txt b/src/graph_sitter/system-prompt.txt
index f7d202c9e..a6ec0570c 100644
--- a/src/graph_sitter/system-prompt.txt
+++ b/src/graph_sitter/system-prompt.txt
@@ -650,7 +650,7 @@ gs creates a custom Python environment in `.codegen/.venv`. Configure your IDE t
        ```bash
        .codegen/.venv/bin/python
        ```
-    
+
     Alternatively, create a `.vscode/settings.json`:
         ```json
         {
@@ -672,7 +672,7 @@ gs creates a custom Python environment in `.codegen/.venv`. Configure your IDE t
        .codegen/.venv/bin/python
        ```
   </Accordion>
-  
+
 </AccordionGroup>
 
 
@@ -1618,8 +1618,8 @@ iconType: "solid"
     </Tip>
   </Accordion>
   <Accordion title="Can I use Graph-sitter with my existing tools?" icon="screwdriver-wrench">
-    Yes - [by design](/introduction/guiding-principles#python-first-composability). 
-    
+    Yes - [by design](/introduction/guiding-principles#python-first-composability).
+
     Graph-sitter works like any other python package. It works alongside your IDE, version control system, and other development tools.
   </Accordion>
   <Accordion
@@ -1931,7 +1931,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
 
@@ -1944,10 +1944,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
@@ -1990,15 +1990,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:
@@ -2019,7 +2019,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)
@@ -2171,7 +2171,7 @@ gs run my-codemod --apply-local
 
 ---
 title: Function Decorator
-sidebarTitle: "@codegen.function"
+sidebarTitle: "@graph_sitter.function"
 icon: "at"
 iconType: "solid"
 ---
@@ -2182,17 +2182,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
 
@@ -2568,7 +2568,7 @@ if not isinstance(mdx_file, SourceFile):
 ```
 
 <Note>
-  Currently, the codebase object can only parse source code files of one language at a time. This means that if you want to work with both Python and TypeScript files, you will need to create two separate codebase objects. 
+  Currently, the codebase object can only parse source code files of one language at a time. This means that if you want to work with both Python and TypeScript files, you will need to create two separate codebase objects.
 </Note>
 
 ## Accessing Code
@@ -3388,7 +3388,7 @@ for module, imports in module_imports.items():
 
 <Note>
 Always check if imports resolve to external modules before modification to avoid breaking third-party package imports.
-</Note> 
+</Note>
 
 ## Import Statements vs Imports
 
@@ -3590,7 +3590,7 @@ for exp in file.exports:
         # Get original and current symbols
         current = exp.exported_symbol
         original = exp.resolved_symbol
-        
+
         print(f"Re-exporting {original.name} from {exp.from_file.filepath}")
         print(f"Through: {' -> '.join(e.file.filepath for e in exp.export_chain)}")
 ```
@@ -3640,7 +3640,7 @@ for from_file, exports in file_exports.items():
 
 <Note>
 When managing exports, consider the impact on your module's public API. Not all symbols that can be exported should be exported.
-</Note> 
+</Note>
 
 ---
 title: "Inheritable Behaviors"
@@ -4130,9 +4130,9 @@ If `A` depends on `B`, then `B` is used by `A`. This relationship is tracked in
 
 flowchart LR
     B(BaseClass)
-    
-    
-    
+
+
+
     A(MyClass)
     B ---| used by |A
     A ---|depends on |B
@@ -4301,7 +4301,7 @@ class A:
     def method_a(self): pass
 
 class B(A):
-    def method_b(self): 
+    def method_b(self):
         self.method_a()
 
 class C(B):
@@ -5191,7 +5191,7 @@ for attr in class_def.attributes:
     # Each attribute has an assignment property
     attr_type = attr.assignment.type  # -> TypeAnnotation
     print(f"{attr.name}: {attr_type.source}")  # e.g. "x: int"
-    
+
     # Set attribute type
     attr.assignment.set_type("int")
 
@@ -5208,7 +5208,7 @@ Union types ([UnionType](/api-reference/core/UnionType)) can be manipulated as c
 
 ```python
 # Get union type
-union_type = function.return_type  # -> A | B 
+union_type = function.return_type  # -> A | B
 print(union_type.symbols)  # ["A", "B"]
 
 # Add/remove options
@@ -6059,13 +6059,13 @@ Here's an example of using flags during code analysis:
 
 ```python
 def analyze_codebase(codebase):
-    for function in codebase.functions:            
+    for function in codebase.functions:
         # Check documentation
         if not function.docstring:
             function.flag(
                 message="Missing docstring",
             )
-            
+
         # Check error handling
         if function.is_async and not function.has_try_catch:
             function.flag(
@@ -6768,7 +6768,7 @@ Explore our tutorials to learn how to use Graph-sitter for various code transfor
   >
     Update API calls, handle breaking changes, and manage bulk updates across your codebase.
   </Card>
-  <Card 
+  <Card
     title="SQLAlchemy 1.4 to 2.0"
     icon="layer-group"
     href="/tutorials/sqlalchemy-1.6-to-2.0"
@@ -6782,9 +6782,9 @@ Explore our tutorials to learn how to use Graph-sitter for various code transfor
   >
     Convert Flask applications to FastAPI, updating routes and dependencies.
   </Card>
-  <Card 
-    title="Python 2 to 3" 
-    icon="snake" 
+  <Card
+    title="Python 2 to 3"
+    icon="snake"
     href="/tutorials/python2-to-python3"
   >
     Migrate Python 2 code to Python 3, updating syntax and modernizing APIs.
@@ -6801,9 +6801,9 @@ Explore our tutorials to learn how to use Graph-sitter for various code transfor
   >
     Restructure files, enforce naming conventions, and improve project layout.
   </Card>
-  <Card 
-    title="Improve Modularity" 
-    icon="cubes" 
+  <Card
+    title="Improve Modularity"
+    icon="cubes"
     href="/tutorials/modularity"
   >
     Split large files, extract shared logic, and manage dependencies.
@@ -6994,21 +6994,21 @@ def research(repo_name: Optional[str] = None, query: Optional[str] = None):
     """Start a code research session."""
     # Initialize codebase
     codebase = initialize_codebase(repo_name)
-    
+
     # Create and run the agent
     agent = create_research_agent(codebase)
-    
+
     # Main research loop
     while True:
         if not query:
             query = Prompt.ask("[bold cyan]Research query[/bold cyan]")
-            
+
         result = agent.invoke(
             {"input": query},
             config={"configurable": {"thread_id": 1}}
         )
         console.print(Markdown(result["messages"][-1].content))
-        
+
         query = None  # Clear for next iteration
 ```
 
@@ -7056,7 +7056,7 @@ class CustomAnalysisTool(BaseTool):
     """Custom tool for specialized code analysis."""
     name = "custom_analysis"
     description = "Performs specialized code analysis"
-    
+
     def _run(self, query: str) -> str:
         # Custom analysis logic
         return results
@@ -7196,7 +7196,7 @@ def calculate_maintainability_index(
 
 ## Line Metrics
 
-Line metrics provide insights into the size, complexity, and maintainability of a codebase. These measurements help determine the scale of a project, identify areas that may need refactoring, and track the growth of the codebase over time. 
+Line metrics provide insights into the size, complexity, and maintainability of a codebase. These measurements help determine the scale of a project, identify areas that may need refactoring, and track the growth of the codebase over time.
 
 ### Lines of Code
 Lines of Code refers to the total number of lines in the source code, including blank lines and comments. This is accomplished with a simple count of all lines in the source file.
@@ -7277,7 +7277,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
@@ -7493,7 +7493,7 @@ from graph_sitter import Codebase
 # Initialize codebase
 codebase = Codebase("path/to/posthog/")
 
-# Create a directed graph for representing call relationships 
+# Create a directed graph for representing call relationships
 G = nx.DiGraph()
 
 # Configuration flags
@@ -7515,7 +7515,7 @@ We'll create a function that will recursively traverse the call trace of a funct
 ```python
 def create_downstream_call_trace(src_func: Function, depth: int = 0):
     """Creates call graph by recursively traversing function calls
-    
+
     Args:
         src_func (Function): Starting function for call graph
         depth (int): Current recursion depth
@@ -7523,7 +7523,7 @@ def create_downstream_call_trace(src_func: Function, depth: int = 0):
     # Prevent infinite recursion
     if MAX_DEPTH <= depth:
         return
-        
+
     # External modules are not functions
     if isinstance(src_func, ExternalModule):
         return
@@ -7533,12 +7533,12 @@ def create_downstream_call_trace(src_func: Function, depth: int = 0):
         # Skip self-recursive calls
         if call.name == src_func.name:
             continue
-            
+
         # Get called function definition
         func = call.function_definition
         if not func:
             continue
-            
+
         # Apply configured filters
         if isinstance(func, ExternalModule) and IGNORE_EXTERNAL_MODULE_CALLS:
             continue
@@ -7552,7 +7552,7 @@ def create_downstream_call_trace(src_func: Function, depth: int = 0):
             func_name = f"{func.parent_class.name}.{func.name}" if func.is_method else func.name
 
         # Add node and edge with metadata
-        G.add_node(func, name=func_name, 
+        G.add_node(func, name=func_name,
                   color=COLOR_PALETTE.get(func.__class__.__name__))
         G.add_edge(src_func, func, **generate_edge_meta(call))
 
@@ -7567,10 +7567,10 @@ We can enrich our edges with metadata about the function calls:
 ```python
 def generate_edge_meta(call: FunctionCall) -> dict:
     """Generate metadata for call graph edges
-    
+
     Args:
         call (FunctionCall): Function call information
-        
+
     Returns:
         dict: Edge metadata including name and location
     """
@@ -7589,8 +7589,8 @@ Finally, we can visualize our call graph starting from a specific function:
 target_class = codebase.get_class('SharingConfigurationViewSet')
 target_method = target_class.get_method('patch')
 
-# Add root node 
-G.add_node(target_method, 
+# Add root node
+G.add_node(target_method,
            name=f"{target_class.name}.{target_method.name}",
            color=COLOR_PALETTE["StartFunction"])
 
@@ -7640,7 +7640,7 @@ The core function for building our dependency graph:
 ```python
 def create_dependencies_visualization(symbol: Symbol, depth: int = 0):
     """Creates visualization of symbol dependencies
-    
+
     Args:
         symbol (Symbol): Starting symbol to analyze
         depth (int): Current recursion depth
@@ -7648,11 +7648,11 @@ def create_dependencies_visualization(symbol: Symbol, depth: int = 0):
     # Prevent excessive recursion
     if depth >= MAX_DEPTH:
         return
-    
+
     # Process each dependency
     for dep in symbol.dependencies:
         dep_symbol = None
-        
+
         # Handle different dependency types
         if isinstance(dep, Symbol):
             # Direct symbol reference
@@ -7663,13 +7663,13 @@ def create_dependencies_visualization(symbol: Symbol, depth: int = 0):
 
         if dep_symbol:
             # Add node with appropriate styling
-            G.add_node(dep_symbol, 
-                      color=COLOR_PALETTE.get(dep_symbol.__class__.__name__, 
+            G.add_node(dep_symbol,
+                      color=COLOR_PALETTE.get(dep_symbol.__class__.__name__,
                                             "#f694ff"))
-            
+
             # Add dependency relationship
             G.add_edge(symbol, dep_symbol)
-            
+
             # Recurse unless it's a class (avoid complexity)
             if not isinstance(dep_symbol, PyClass):
                 create_dependencies_visualization(dep_symbol, depth + 1)
@@ -7681,7 +7681,7 @@ Finally, we can visualize our dependency graph starting from a specific symbol:
 # Get target symbol
 target_func = codebase.get_function("get_query_runner")
 
-# Add root node 
+# Add root node
 G.add_node(target_func, color=COLOR_PALETTE["StartFunction"])
 
 # Generate dependency graph
@@ -7724,16 +7724,16 @@ HTTP_METHODS = ["get", "put", "patch", "post", "head", "delete"]
 
 def generate_edge_meta(usage: Usage) -> dict:
     """Generate metadata for graph edges
-    
+
     Args:
         usage (Usage): Usage relationship information
-        
+
     Returns:
         dict: Edge metadata including name and location
     """
     return {
         "name": usage.match.source,
-        "file_path": usage.match.filepath, 
+        "file_path": usage.match.filepath,
         "start_point": usage.match.start_point,
         "end_point": usage.match.end_point,
         "symbol_name": usage.match.__class__.__name__
@@ -7741,10 +7741,10 @@ def generate_edge_meta(usage: Usage) -> dict:
 
 def is_http_method(symbol: PySymbol) -> bool:
     """Check if a symbol is an HTTP endpoint method
-    
+
     Args:
         symbol (PySymbol): Symbol to check
-        
+
     Returns:
         bool: True if symbol is an HTTP method
     """
@@ -7758,7 +7758,7 @@ The main function for creating our blast radius visualization:
 ```python
 def create_blast_radius_visualization(symbol: PySymbol, depth: int = 0):
     """Create visualization of symbol usage relationships
-    
+
     Args:
         symbol (PySymbol): Starting symbol to analyze
         depth (int): Current recursion depth
@@ -7766,11 +7766,11 @@ def create_blast_radius_visualization(symbol: PySymbol, depth: int = 0):
     # Prevent excessive recursion
     if depth >= MAX_DEPTH:
         return
-    
+
     # Process each usage of the symbol
     for usage in symbol.usages:
         usage_symbol = usage.usage_symbol
-        
+
         # Determine node color based on type
         if is_http_method(usage_symbol):
             color = COLOR_PALETTE.get("HTTP_METHOD")
@@ -7780,7 +7780,7 @@ def create_blast_radius_visualization(symbol: PySymbol, depth: int = 0):
         # Add node and edge to graph
         G.add_node(usage_symbol, color=color)
         G.add_edge(symbol, usage_symbol, **generate_edge_meta(usage))
-        
+
         # Recursively process usage symbol
         create_blast_radius_visualization(usage_symbol, depth + 1)
 ```
@@ -7931,7 +7931,7 @@ for call in old_api.call_sites:
         f"data={call.get_arg_by_parameter_name('input').value}",
         f"timeout={call.get_arg_by_parameter_name('wait').value}"
     ]
-    
+
     # Replace the old call with the new API
     call.replace(f"new_process_data({', '.join(args)})")
 ```
@@ -7945,10 +7945,10 @@ When updating chained method calls, like database queries or builder patterns:
 for execute_call in codebase.function_calls:
     if execute_call.name != "execute":
         continue
-        
+
     # Get the full chain
     chain = execute_call.call_chain
-    
+
     # Example: Add .timeout() before .execute()
     if "timeout" not in {call.name for call in chain}:
         execute_call.insert_before("timeout(30)")
@@ -7967,45 +7967,45 @@ Here's a comprehensive example:
 ```python
 def migrate_api_v1_to_v2(codebase):
     old_api = codebase.get_function("create_user_v1")
-    
+
     # Document all existing call patterns
     call_patterns = {}
     for call in old_api.call_sites:
         args = [arg.source for arg in call.args]
         pattern = ", ".join(args)
         call_patterns[pattern] = call_patterns.get(pattern, 0) + 1
-    
+
     print("Found call patterns:")
     for pattern, count in call_patterns.items():
         print(f"  {pattern}: {count} occurrences")
-    
+
     # Create new API version
     new_api = old_api.copy()
     new_api.rename("create_user_v2")
-    
+
     # Update parameter types
     new_api.get_parameter("email").type = "EmailStr"
     new_api.get_parameter("role").type = "UserRole"
-    
+
     # Add new required parameters
     new_api.add_parameter("tenant_id: UUID")
-    
+
     # Update all call sites
     for call in old_api.call_sites:
         # Get current arguments
         email_arg = call.get_arg_by_parameter_name("email")
         role_arg = call.get_arg_by_parameter_name("role")
-        
+
         # Build new argument list with type conversions
         new_args = [
             f"email=EmailStr({email_arg.value})",
             f"role=UserRole({role_arg.value})",
             "tenant_id=get_current_tenant_id()"
         ]
-        
+
         # Replace old call with new version
         call.replace(f"create_user_v2({', '.join(new_args)})")
-    
+
     # Add deprecation notice to old version
     old_api.add_decorator('@deprecated("Use create_user_v2 instead")')
 
@@ -8027,10 +8027,10 @@ migrate_api_v1_to_v2(codebase)
    ```python
    # First update parameter names
    param.rename("new_name")
-   
+
    # Then update types
    param.type = "new_type"
-   
+
    # Finally update call sites
    for call in api.call_sites:
        # ... update calls
@@ -8040,7 +8040,7 @@ migrate_api_v1_to_v2(codebase)
    ```python
    # Add new parameter with default
    api.add_parameter("new_param: str = None")
-   
+
    # Later make it required
    api.get_parameter("new_param").remove_default()
    ```
@@ -8055,7 +8055,7 @@ migrate_api_v1_to_v2(codebase)
 
 <Note>
 Remember to test thoroughly after making bulk changes to APIs. While Graph-sitter ensures syntactic correctness, you'll want to verify the semantic correctness of the changes.
-</Note> 
+</Note>
 
 ---
 title: "Organizing Your Codebase"
@@ -8619,16 +8619,16 @@ from collections import defaultdict
 # Create a graph of file dependencies
 def create_dependency_graph():
     G = nx.DiGraph()
-    
+
     for file in codebase.files:
         # Add node for this file
         G.add_node(file.filepath)
-        
+
         # Add edges for each import
         for imp in file.imports:
             if imp.from_file:  # Skip external imports
                 G.add_edge(file.filepath, imp.from_file.filepath)
-    
+
     return G
 
 # Create and analyze the graph
@@ -8657,18 +8657,18 @@ def break_circular_dependency(cycle):
     # Get the first two files in the cycle
     file1 = codebase.get_file(cycle[0])
     file2 = codebase.get_file(cycle[1])
-    
+
     # Create a shared module for common code
     shared_dir = "shared"
     if not codebase.has_directory(shared_dir):
         codebase.create_directory(shared_dir)
-    
+
     # Find symbols used by both files
     shared_symbols = []
     for symbol in file1.symbols:
         if any(usage.file == file2 for usage in symbol.usages):
             shared_symbols.append(symbol)
-    
+
     # Move shared symbols to a new file
     if shared_symbols:
         shared_file = codebase.create_file(f"{shared_dir}/shared_types.py")
@@ -8690,7 +8690,7 @@ def organize_file_imports(file):
     std_lib_imports = []
     third_party_imports = []
     local_imports = []
-    
+
     for imp in file.imports:
         if imp.is_standard_library:
             std_lib_imports.append(imp)
@@ -8698,26 +8698,26 @@ def organize_file_imports(file):
             third_party_imports.append(imp)
         else:
             local_imports.append(imp)
-    
+
     # Sort each group
     for group in [std_lib_imports, third_party_imports, local_imports]:
         group.sort(key=lambda x: x.module_name)
-    
+
     # Remove all existing imports
     for imp in file.imports:
         imp.remove()
-    
+
     # Add imports back in organized groups
     if std_lib_imports:
         for imp in std_lib_imports:
             file.add_import(imp.source)
         file.insert_after_imports("")  # Add newline
-        
+
     if third_party_imports:
         for imp in third_party_imports:
             file.add_import(imp.source)
         file.insert_after_imports("")  # Add newline
-        
+
     if local_imports:
         for imp in local_imports:
             file.add_import(imp.source)
@@ -8736,22 +8736,22 @@ from collections import defaultdict
 
 def analyze_module_coupling():
     coupling_scores = defaultdict(int)
-    
+
     for file in codebase.files:
         # Count unique files imported from
         imported_files = {imp.from_file for imp in file.imports if imp.from_file}
         coupling_scores[file.filepath] = len(imported_files)
-        
+
         # Count files that import this file
-        importing_files = {usage.file for symbol in file.symbols 
+        importing_files = {usage.file for symbol in file.symbols
                          for usage in symbol.usages if usage.file != file}
         coupling_scores[file.filepath] += len(importing_files)
-    
+
     # Sort by coupling score
-    sorted_files = sorted(coupling_scores.items(), 
-                         key=lambda x: x[1], 
+    sorted_files = sorted(coupling_scores.items(),
+                         key=lambda x: x[1],
                          reverse=True)
-    
+
     print("\n🔍 Module Coupling Analysis:")
     print("\nMost coupled files:")
     for filepath, score in sorted_files[:5]:
@@ -8769,9 +8769,9 @@ def extract_shared_code(file, min_usages=3):
     # Find symbols used by multiple files
     for symbol in file.symbols:
         # Get unique files using this symbol
-        using_files = {usage.file for usage in symbol.usages 
+        using_files = {usage.file for usage in symbol.usages
                       if usage.file != file}
-        
+
         if len(using_files) >= min_usages:
             # Create appropriate shared module
             module_name = determine_shared_module(symbol)
@@ -8779,7 +8779,7 @@ def extract_shared_code(file, min_usages=3):
                 shared_file = codebase.create_file(f"shared/{module_name}.py")
             else:
                 shared_file = codebase.get_file(f"shared/{module_name}.py")
-            
+
             # Move symbol to shared module
             symbol.move_to_file(shared_file, strategy="update_all_imports")
 
@@ -8833,7 +8833,7 @@ if feature_flag_class:
     # Initialize usage count for all attributes
     for attr in feature_flag_class.attributes:
         feature_flag_usage[attr.name] = 0
-    
+
     # Get all usages of the FeatureFlag class
     for usage in feature_flag_class.usages:
         usage_source = usage.usage_symbol.source if hasattr(usage, 'usage_symbol') else str(usage)
@@ -9578,7 +9578,7 @@ Let's break down how this works:
         if export.is_reexport() and export.is_default_export():
             print(f"  🔄 Converting default export '{export.name}'")
     ```
-    
+
     The code identifies default exports by checking:
     1. If it's a re-export (`is_reexport()`)
     2. If it's a default export (`is_default_export()`)
@@ -9686,7 +9686,7 @@ for file in codebase.files:
 
     print(f"✨ Fixed exports in {target_file.filepath}")
 
-``` 
+```
 
 ---
 title: "Creating Documentation"
@@ -9775,11 +9775,11 @@ for directory in codebase.directories:
     # Skip test, sql and alembic directories
     if any(x in directory.path.lower() for x in ['test', 'sql', 'alembic']):
         continue
-        
+
     # Get undecorated functions
     funcs = [f for f in directory.functions if not f.is_decorated]
     total = len(funcs)
-    
+
     # Only analyze dirs with >10 functions
     if total > 10:
         documented = sum(1 for f in funcs if f.docstring)
@@ -9794,12 +9794,12 @@ for directory in codebase.directories:
 if dir_stats:
     lowest_dir = min(dir_stats.items(), key=lambda x: x[1]['coverage'])
     path, stats = lowest_dir
-    
+
     print(f"📉 Lowest coverage directory: '{path}'")
     print(f"  • Total functions: {stats['total']}")
     print(f"  • Documented: {stats['documented']}")
     print(f"  • Coverage: {stats['coverage']:.1f}%")
-    
+
     # Print all directory stats for comparison
     print("\n📊 All directory coverage rates:")
     for path, stats in sorted(dir_stats.items(), key=lambda x: x[1]['coverage']):
@@ -10587,7 +10587,7 @@ iconType: "solid"
 </Frame>
 
 
-Import loops occur when two or more Python modules depend on each other, creating a circular dependency. While some import cycles can be harmless, others can lead to runtime errors and make code harder to maintain. 
+Import loops occur when two or more Python modules depend on each other, creating a circular dependency. While some import cycles can be harmless, others can lead to runtime errors and make code harder to maintain.
 
 In this tutorial, we'll explore how to identify and fix problematic import cycles using Codegen.
 
@@ -11483,7 +11483,7 @@ Match (s: Func )-[r: CALLS]-> (e:Func) RETURN s, e LIMIT 10
 
 ```cypher
 Match path = (:(Method|Func)) -[:CALLS*5..10]-> (:(Method|Func))
-Return path 
+Return path
 LIMIT 20
 ```
 
@@ -11501,8 +11501,8 @@ iconType: "solid"
 
 # AI Impact Analysis
 
-This tutorial shows how to use Codegen's attribution extension to analyze the impact of AI on your 
-codebase. You'll learn how to identify which parts of your code were written by AI tools like 
+This tutorial shows how to use Codegen's attribution extension to analyze the impact of AI on your
+codebase. You'll learn how to identify which parts of your code were written by AI tools like
 GitHub Copilot, Devin, or other AI assistants.
 
 Note: the code is flexible - you can track CI pipeline bots, or any other contributor you want.
@@ -11581,7 +11581,7 @@ for symbol in codebase.symbols:
 
 ### Customizing AI Author Detection
 
-By default, the analysis looks for common AI bot names in commit authors. 
+By default, the analysis looks for common AI bot names in commit authors.
 You can customize this by providing your own list of AI authors:
 
 ```python
@@ -11622,34 +11622,34 @@ from codegen.shared.enums.programming_language import ProgrammingLanguage
 def analyze_contributors(codebase):
     """Analyze contributors to the codebase and their impact."""
     print("\n🔍 Contributor Analysis:")
-    
+
     # Define which authors are considered AI
     ai_authors = ['devin[bot]', 'codegen[bot]', 'github-actions[bot]', 'dependabot[bot]']
-    
+
     # Add attribution information to all symbols
     print("Adding attribution information to symbols...")
     add_attribution_to_symbols(codebase, ai_authors)
-    
+
     # Collect statistics about contributors
     contributor_stats = Counter()
     ai_contributor_stats = Counter()
-    
+
     print("Analyzing symbol attributions...")
     for symbol in codebase.symbols:
         if hasattr(symbol, 'last_editor') and symbol.last_editor:
             contributor_stats[symbol.last_editor] += 1
-            
+
             # Track if this is an AI contributor
             if any(ai in symbol.last_editor for ai in ai_authors):
                 ai_contributor_stats[symbol.last_editor] += 1
-    
+
     # Print top contributors overall
     print("\n👥 Top Contributors by Symbols Authored:")
     for contributor, count in contributor_stats.most_common(10):
         is_ai = any(ai in contributor for ai in ai_authors)
         ai_indicator = "🤖" if is_ai else "👤"
         print(f"  {ai_indicator} {contributor}: {count} symbols")
-    
+
     # Print top AI contributors if any
     if ai_contributor_stats:
         print("\n🤖 Top AI Contributors:")
@@ -11661,24 +11661,23 @@ if os.path.exists(".git"):
     repo_path = os.getcwd()
     repo_config = RepoConfig.from_repo_path(repo_path)
     repo_operator = RepoOperator(repo_config=repo_config)
-    
+
     project = ProjectConfig.from_repo_operator(
         repo_operator=repo_operator,
         programming_language=ProgrammingLanguage.PYTHON
     )
     codebase = Codebase(projects=[project])
-    
+
     # Run the contributor analysis
     analyze_contributors(codebase)
 ```
 
 ## Conclusion
 
-The attribution extension provides valuable insights into how AI tools are being used in your 
+The attribution extension provides valuable insights into how AI tools are being used in your
 development process. By understanding which parts of your codebase are authored by AI, you can:
 
 - Track the adoption of AI coding assistants in your team
 - Identify areas where AI is most effective
 - Ensure appropriate review of AI-generated code
 - Measure the impact of AI on developer productivity
-
diff --git a/uv.lock b/uv.lock
index af429489f..e4103ccb6 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1245,6 +1245,7 @@ dev = [
     { name = "pytest-cov" },
     { name = "pytest-lsp" },
     { name = "pytest-mock" },
+    { name = "pytest-split" },
     { name = "pytest-timeout" },
     { name = "pytest-xdist" },
     { name = "ruff" },
@@ -1360,6 +1361,7 @@ dev = [
     { name = "pytest-cov", specifier = ">=6.1.0,<6.1.1" },
     { name = "pytest-lsp", specifier = ">=1.0.0b1" },
     { name = "pytest-mock", specifier = ">=3.14.0,<4.0.0" },
+    { name = "pytest-split", specifier = ">=0.10.0" },
     { name = "pytest-timeout", specifier = ">=2.3.1" },
     { name = "pytest-xdist", specifier = ">=3.6.1,<4.0.0" },
     { name = "ruff", specifier = ">=0.6.8" },
@@ -3318,6 +3320,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/cb/29/518f32faf6edad9f56d6e0107217f7de6b79f297a47170414a2bd4be7f01/pytest_snapshot-0.9.0-py3-none-any.whl", hash = "sha256:4b9fe1c21c868fe53a545e4e3184d36bc1c88946e3f5c1d9dd676962a9b3d4ab", size = 10715 },
 ]
 
+[[package]]
+name = "pytest-split"
+version = "0.10.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pytest" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/46/d7/e30ba44adf83f15aee3f636daea54efadf735769edc0f0a7d98163f61038/pytest_split-0.10.0.tar.gz", hash = "sha256:adf80ba9fef7be89500d571e705b4f963dfa05038edf35e4925817e6b34ea66f", size = 13903 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d6/a7/cad88e9c1109a5c2a320d608daa32e5ee008ccbc766310f54b1cd6b3d69c/pytest_split-0.10.0-py3-none-any.whl", hash = "sha256:466096b086a7147bcd423c6e6c2e57fc62af1c5ea2e256b4ed50fc030fc3dddc", size = 11961 },
+]
+
 [[package]]
 name = "pytest-timeout"
 version = "2.4.0"