diff --git a/README.md b/README.md index faae1a232..7fad15c74 100644 --- a/README.md +++ b/README.md @@ -16,9 +16,9 @@
-[![PyPI](https://img.shields.io/pypi/v/codegen?style=flat-square&color=blue)](https://pypi.org/project/codegen/) -[![Documentation](https://img.shields.io/badge/docs-docs.codegen.com-purple?style=flat-square)](https://docs.codegen.com) -[![Slack Community](https://img.shields.io/badge/slack-community-4A154B?logo=slack&style=flat-square)](https://community.codegen.com) +[![PyPI](https://img.shields.io/pypi/v/codegen?style=flat-square&color=blue)](https://pypi.org/project/codegen/) +[![Documentation](https://img.shields.io/badge/docs-docs.codegen.com-purple?style=flat-square)](https://docs.codegen.com) +[![Slack Community](https://img.shields.io/badge/slack-community-4A154B?logo=slack&style=flat-square)](https://community.codegen.com) [![Follow on X](https://img.shields.io/twitter/follow/codegen?style=social)](https://x.com/codegen)
diff --git a/docs/cli/create.mdx b/docs/cli/create.mdx index 5b7b5dd55..4a6e14b7d 100644 --- a/docs/cli/create.mdx +++ b/docs/cli/create.mdx @@ -33,8 +33,8 @@ When you run `codegen create rename-function`, it creates: .codegen/ └── codemods/ └── rename_function/ - ├── rename_function.py # The codemod implementation - └── rename_function_prompt.md # System prompt (if --description used) + ├── rename_function.py # The codemod implementation + └── rename_function-system-prompt.txt # System prompt (if --description used) ``` The generated codemod will have this structure: @@ -47,7 +47,16 @@ from codegen import Codebase def run(codebase: Codebase): """Add a description of what this codemod does.""" # Add your code here - pass + print('Total files: ', len(codebase.files)) + print('Total functions: ', len(codebase.functions)) + print('Total imports: ', len(codebase.imports)) + +if __name__ == "__main__": + print('Parsing codebase...') + codebase = Codebase("./") + + print('Running...') + run(codebase) ``` ## Examples diff --git a/docs/cli/notebook.mdx b/docs/cli/notebook.mdx index 343856970..4941f1a94 100644 --- a/docs/cli/notebook.mdx +++ b/docs/cli/notebook.mdx @@ -1,5 +1,5 @@ --- -title: "notebook" +title: "Notebook Command" sidebarTitle: "notebook" description: "Open a Jupyter notebook with the current codebase loaded" icon: "book" diff --git a/src/codegen/cli/api/client.py b/src/codegen/cli/api/client.py index bef1fcbde..ff81459ab 100644 --- a/src/codegen/cli/api/client.py +++ b/src/codegen/cli/api/client.py @@ -174,7 +174,7 @@ def create(self, name: str, query: str) -> CreateResponse: return self._make_request( "GET", CREATE_ENDPOINT, - CreateInput(input=CreateInput.BaseCreateInput(name=name, query=query, repo_full_name=session.repo_name)), + CreateInput(input=CreateInput.BaseCreateInput(name=name, query=query, language=session.language)), CreateResponse, ) diff --git a/src/codegen/cli/api/schemas.py b/src/codegen/cli/api/schemas.py index 34583200e..9fffd418e 100644 --- a/src/codegen/cli/api/schemas.py +++ b/src/codegen/cli/api/schemas.py @@ -92,8 +92,8 @@ class DocsResponse(SafeBaseModel): class CreateInput(SafeBaseModel): class BaseCreateInput(SafeBaseModel): name: str - query: str | None = None - repo_full_name: str | None = None + query: str + language: ProgrammingLanguage input: BaseCreateInput @@ -102,8 +102,7 @@ class CreateResponse(SafeBaseModel): success: bool response: str code: str - codemod_id: int - context: str | None = None + context: str ########################################################################### diff --git a/src/codegen/cli/auth/session.py b/src/codegen/cli/auth/session.py index bdf92a8f6..07dddb3b1 100644 --- a/src/codegen/cli/auth/session.py +++ b/src/codegen/cli/auth/session.py @@ -8,6 +8,7 @@ from codegen.cli.errors import AuthError, NoTokenError from codegen.cli.git.repo import get_git_repo from codegen.cli.utils.config import Config, get_config, write_config +from codegen.sdk.enums import ProgrammingLanguage @dataclass @@ -109,6 +110,13 @@ def repo_name(self) -> str: """Get the current repository name""" return self.config.repo_full_name + @property + def language(self) -> ProgrammingLanguage: + """Get the current language""" + # TODO(jayhack): This is a temporary solution to get the language. + # We should eventually get the language on init. + return self.config.programming_language or ProgrammingLanguage.PYTHON + @property def codegen_dir(self) -> Path: """Get the path to the codegen-sh directory""" diff --git a/src/codegen/cli/codemod/convert.py b/src/codegen/cli/codemod/convert.py index 512eeebd4..15c80f4d5 100644 --- a/src/codegen/cli/codemod/convert.py +++ b/src/codegen/cli/codemod/convert.py @@ -2,19 +2,21 @@ def convert_to_cli(input: str, language: str, name: str) -> str: - codebase_type = "PyCodebaseType" if language.lower() == "python" else "TSCodebaseType" - return f"""import codegen.cli.sdk.decorator -# from app.codemod.compilation.models.context import CodemodContext -#from app.codemod.compilation.models.pr_options import PROptions + return f"""import codegen +from codegen import Codebase -from codegen.sdk import {codebase_type} -context: Any +@codegen.function('{name}') +def run(codebase: Codebase): +{indent(input, " ")} -@codegen.cli.sdk.decorator.function('{name}') -def run(codebase: {codebase_type}, pr_options: Any): -{indent(input, " ")} +if __name__ == "__main__": + print('Parsing codebase...') + codebase = Codebase("./") + + print('Running function...') + codegen.run(run) """ diff --git a/src/codegen/cli/commands/create/main.py b/src/codegen/cli/commands/create/main.py index 14e34510c..390276a2d 100644 --- a/src/codegen/cli/commands/create/main.py +++ b/src/codegen/cli/commands/create/main.py @@ -11,7 +11,6 @@ from codegen.cli.rich.codeblocks import format_command, format_path from codegen.cli.rich.pretty_print import pretty_print_error from codegen.cli.rich.spinners import create_spinner -from codegen.cli.utils.constants import ProgrammingLanguage from codegen.cli.utils.default_code import DEFAULT_CODEMOD from codegen.cli.workspace.decorators import requires_init @@ -29,7 +28,7 @@ def get_prompts_dir() -> Path: return PROMPTS_DIR -def get_target_path(name: str, path: Path) -> Path: +def get_target_paths(name: str, path: Path) -> tuple[Path, Path]: """Get the target path for the new function file. Creates a directory structure like: @@ -47,7 +46,9 @@ def get_target_path(name: str, path: Path) -> Path: # Create path within .codegen/codemods codemods_dir = base_dir / ".codegen" / "codemods" function_dir = codemods_dir / name_snake - return function_dir / f"{name_snake}.py" + codemod_path = function_dir / f"{name_snake}.py" + prompt_path = function_dir / f"{name_snake}-system-prompt.txt" + return codemod_path, prompt_path def make_relative(path: Path) -> str: @@ -76,11 +77,11 @@ def create_command(session: CodegenSession, name: str, path: Path, description: PATH is where to create the function (default: current directory) """ # Get the target path for the function - target_path = get_target_path(name, path) + codemod_path, prompt_path = get_target_paths(name, path) # Check if file exists - if target_path.exists() and not overwrite: - rel_path = make_relative(target_path) + if codemod_path.exists() and not overwrite: + rel_path = make_relative(codemod_path) pretty_print_error(f"File already exists at {format_path(rel_path)}\n\nTo overwrite the file:\n{format_command(f'codegen create {name} {rel_path} --overwrite')}") return @@ -89,34 +90,30 @@ def create_command(session: CodegenSession, name: str, path: Path, description: try: if description: # Use API to generate implementation - with create_spinner("Generating function (using LLM, this will take ~30s)") as status: + with create_spinner("Generating function (using LLM, this will take ~10s)") as status: response = RestAPI(session.token).create(name=name, query=description) - code = convert_to_cli(response.code, session.config.programming_language or ProgrammingLanguage.PYTHON, name) - - # Write the system prompt if provided - if response.context: - prompt_path = get_prompts_dir() / f"{name.lower().replace(' ', '-')}-system-prompt.md" - prompt_path.write_text(response.context) + code = convert_to_cli(response.code, session.language, name) else: # Use default implementation code = DEFAULT_CODEMOD.format(name=name) # Create the target directory if needed - target_path.parent.mkdir(parents=True, exist_ok=True) + codemod_path.parent.mkdir(parents=True, exist_ok=True) # Write the function code - target_path.write_text(code) + codemod_path.write_text(code) + prompt_path.write_text(response.context) except (ServerError, ValueError) as e: raise click.ClickException(str(e)) # Success message - rich.print(f"\n✅ {'Overwrote' if overwrite and target_path.exists() else 'Created'} function '{name}'") + rich.print(f"\n✅ {'Overwrote' if overwrite and codemod_path.exists() else 'Created'} function '{name}'") rich.print("") rich.print("📁 Files Created:") - rich.print(f" [dim]Function:[/dim] {make_relative(target_path)}") + rich.print(f" [dim]Function:[/dim] {make_relative(codemod_path)}") if description and response.context: - rich.print(f" [dim]Prompt:[/dim] {make_relative(get_prompts_dir() / f'{name.lower().replace(" ", "-")}-system-prompt.md')}") + rich.print(f" [dim]Prompt:[/dim] {make_relative(prompt_path)}") # Next steps rich.print("\n[bold]What's next?[/bold]\n")