diff --git a/pyproject.toml b/pyproject.toml index 40a09e254..afd6ab8e6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,13 @@ dependencies = [ "pyinstrument>=5.0.0", "pip>=24.3.1", # This is needed for some NPM/YARN/PNPM post-install scripts to work! "emoji>=2.14.0", + "rich-click>=1.8.5", + "python-dotenv>=1.0.1", + "giturlparse", + "pygit2>=1.16.0", + "unidiff>=0.7.5", + "datamodel-code-generator>=0.26.5", + "toml>=0.10.2", "PyGithub==2.5.0", "GitPython==3.1.44", ] @@ -62,8 +69,10 @@ classifiers = [ "Intended Audience :: Developers", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Topic :: Software Development", "Development Status :: 4 - Beta", "Environment :: MacOS X", "Programming Language :: Python :: 3", "Programming Language :: Python", ] [project.scripts] +codegen = "codegen.cli.cli:main" gs = "codegen.gscli.main:main" run_string = "codegen.sdk.core.main:main" + [project.optional-dependencies] types = [ "types-networkx>=3.2.1.20240918", diff --git a/src/codegen/cli/__init__.py b/src/codegen/cli/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/codegen/cli/_env.py b/src/codegen/cli/_env.py new file mode 100644 index 000000000..5a12ba1d0 --- /dev/null +++ b/src/codegen/cli/_env.py @@ -0,0 +1 @@ +ENV = "" diff --git a/src/codegen/cli/api/client.py b/src/codegen/cli/api/client.py new file mode 100644 index 000000000..bef1fcbde --- /dev/null +++ b/src/codegen/cli/api/client.py @@ -0,0 +1,245 @@ +import json +from typing import ClassVar, TypeVar + +import requests +from pydantic import BaseModel +from rich import print as rprint + +from codegen.cli.api.endpoints import ( + CREATE_ENDPOINT, + DEPLOY_ENDPOINT, + DOCS_ENDPOINT, + EXPERT_ENDPOINT, + IDENTIFY_ENDPOINT, + LOOKUP_ENDPOINT, + PR_LOOKUP_ENDPOINT, + RUN_ENDPOINT, + RUN_ON_PR_ENDPOINT, +) +from codegen.cli.api.schemas import ( + AskExpertInput, + AskExpertResponse, + CodemodRunType, + CreateInput, + CreateResponse, + DeployInput, + DeployResponse, + DocsInput, + DocsResponse, + IdentifyResponse, + LookupInput, + LookupOutput, + PRLookupInput, + PRLookupResponse, + PRSchema, + RunCodemodInput, + RunCodemodOutput, + RunOnPRInput, + RunOnPRResponse, +) +from codegen.cli.auth.session import CodegenSession +from codegen.cli.codemod.convert import convert_to_ui +from codegen.cli.env.global_env import global_env +from codegen.cli.errors import InvalidTokenError, ServerError +from codegen.cli.utils.codemods import Codemod +from codegen.cli.utils.function_finder import DecoratedFunction + +InputT = TypeVar("InputT", bound=BaseModel) +OutputT = TypeVar("OutputT", bound=BaseModel) + + +class RestAPI: + """Handles auth + validation with the codegen API.""" + + _session: ClassVar[requests.Session] = requests.Session() + + auth_token: str | None = None + + def __init__(self, auth_token: str): + self.auth_token = auth_token + + def _get_headers(self) -> dict[str, str]: + """Get headers with authentication token.""" + return {"Authorization": f"Bearer {self.auth_token}"} + + def _make_request( + self, + method: str, + endpoint: str, + input_data: InputT | None, + output_model: type[OutputT], + ) -> OutputT: + """Make an API request with input validation and response handling.""" + if global_env.DEBUG: + rprint(f"[purple]{method}[/purple] {endpoint}") + if input_data: + rprint(f"{json.dumps(input_data.model_dump(), indent=4)}") + + try: + headers = self._get_headers() + + json_data = input_data.model_dump() if input_data else None + + response = self._session.request( + method, + endpoint, + json=json_data, + headers=headers, + ) + + if response.status_code == 200: + try: + return output_model.model_validate(response.json()) + except ValueError as e: + raise ServerError(f"Invalid response format: {e}") + elif response.status_code == 401: + raise InvalidTokenError("Invalid or expired authentication token") + elif response.status_code == 500: + raise ServerError("The server encountered an error while processing your request") + else: + try: + error_json = response.json() + error_msg = error_json.get("detail", error_json) + except Exception: + error_msg = response.text + raise ServerError(f"Error ({response.status_code}): {error_msg}") + + except requests.RequestException as e: + raise ServerError(f"Network error: {e!s}") + + def run( + self, + function: DecoratedFunction | Codemod, + include_source: bool = True, + run_type: CodemodRunType = CodemodRunType.DIFF, + template_context: dict[str, str] | None = None, + ) -> RunCodemodOutput: + """Run a codemod transformation. + + Args: + function: The function or codemod to run + include_source: Whether to include the source code in the request. + If False, uses the deployed version. + run_type: Type of run (diff or pr) + template_context: Context variables to pass to the codemod + + """ + session = CodegenSession() + + base_input = { + "codemod_name": function.name, + "repo_full_name": session.repo_name, + "codemod_run_type": run_type, + } + + # Only include source if requested + if include_source: + source = function.get_current_source() if isinstance(function, Codemod) else function.source + base_input["codemod_source"] = convert_to_ui(source) + + # Add template context if provided + if template_context: + base_input["template_context"] = template_context + + input_data = RunCodemodInput(input=RunCodemodInput.BaseRunCodemodInput(**base_input)) + return self._make_request( + "POST", + RUN_ENDPOINT, + input_data, + RunCodemodOutput, + ) + + def get_docs(self) -> dict: + """Search documentation.""" + session = CodegenSession() + return self._make_request( + "GET", + DOCS_ENDPOINT, + DocsInput(docs_input=DocsInput.BaseDocsInput(repo_full_name=session.repo_name)), + DocsResponse, + ) + + def ask_expert(self, query: str) -> AskExpertResponse: + """Ask the expert system a question.""" + return self._make_request( + "GET", + EXPERT_ENDPOINT, + AskExpertInput(input=AskExpertInput.BaseAskExpertInput(query=query)), + AskExpertResponse, + ) + + def create(self, name: str, query: str) -> CreateResponse: + """Get AI-generated starter code for a codemod.""" + session = CodegenSession() + return self._make_request( + "GET", + CREATE_ENDPOINT, + CreateInput(input=CreateInput.BaseCreateInput(name=name, query=query, repo_full_name=session.repo_name)), + CreateResponse, + ) + + def identify(self) -> IdentifyResponse | None: + """Identify the user's codemod.""" + return self._make_request( + "POST", + IDENTIFY_ENDPOINT, + None, + IdentifyResponse, + ) + + def deploy( + self, codemod_name: str, codemod_source: str, lint_mode: bool = False, lint_user_whitelist: list[str] | None = None, message: str | None = None, arguments_schema: dict | None = None + ) -> DeployResponse: + """Deploy a codemod to the Modal backend.""" + session = CodegenSession() + return self._make_request( + "POST", + DEPLOY_ENDPOINT, + DeployInput( + input=DeployInput.BaseDeployInput( + codemod_name=codemod_name, + codemod_source=codemod_source, + repo_full_name=session.repo_name, + lint_mode=lint_mode, + lint_user_whitelist=lint_user_whitelist or [], + message=message, + arguments_schema=arguments_schema, + ) + ), + DeployResponse, + ) + + def lookup(self, codemod_name: str) -> LookupOutput: + """Look up a codemod by name.""" + session = CodegenSession() + return self._make_request( + "GET", + LOOKUP_ENDPOINT, + LookupInput(input=LookupInput.BaseLookupInput(codemod_name=codemod_name, repo_full_name=session.repo_name)), + LookupOutput, + ) + + def run_on_pr(self, codemod_name: str, repo_full_name: str, github_pr_number: int, language: str | None = None) -> RunOnPRResponse: + """Test a webhook against a specific PR.""" + return self._make_request( + "POST", + RUN_ON_PR_ENDPOINT, + RunOnPRInput( + input=RunOnPRInput.BaseRunOnPRInput( + codemod_name=codemod_name, + repo_full_name=repo_full_name, + github_pr_number=github_pr_number, + language=language, + ) + ), + RunOnPRResponse, + ) + + def lookup_pr(self, repo_full_name: str, github_pr_number: int) -> PRSchema: + """Look up a PR by repository and PR number.""" + return self._make_request( + "GET", + PR_LOOKUP_ENDPOINT, + PRLookupInput(input=PRLookupInput.BasePRLookupInput(repo_full_name=repo_full_name, github_pr_number=github_pr_number)), + PRLookupResponse, + ) diff --git a/src/codegen/cli/api/endpoints.py b/src/codegen/cli/api/endpoints.py new file mode 100644 index 000000000..c44e0b81a --- /dev/null +++ b/src/codegen/cli/api/endpoints.py @@ -0,0 +1,11 @@ +from codegen.cli.api.modal import MODAL_PREFIX + +RUN_ENDPOINT = f"https://{MODAL_PREFIX}--cli-run.modal.run" +DOCS_ENDPOINT = f"https://{MODAL_PREFIX}--cli-docs.modal.run" +EXPERT_ENDPOINT = f"https://{MODAL_PREFIX}--cli-ask-expert.modal.run" +IDENTIFY_ENDPOINT = f"https://{MODAL_PREFIX}--cli-identify.modal.run" +CREATE_ENDPOINT = f"https://{MODAL_PREFIX}--cli-create.modal.run" +DEPLOY_ENDPOINT = f"https://{MODAL_PREFIX}--cli-deploy.modal.run" +LOOKUP_ENDPOINT = f"https://{MODAL_PREFIX}--cli-lookup.modal.run" +RUN_ON_PR_ENDPOINT = f"https://{MODAL_PREFIX}--cli-run-on-pull-request.modal.run" +PR_LOOKUP_ENDPOINT = f"https://{MODAL_PREFIX}--cli-pr-lookup.modal.run" diff --git a/src/codegen/cli/api/modal.py b/src/codegen/cli/api/modal.py new file mode 100644 index 000000000..534717407 --- /dev/null +++ b/src/codegen/cli/api/modal.py @@ -0,0 +1,24 @@ +from codegen.cli.env.enums import Environment +from codegen.cli.env.global_env import global_env + + +def get_modal_workspace(): + match global_env.ENV: + case Environment.PRODUCTION: + return "codegen-sh" + case Environment.STAGING: + return "codegen-sh-staging" + case Environment.DEVELOP: + return "codegen-sh-develop" + case _: + raise ValueError(f"Invalid environment: {global_env.ENV}") + + +def get_modal_prefix(): + workspace = get_modal_workspace() + if global_env.ENV == Environment.DEVELOP and global_env.MODAL_ENVIRONMENT: + return f"{workspace}-{global_env.MODAL_ENVIRONMENT}" + return workspace + + +MODAL_PREFIX = get_modal_prefix() diff --git a/src/codegen/cli/api/schemas.py b/src/codegen/cli/api/schemas.py new file mode 100644 index 000000000..34583200e --- /dev/null +++ b/src/codegen/cli/api/schemas.py @@ -0,0 +1,236 @@ +from enum import Enum +from typing import TypeVar + +from pydantic import BaseModel, Field + +from codegen.cli.utils.constants import ProgrammingLanguage +from codegen.cli.utils.schema import SafeBaseModel + +T = TypeVar("T") + + +########################################################################### +# RUN +########################################################################### + + +class CodemodRunType(str, Enum): + """Type of codemod run.""" + + DIFF = "diff" + PR = "pr" + + +class RunCodemodInput(SafeBaseModel): + class BaseRunCodemodInput(SafeBaseModel): + repo_full_name: str + codemod_id: int | None = None + codemod_name: str | None = None + codemod_source: str | None = None + codemod_run_type: CodemodRunType = CodemodRunType.DIFF + template_context: dict[str, str] = Field(default_factory=dict) + + input: BaseRunCodemodInput + + +class RunCodemodOutput(SafeBaseModel): + success: bool = False + web_link: str | None = None + logs: str | None = None + observation: str | None = None + error: str | None = None + + +########################################################################### +# EXPERT +########################################################################### + + +class AskExpertInput(SafeBaseModel): + class BaseAskExpertInput(SafeBaseModel): + query: str + + input: BaseAskExpertInput + + +class AskExpertResponse(SafeBaseModel): + response: str + success: bool + + +########################################################################### +# DOCS +########################################################################### + + +class SerializedExample(SafeBaseModel): + name: str | None + description: str | None + source: str + language: ProgrammingLanguage + docstring: str = "" + + +class DocsInput(SafeBaseModel): + class BaseDocsInput(SafeBaseModel): + repo_full_name: str + + docs_input: BaseDocsInput + + +class DocsResponse(SafeBaseModel): + docs: dict[str, str] + examples: list[SerializedExample] + language: ProgrammingLanguage + + +########################################################################### +# CREATE +########################################################################### + + +class CreateInput(SafeBaseModel): + class BaseCreateInput(SafeBaseModel): + name: str + query: str | None = None + repo_full_name: str | None = None + + input: BaseCreateInput + + +class CreateResponse(SafeBaseModel): + success: bool + response: str + code: str + codemod_id: int + context: str | None = None + + +########################################################################### +# IDENTIFY +########################################################################### + + +class IdentifyResponse(SafeBaseModel): + class AuthContext(SafeBaseModel): + token_id: int + expires_at: str + status: str + user_id: int + + class User(SafeBaseModel): + github_user_id: str + avatar_url: str + auth_user_id: str + created_at: str + email: str + is_contractor: str | None + github_username: str + full_name: str | None + id: int + last_updated_at: str | None + + auth_context: AuthContext + user: User + + +########################################################################### +# DEPLOY +########################################################################### + + +class DeployInput(BaseModel): + """Input for deploying a codemod.""" + + class BaseDeployInput(BaseModel): + codemod_name: str = Field(..., description="Name of the codemod to deploy") + codemod_source: str = Field(..., description="Source code of the codemod") + repo_full_name: str = Field(..., description="Full name of the repository") + lint_mode: bool = Field(default=False, description="Whether this is a PR check/lint mode function") + lint_user_whitelist: list[str] = Field(default_factory=list, description="List of GitHub usernames to notify") + message: str | None = Field(default=None, description="Optional message describing the codemod being deployed.") + arguments_schema: dict | None = Field(default=None, description="Schema of the arguments parameter") + + input: BaseDeployInput = Field(..., description="Input data for deployment") + + +class DeployResponse(BaseModel): + """Response from deploying a codemod.""" + + success: bool = Field(..., description="Whether the deployment was successful") + new: bool = Field(..., description="Whether the codemod is newly created") + codemod_id: int = Field(..., description="ID of the deployed codemod") + version_id: int = Field(..., description="Version ID of the deployed codemod") + url: str = Field(..., description="URL of the deployed codemod") + + +########################################################################### +# LOOKUP +########################################################################### + + +class LookupInput(BaseModel): + """Input for looking up a codemod.""" + + class BaseLookupInput(BaseModel): + codemod_name: str = Field(..., description="Name of the codemod to look up") + repo_full_name: str = Field(..., description="Full name of the repository") + + input: BaseLookupInput = Field(..., description="Input data for lookup") + + +class LookupOutput(BaseModel): + """Response from looking up a codemod.""" + + codemod_id: int = Field(..., description="ID of the codemod") + version_id: int = Field(..., description="Version ID of the codemod") + + +########################################################################### +# PR LOOKUP +########################################################################### + + +class PRSchema(BaseModel): + url: str + title: str + body: str + github_pr_number: int + codegen_pr_id: int + + +class PRLookupInput(BaseModel): + class BasePRLookupInput(BaseModel): + repo_full_name: str + github_pr_number: int + + input: BasePRLookupInput + + +class PRLookupResponse(BaseModel): + pr: PRSchema + + +########################################################################### +# TEST WEBHOOK +########################################################################### + + +class RunOnPRInput(BaseModel): + """Input for testing a webhook against a PR.""" + + class BaseRunOnPRInput(BaseModel): + codemod_name: str = Field(..., description="Name of the codemod to test") + repo_full_name: str = Field(..., description="Full name of the repository") + github_pr_number: int = Field(..., description="GitHub PR number to test against") + language: str | None = Field(..., description="Language of the codemod") + + input: BaseRunOnPRInput = Field(..., description="Input data for webhook test") + + +class RunOnPRResponse(BaseModel): + """Response from testing a webhook.""" + + codemod_id: int = Field(..., description="ID of the codemod") + codemod_run_id: int = Field(..., description="ID of the codemod run") + web_url: str = Field(..., description="URL to view the test results") diff --git a/src/codegen/cli/api/webapp_routes.py b/src/codegen/cli/api/webapp_routes.py new file mode 100644 index 000000000..3588fcc0f --- /dev/null +++ b/src/codegen/cli/api/webapp_routes.py @@ -0,0 +1,4 @@ +# Urls linking to the webapp +from codegen.cli.utils.url import generate_webapp_url + +USER_SECRETS_ROUTE = generate_webapp_url(path="cli-token") diff --git a/src/codegen/cli/auth/constants.py b/src/codegen/cli/auth/constants.py new file mode 100644 index 000000000..84849c81c --- /dev/null +++ b/src/codegen/cli/auth/constants.py @@ -0,0 +1,13 @@ +from pathlib import Path + +# Base directories +CONFIG_DIR = Path("~/.config/codegen-sh").expanduser() +CODEGEN_DIR = Path(".codegen") +PROMPTS_DIR = CODEGEN_DIR / "prompts" + +# Subdirectories +DOCS_DIR = CODEGEN_DIR / "docs" +EXAMPLES_DIR = CODEGEN_DIR / "examples" + +# Files +AUTH_FILE = CONFIG_DIR / "auth.json" diff --git a/src/codegen/cli/auth/decorators.py b/src/codegen/cli/auth/decorators.py new file mode 100644 index 000000000..b983fea89 --- /dev/null +++ b/src/codegen/cli/auth/decorators.py @@ -0,0 +1,31 @@ +import functools +from collections.abc import Callable + +import click +import rich + +from codegen.cli.auth.login import login_routine +from codegen.cli.auth.session import CodegenSession +from codegen.cli.errors import AuthError, InvalidTokenError, NoTokenError + + +def requires_auth(f: Callable) -> Callable: + """Decorator that ensures a user is authenticated and injects a CodegenSession.""" + + @functools.wraps(f) + def wrapper(*args, **kwargs): + session = CodegenSession() + + try: + if not session.is_authenticated(): + rich.print("[yellow]Not authenticated. Let's get you logged in first![/yellow]\n") + session = login_routine() + except (InvalidTokenError, NoTokenError) as e: + rich.print("[yellow]Authentication token is invalid or expired. Let's get you logged in again![/yellow]\n") + session = login_routine() + except AuthError as e: + raise click.ClickException(str(e)) + + return f(*args, session=session, **kwargs) + + return wrapper diff --git a/src/codegen/cli/auth/login.py b/src/codegen/cli/auth/login.py new file mode 100644 index 000000000..0f1cc78c1 --- /dev/null +++ b/src/codegen/cli/auth/login.py @@ -0,0 +1,49 @@ +import webbrowser + +import rich +import rich_click as click + +from codegen.cli.api.webapp_routes import USER_SECRETS_ROUTE +from codegen.cli.auth.session import CodegenSession +from codegen.cli.auth.token_manager import TokenManager +from codegen.cli.env.global_env import global_env +from codegen.cli.errors import AuthError + + +def login_routine(token: str | None = None) -> CodegenSession: + """Guide user through login flow and return authenticated session. + + Args: + console: Optional console for output. Creates new one if not provided. + + Returns: + CodegenSession: Authenticated session + + Raises: + click.ClickException: If login fails + + """ + # Try environment variable first + + _token = token or global_env.CODEGEN_USER_ACCESS_TOKEN + + # If no token provided, guide user through browser flow + if not _token: + rich.print(f"Opening {USER_SECRETS_ROUTE} to get your authentication token...") + webbrowser.open_new(USER_SECRETS_ROUTE) + _token = click.prompt("Please enter your authentication token from the browser", hide_input=False) + + if not _token: + raise click.ClickException("Token must be provided via CODEGEN_USER_ACCESS_TOKEN environment variable or manual input") + + # Validate and store token + token_manager = TokenManager() + session = CodegenSession(_token) + + try: + session.assert_authenticated() + token_manager.save_token(_token) + rich.print(f"[green]āœ“ Stored token to:[/green] {token_manager.token_file}") + return session + except AuthError as e: + raise click.ClickException(f"Error: {e!s}") diff --git a/src/codegen/cli/auth/session.ipynb b/src/codegen/cli/auth/session.ipynb new file mode 100644 index 000000000..ffeb6feb8 --- /dev/null +++ b/src/codegen/cli/auth/session.ipynb @@ -0,0 +1,43 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from codegen.cli.auth.session import CodegenSession\n", + "\n", + "\n", + "session = CodegenSession()\n", + "print(session.identity)\n", + "print(session.identity)\n", + "print(session.config)\n", + "print(session.config)\n", + "print(session.profile)\n", + "print(session.profile)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.0" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/src/codegen/cli/auth/session.py b/src/codegen/cli/auth/session.py new file mode 100644 index 000000000..bdf92a8f6 --- /dev/null +++ b/src/codegen/cli/auth/session.py @@ -0,0 +1,133 @@ +from dataclasses import dataclass +from pathlib import Path + +from pygit2.repository import Repository + +from codegen.cli.auth.constants import CODEGEN_DIR +from codegen.cli.auth.token_manager import get_current_token +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 + + +@dataclass +class Identity: + token: str + expires_at: str + status: str + user: "User" + + +@dataclass +class User: + full_name: str + email: str + github_username: str + + +@dataclass +class UserProfile: + """User profile populated from /identity endpoint""" + + name: str + email: str + username: str + + +class CodegenSession: + """Represents an authenticated codegen session with user and repository context""" + + # =====[ Instance attributes ]===== + token: str | None = None + + # =====[ Lazy instance attributes ]===== + _config: Config | None = None + _identity: Identity | None = None + _profile: UserProfile | None = None + + def __init__(self, token: str | None = None): + self.token = token or get_current_token() + + @property + def config(self) -> Config: + """Get the config for the current session""" + if self._config: + return self._config + self._config = get_config(self.codegen_dir) + return self._config + + @property + def identity(self) -> Identity | None: + """Get the identity of the user, if a token has been provided""" + if self._identity: + return self._identity + if not self.token: + raise NoTokenError("No authentication token found") + + from codegen.cli.api.client import RestAPI + + identity = RestAPI(self.token).identify() + if not identity: + return None + + self._identity = Identity( + token=self.token, + expires_at=identity.auth_context.expires_at, + status=identity.auth_context.status, + user=User( + full_name=identity.user.full_name, + email=identity.user.email, + github_username=identity.user.github_username, + ), + ) + return self._identity + + @property + def profile(self) -> UserProfile | None: + """Get the user profile information""" + if self._profile: + return self._profile + if not self.identity: + return None + + self._profile = UserProfile( + name=self.identity.user.full_name, + email=self.identity.user.email, + username=self.identity.user.github_username, + ) + return self._profile + + @property + def git_repo(self) -> Repository: + git_repo = get_git_repo(Path.cwd()) + if not git_repo: + raise ValueError("No git repository found") + return git_repo + + @property + def repo_name(self) -> str: + """Get the current repository name""" + return self.config.repo_full_name + + @property + def codegen_dir(self) -> Path: + """Get the path to the codegen-sh directory""" + return Path.cwd() / CODEGEN_DIR + + def __str__(self) -> str: + return f"CodegenSession(user={self.profile.name}, repo={self.repo_name})" + + def is_authenticated(self) -> bool: + """Check if the session is fully authenticated, including token expiration""" + return bool(self.identity and self.identity.status == "active") + + def assert_authenticated(self) -> None: + """Raise an AuthError if the session is not fully authenticated""" + if not self.identity: + raise AuthError("No identity found for session") + if self.identity.status != "active": + raise AuthError("Current session is not active. API Token may be invalid or may have expired.") + + def write_config(self) -> None: + """Write the config to the codegen-sh/config.toml file""" + write_config(self.config, self.codegen_dir) diff --git a/src/codegen/cli/auth/token_manager.ipynb b/src/codegen/cli/auth/token_manager.ipynb new file mode 100644 index 000000000..145d47608 --- /dev/null +++ b/src/codegen/cli/auth/token_manager.ipynb @@ -0,0 +1,38 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from codegen.cli.auth.token_manager import TokenManager\n", + "\n", + "\n", + "token_manager = TokenManager()\n", + "print(token_manager.get_token())" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.0" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/src/codegen/cli/auth/token_manager.py b/src/codegen/cli/auth/token_manager.py new file mode 100644 index 000000000..04d61c823 --- /dev/null +++ b/src/codegen/cli/auth/token_manager.py @@ -0,0 +1,73 @@ +import json +import os +from pathlib import Path + +from codegen.cli.auth.constants import AUTH_FILE, CONFIG_DIR + + +class TokenManager: + # Simple token manager to store and retrieve tokens. + # This manager checks if the token is expired before retrieval. + # TODO: add support for refreshing token and re authorization via supabase oauth + def __init__(self): + self.config_dir = CONFIG_DIR + self.token_file = AUTH_FILE + self._ensure_config_dir() + + def _ensure_config_dir(self): + """Create config directory if it doesn't exist.""" + if not os.path.exists(self.config_dir): + Path(self.config_dir).mkdir(parents=True, exist_ok=True) + + def save_token(self, token: str) -> None: + """Save api token to disk.""" + try: + with open(self.token_file, "w") as f: + json.dump({"token": token}, f) + + # Secure the file permissions (read/write for owner only) + os.chmod(self.token_file, 0o600) + except Exception as e: + print(f"Error saving token: {e!s}") + raise + + def get_token(self) -> str | None: + """Retrieve token from disk if it exists and is valid.""" + try: + if not os.access(self.config_dir, os.R_OK): + return None + + if not os.path.exists(self.token_file): + return None + + with open(self.token_file) as f: + data = json.load(f) + token = data.get("token") + if not token: + return None + + return token + + except (KeyError, OSError) as e: + print(e) + return None + + def clear_token(self) -> None: + """Remove stored token.""" + if os.path.exists(self.token_file): + os.remove(self.token_file) + + +def get_current_token() -> str | None: + """Get the current authentication token if one exists. + + This is a helper function that creates a TokenManager instance and retrieves + the stored token. The token is validated before being returned. + + Returns: + Optional[str]: The current valid api token if one exists. + Returns None if no token exists. + + """ + token_manager = TokenManager() + return token_manager.get_token() diff --git a/src/codegen/cli/cli.py b/src/codegen/cli/cli.py new file mode 100644 index 000000000..0d115c820 --- /dev/null +++ b/src/codegen/cli/cli.py @@ -0,0 +1,41 @@ +import rich_click as click +from rich.traceback import install + +from codegen.cli.commands.create.main import create_command +from codegen.cli.commands.deploy.main import deploy_command +from codegen.cli.commands.expert.main import expert_command +from codegen.cli.commands.init.main import init_command +from codegen.cli.commands.list.main import list_command +from codegen.cli.commands.login.main import login_command +from codegen.cli.commands.logout.main import logout_command +from codegen.cli.commands.profile.main import profile_command +from codegen.cli.commands.run.main import run_command +from codegen.cli.commands.run_on_pr.main import run_on_pr_command +from codegen.cli.commands.style_debug.main import style_debug_command + +click.rich_click.USE_RICH_MARKUP = True +install(show_locals=True) + + +@click.group() +@click.version_option(prog_name="codegen", message="%(version)s") +def main(): + """Codegen CLI - Transform your code with AI.""" + + +# Wrap commands with error handler +main.add_command(init_command) +main.add_command(logout_command) +main.add_command(login_command) +main.add_command(run_command) +main.add_command(profile_command) +main.add_command(create_command) +main.add_command(expert_command) +main.add_command(list_command) +main.add_command(deploy_command) +main.add_command(style_debug_command) +main.add_command(run_on_pr_command) + + +if __name__ == "__main__": + main() diff --git a/src/codegen/cli/codemod/convert.py b/src/codegen/cli/codemod/convert.py new file mode 100644 index 000000000..fe5d62337 --- /dev/null +++ b/src/codegen/cli/codemod/convert.py @@ -0,0 +1,22 @@ +from textwrap import indent + + +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 + +from graph_sitter import {codebase_type} + +context: Any + + +@codegen.cli.sdk.decorator.function('{name}') +def run(codebase: {codebase_type}, pr_options: Any): +{indent(input, " ")} +""" + + +def convert_to_ui(input: str) -> str: + return input diff --git a/src/codegen/cli/commands/create/main.py b/src/codegen/cli/commands/create/main.py new file mode 100644 index 000000000..d6e1e5416 --- /dev/null +++ b/src/codegen/cli/commands/create/main.py @@ -0,0 +1,128 @@ +from pathlib import Path + +import rich +import rich_click as click + +from codegen.cli.api.client import RestAPI +from codegen.cli.auth.constants import PROMPTS_DIR +from codegen.cli.auth.decorators import requires_auth +from codegen.cli.auth.session import CodegenSession +from codegen.cli.codemod.convert import convert_to_cli +from codegen.cli.errors import ServerError +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.workspace.decorators import requires_init + + +def get_prompts_dir() -> Path: + """Get the directory for storing prompts, creating it if needed.""" + PROMPTS_DIR.mkdir(parents=True, exist_ok=True) + + # Ensure .gitignore exists and contains the prompts directory + gitignore = Path.cwd() / ".gitignore" + if not gitignore.exists() or "codegen-sh/prompts" not in gitignore.read_text(): + with open(gitignore, "a") as f: + f.write("\n# Codegen prompts\ncodegen-sh/prompts/\n") + + return PROMPTS_DIR + + +def get_target_path(name: str, path: Path) -> Path: + """Get the target path for the new function file.""" + # Convert name to snake case for filename + name_snake = name.lower().replace("-", "_").replace(" ", "_") + + if path.suffix == ".py": + # If path is a file, use it directly + return path + else: + # If path is a directory, create name_snake.py in it + return path / f"{name_snake}.py" + + +def make_relative(path: Path) -> str: + """Convert a path to a relative path from cwd, handling non-existent paths.""" + # If it's just a filename in the current directory, return it directly + if str(path.parent) == ".": + return f"./{path.name}" + + try: + return f"./{path.relative_to(Path.cwd())}" + except ValueError: + # For paths in subdirectories, try to make the parent relative + try: + parent_rel = path.parent.relative_to(Path.cwd()) + return f"./{parent_rel}/{path.name}" + except ValueError: + # If all else fails, just return the filename + return f"./{path.name}" + + +@click.command(name="create") +@requires_auth +@requires_init +@click.argument("name", type=str) +@click.argument("path", type=click.Path(path_type=Path), default=Path.cwd()) +@click.option("--description", "-d", default=None, help="Description of what this codemod does.") +@click.option("--overwrite", is_flag=True, help="Overwrites function if it already exists.") +def create_command(session: CodegenSession, name: str, path: Path, description: str | None = None, overwrite: bool = False): + """Create a new codegen function. + + NAME is the name/label for the function + PATH is where to create the function (default: current directory) + """ + # Get the target path for the function + target_path = get_target_path(name, path) + + # Check if file exists + if target_path.exists() and not overwrite: + rel_path = make_relative(target_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 + + if description: + status_message = "Generating function (using LLM, this will take ~30s)" + else: + status_message = "Setting up function" + + rich.print("") # Add a newline before the spinner + with create_spinner(status_message) as status: + try: + # Get code from API + response = RestAPI(session.token).create(name=name, query=description if description else None) + + # Convert the code to include the decorator + code = convert_to_cli(response.code, session.config.programming_language or ProgrammingLanguage.PYTHON, name) + + # Create the target directory if needed + target_path.parent.mkdir(parents=True, exist_ok=True) + + # Write the function code + target_path.write_text(code) + + # Write the system prompt to the prompts directory + if response.context: + prompt_path = get_prompts_dir() / f"{name.lower().replace(' ', '-')}-system-prompt.md" + prompt_path.write_text(response.context) + + except ServerError as e: + status.stop() + raise click.ClickException(str(e)) + except ValueError as e: + status.stop() + raise click.ClickException(str(e)) + + # Success message + rich.print(f"\nāœ… {'Overwrote' if overwrite and target_path.exists() else 'Created'} function '{name}'") + rich.print("") + rich.print("šŸ“ Files Created:") + rich.print(f" [dim]Function:[/dim] {make_relative(target_path)}") + if response.context: + rich.print(f" [dim]Prompt:[/dim] {make_relative(get_prompts_dir() / f'{name.lower().replace(" ", "-")}-system-prompt.md')}") + + # Next steps + rich.print("\n[bold]What's next?[/bold]\n") + rich.print("1. Review and edit the function to customize its behavior") + rich.print(f"2. Run it with: \n{format_command(f'codegen run {name}')}") diff --git a/src/codegen/cli/commands/deploy/main.py b/src/codegen/cli/commands/deploy/main.py new file mode 100644 index 000000000..873b93b55 --- /dev/null +++ b/src/codegen/cli/commands/deploy/main.py @@ -0,0 +1,77 @@ +import time +from pathlib import Path + +import rich +import rich_click as click + +from codegen.cli.api.client import RestAPI +from codegen.cli.auth.decorators import requires_auth +from codegen.cli.auth.session import CodegenSession +from codegen.cli.rich.codeblocks import format_command +from codegen.cli.rich.spinners import create_spinner +from codegen.cli.utils.codemod_manager import CodemodManager +from codegen.cli.utils.function_finder import DecoratedFunction + + +def deploy_functions(session: CodegenSession, functions: list[DecoratedFunction], message: str | None = None) -> None: + """Deploy a list of functions.""" + if not functions: + rich.print("\n[yellow]No @codegen.function decorators found.[/yellow]\n") + return + + # Deploy each function + api_client = RestAPI(session.token) + rich.print() # Add a blank line before deployments + + for func in functions: + with create_spinner(f"Deploying function '{func.name}'...") as status: + start_time = time.time() + response = api_client.deploy( + codemod_name=func.name, + codemod_source=func.source, + lint_mode=func.lint_mode, + lint_user_whitelist=func.lint_user_whitelist, + message=message, + arguments_schema=func.arguments_type_schema, + ) + deploy_time = time.time() - start_time + + func_type = "Webhook" if func.lint_mode else "Function" + rich.print(f"āœ… {func_type} '{func.name}' deployed in {deploy_time:.3f}s! šŸŽ‰") + rich.print(" [dim]View deployment:[/dim]") + rich.print(format_command(f"codegen run {func.name}")) + + +@click.command(name="deploy") +@requires_auth +@click.argument("name", required=False) +@click.option("-d", "--directory", type=click.Path(exists=True, path_type=Path), help="Directory to search for functions") +@click.option("-m", "--message", help="Optional message to include with the deploy") +def deploy_command(session: CodegenSession, name: str | None = None, directory: Path | None = None, message: str | None = None): + """Deploy codegen functions. + + If NAME is provided, deploys a specific function by that name. + If no NAME is provided, deploys all functions in the current directory or specified directory. + """ + try: + search_path = directory or Path.cwd() + + if name: + # Find and deploy specific function by name + functions = CodemodManager.get_decorated(search_path) + matching = [f for f in functions if f.name == name] + if not matching: + raise click.ClickException(f"No function found with name '{name}'") + if len(matching) > 1: + # If multiple matches, show their locations + rich.print(f"[yellow]Multiple functions found with name '{name}':[/yellow]") + for func in matching: + rich.print(f" ā€¢ {func.filepath}") + raise click.ClickException("Please specify the exact directory with --directory") + deploy_functions(session, matching, message=message) + else: + # Deploy all functions in the directory + functions = CodemodManager.get_decorated(search_path) + deploy_functions(session, functions) + except Exception as e: + raise click.ClickException(f"Failed to deploy: {e!s}") diff --git a/src/codegen/cli/commands/expert/main.py b/src/codegen/cli/commands/expert/main.py new file mode 100644 index 000000000..d8e55d9d2 --- /dev/null +++ b/src/codegen/cli/commands/expert/main.py @@ -0,0 +1,30 @@ +import rich +import rich_click as click +from rich.status import Status + +from codegen.cli.api.client import RestAPI +from codegen.cli.auth.decorators import requires_auth +from codegen.cli.auth.session import CodegenSession +from codegen.cli.errors import ServerError +from codegen.cli.workspace.decorators import requires_init + + +@click.command(name="expert") +@click.option("--query", "-q", help="The question to ask the expert.") +@requires_auth +@requires_init +def expert_command(session: CodegenSession, query: str): + """Asks a codegen expert a question.""" + status = Status("Asking expert...", spinner="dots", spinner_style="purple") + status.start() + + try: + response = RestAPI(session.token).ask_expert(query) + status.stop() + rich.print("[bold green]āœ“ Response received[/bold green]") + rich.print(response.response) + except ServerError as e: + status.stop() + raise click.ClickException(str(e)) + finally: + status.stop() diff --git a/src/codegen/cli/commands/init/main.py b/src/codegen/cli/commands/init/main.py new file mode 100644 index 000000000..baee5dc94 --- /dev/null +++ b/src/codegen/cli/commands/init/main.py @@ -0,0 +1,62 @@ +import subprocess +import sys + +import rich +import rich_click as click + +from codegen.cli.auth.decorators import requires_auth +from codegen.cli.auth.session import CodegenSession +from codegen.cli.commands.init.render import get_success_message +from codegen.cli.git.url import get_git_organization_and_repo +from codegen.cli.rich.codeblocks import format_command +from codegen.cli.workspace.initialize_workspace import initialize_codegen + + +@click.command(name="init") +@click.option("--repo-name", type=str, help="The name of the repository") +@click.option("--organization-name", type=str, help="The name of the organization") +@requires_auth +def init_command(session: CodegenSession, repo_name: str | None = None, organization_name: str | None = None): + """Initialize or update the Codegen folder.""" + # Print a message if not in a git repo + try: + subprocess.run(["git", "rev-parse", "--is-inside-work-tree"], capture_output=True, check=True, text=True) + except (subprocess.CalledProcessError, FileNotFoundError): + rich.print("\n[bold red]Error:[/bold red] Not in a git repository") + rich.print("[white]Please run this command from within a git repository.[/white]") + rich.print("\n[dim]To initialize a new git repository:[/dim]") + rich.print(format_command("git init")) + rich.print(format_command("git remote add origin ")) + rich.print(format_command("codegen init")) + sys.exit(1) + + codegen_dir = session.codegen_dir + is_update = codegen_dir.exists() + + if organization_name is not None: + session.config.organization_name = organization_name + if repo_name is not None: + session.config.repo_name = repo_name + if not session.config.organization_name or not session.config.repo_name: + cwd_org, cwd_repo = get_git_organization_and_repo(session.git_repo) + session.config.organization_name = session.config.organization_name or cwd_org + session.config.repo_name = session.config.repo_name or cwd_repo + session.write_config() + + action = "Updating" if is_update else "Initializing" + rich.print("") # Add a newline before the spinner + codegen_dir, docs_dir, examples_dir = initialize_codegen(action=action) + + # Print success message + rich.print(f"āœ… {action} complete") + rich.print(f" [dim]Organization:[/dim] {session.config.organization_name}") + rich.print(f" [dim]Repository:[/dim] {session.config.repo_name}") + rich.print("") + rich.print(get_success_message(codegen_dir, docs_dir, examples_dir)) + + # Print next steps + rich.print("\n[bold]What's next?[/bold]\n") + rich.print("1. Create a function:") + rich.print(format_command('codegen create my-function -d "describe what you want to do"')) + rich.print("2. Run it:") + rich.print(format_command("codegen run my-function --apply-local")) diff --git a/src/codegen/cli/commands/init/render.py b/src/codegen/cli/commands/init/render.py new file mode 100644 index 000000000..c92738b91 --- /dev/null +++ b/src/codegen/cli/commands/init/render.py @@ -0,0 +1,9 @@ +from pathlib import Path + + +def get_success_message(codegen_dir: Path, docs_dir: Path, examples_dir: Path) -> str: + """Get the success message to display after initialization.""" + return """šŸ“ Folders Created: + [dim] Location:[/dim] .codegen-sh + [dim] Docs:[/dim] .codegen-sh/docs + [dim] Examples:[/dim] .codegen-sh/examples""" diff --git a/src/codegen/cli/commands/list/main.py b/src/codegen/cli/commands/list/main.py new file mode 100644 index 000000000..046c1a17e --- /dev/null +++ b/src/codegen/cli/commands/list/main.py @@ -0,0 +1,31 @@ +from pathlib import Path + +import rich +import rich_click as click +from rich.table import Table + +from codegen.cli.rich.codeblocks import format_codeblock, format_command +from codegen.cli.utils.codemod_manager import CodemodManager + + +@click.command(name="list") +def list_command(): + """List available codegen functions.""" + functions = CodemodManager.get_decorated() + if functions: + table = Table(title="Codegen Functions", border_style="blue") + table.add_column("Name", style="cyan") + table.add_column("Type", style="magenta") + table.add_column("Path", style="dim") + + for func in functions: + func_type = "Webhook" if func.lint_mode else "Function" + table.add_row(func.name, func_type, str(func.filepath.relative_to(Path.cwd())) if func.filepath else "") + + rich.print(table) + rich.print("\nRun a function with:") + rich.print(format_command("codegen run