77from typing import Annotated , Any , Dict , List , Optional
88
99import requests
10- from codegen import Codebase
11- from codegen .cli .api .client import RestAPI
12- from codegen .cli .api .endpoints import CODEGEN_SYSTEM_PROMPT_URL
13- from codegen .cli .auth .token_manager import get_current_token
14- from codegen .cli .codemod .convert import convert_to_cli
15- from codegen .cli .utils .default_code import DEFAULT_CODEMOD
16- from codegen .extensions .tools .reveal_symbol import reveal_symbol
10+ from graph_sitter import Codebase
11+ from graph_sitter .cli .api .client import RestAPI
12+ from graph_sitter .cli .api .endpoints import CODEGEN_SYSTEM_PROMPT_URL
13+ from graph_sitter .cli .auth .token_manager import get_current_token
14+ from graph_sitter .cli .codemod .convert import convert_to_cli
15+ from graph_sitter .cli .utils .default_code import DEFAULT_CODEMOD
16+ from graph_sitter .extensions .tools .reveal_symbol import reveal_symbol
1717from mcp .server .fastmcp import FastMCP
1818
1919logger = getLogger (__name__ )
@@ -168,7 +168,7 @@ def codegen_system_prompt():
168168For example, consider the following script:
169169
170170```python
171- from codegen import Codebase
171+ from graph_sitter import Codebase
172172
173173# Codegen builds a complete graph connecting
174174# functions, classes, imports and their relationships
@@ -195,7 +195,7 @@ def codegen_system_prompt():
195195
196196They live in the `.codegen/codemods/{name}/{name.py}` directory, and take the following form:
197197```python
198- from codegen import Codebase
198+ from graph_sitter import Codebase
199199
200200@codegen.function('{name}')
201201def codemod(codebase: Codebase):
@@ -364,7 +364,7 @@ async def run_codemod(
364364 try :
365365 # Get the codemod using CodemodManager
366366 try :
367- from codegen .cli .utils .codemod_manager import CodemodManager
367+ from graph_sitter .cli .utils .codemod_manager import CodemodManager
368368
369369 codemod = CodemodManager .get_codemod (name , start_path = state .parsed_codebase .repo_path )
370370 except Exception as e :
@@ -378,15 +378,15 @@ async def run_codemod(
378378
379379 # Validate arguments if schema exists
380380 if codemod .arguments_type_schema :
381- from codegen .cli .utils .json_schema import validate_json
381+ from graph_sitter .cli .utils .json_schema import validate_json
382382
383383 if not validate_json (codemod .arguments_type_schema , args_dict ):
384384 return {"error" : f"Invalid arguments format. Expected schema: { codemod .arguments_type_schema } " }
385385 except json .JSONDecodeError :
386386 return {"error" : "Invalid JSON in arguments parameter" }
387387
388388 # Create a session for the codemod
389- from codegen .cli .auth .session import CodegenSession
389+ from graph_sitter .cli .auth .session import CodegenSession
390390
391391 session = CodegenSession (state .parsed_codebase .repo_path )
392392 session .codebase = state .parsed_codebase
@@ -417,8 +417,8 @@ async def run_codemod(
417417async def reset () -> Dict [str , Any ]:
418418 try :
419419 # Import necessary functions from reset command
420- from codegen .cli .commands .reset .main import backup_codegen_files , remove_untracked_files , restore_codegen_files
421- from codegen .cli .git .repo import get_git_repo
420+ from graph_sitter .cli .commands .reset .main import backup_codegen_files , remove_untracked_files , restore_codegen_files
421+ from graph_sitter .cli .git .repo import get_git_repo
422422 from pygit2 .enums import ResetMode
423423
424424 # Get the git repository
0 commit comments