Skip to content

feat: enables codegen create -d #135

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,9 @@

<div align="center">

[![PyPI](https://img.shields.io/pypi/v/codegen?style=flat-square&color=blue)](https://pypi.org/project/codegen/)
[![Documentation](https://img.shields.io/badge/docs-docs.codegen.com-purple?style=flat-square)](https://docs.codegen.com)
[![Slack Community](https://img.shields.io/badge/slack-community-4A154B?logo=slack&style=flat-square)](https://community.codegen.com)
[![PyPI](https://img.shields.io/pypi/v/codegen?style=flat-square&color=blue)](https://pypi.org/project/codegen/)
[![Documentation](https://img.shields.io/badge/docs-docs.codegen.com-purple?style=flat-square)](https://docs.codegen.com)
[![Slack Community](https://img.shields.io/badge/slack-community-4A154B?logo=slack&style=flat-square)](https://community.codegen.com)
[![Follow on X](https://img.shields.io/twitter/follow/codegen?style=social)](https://x.com/codegen)

</div>
Expand Down
15 changes: 12 additions & 3 deletions docs/cli/create.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion docs/cli/notebook.mdx
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
---
title: "notebook"
title: "Notebook Command"
sidebarTitle: "notebook"
description: "Open a Jupyter notebook with the current codebase loaded"
icon: "book"
Expand Down
2 changes: 1 addition & 1 deletion src/codegen/cli/api/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand Down
7 changes: 3 additions & 4 deletions src/codegen/cli/api/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -102,8 +102,7 @@ class CreateResponse(SafeBaseModel):
success: bool
response: str
code: str
codemod_id: int
context: str | None = None
context: str


###########################################################################
Expand Down
8 changes: 8 additions & 0 deletions src/codegen/cli/auth/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"""
Expand Down
20 changes: 11 additions & 9 deletions src/codegen/cli/codemod/convert.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
"""


Expand Down
33 changes: 15 additions & 18 deletions src/codegen/cli/commands/create/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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

Expand All @@ -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")
Expand Down