diff --git a/src/codegen/cli/auth/decorators.py b/src/codegen/cli/auth/decorators.py index d91674121..b12cf35ed 100644 --- a/src/codegen/cli/auth/decorators.py +++ b/src/codegen/cli/auth/decorators.py @@ -19,7 +19,7 @@ def wrapper(*args, **kwargs): session = CodegenSession.from_active_session() # Check for valid session - if session is None or not session.is_valid(): + if session is None: pretty_print_error("There is currently no active session.\nPlease run 'codegen init' to initialize the project.") raise click.Abort() diff --git a/src/codegen/cli/auth/session.py b/src/codegen/cli/auth/session.py index 421f08181..e34b4b982 100644 --- a/src/codegen/cli/auth/session.py +++ b/src/codegen/cli/auth/session.py @@ -1,8 +1,13 @@ from pathlib import Path -from pygit2.repository import Repository +import click +import rich +from github import BadCredentialsException +from github.MainClass import Github from codegen.cli.git.repo import get_git_repo +from codegen.cli.rich.codeblocks import format_command +from codegen.git.repo_operator.local_git_repo import LocalGitRepo from codegen.shared.configs.constants import CODEGEN_DIR_NAME, CONFIG_FILENAME from codegen.shared.configs.models.session import SessionConfig from codegen.shared.configs.session_configs import global_config, load_session_config @@ -11,16 +16,25 @@ class CodegenSession: """Represents an authenticated codegen session with user and repository context""" - repo_path: Path # TODO: rename to root_path + repo_path: Path + local_git: LocalGitRepo codegen_dir: Path config: SessionConfig existing: bool - def __init__(self, repo_path: Path): + def __init__(self, repo_path: Path, git_token: str | None = None) -> None: + if not repo_path.exists() or get_git_repo(repo_path) is None: + rich.print(f"\n[bold red]Error:[/bold red] Path to git repo does not exist at {self.repo_path}") + raise click.Abort() + self.repo_path = repo_path + self.local_git = LocalGitRepo(repo_path=repo_path) self.codegen_dir = repo_path / CODEGEN_DIR_NAME - self.existing = global_config.get_session(repo_path) is not None self.config = load_session_config(self.codegen_dir / CONFIG_FILENAME) + self.config.secrets.github_token = git_token or self.config.secrets.github_token + self.existing = global_config.get_session(repo_path) is not None + + self._initialize() global_config.set_active_session(repo_path) @classmethod @@ -31,19 +45,52 @@ def from_active_session(cls) -> "CodegenSession | None": return cls(active_session) - def is_valid(self) -> bool: - """Validates that the session configuration is correct""" - # TODO: also make sure all the expected prompt, jupyter, codemods are present - # TODO: make sure there is still a git instance here. - return self.repo_path.exists() and self.codegen_dir.exists() and Path(self.config.file_path).exists() - - @property - def git_repo(self) -> Repository: - git_repo = get_git_repo(Path.cwd()) - if not git_repo: - msg = "No git repository found" - raise ValueError(msg) - return git_repo + def _initialize(self) -> None: + """Initialize the codegen session""" + self._validate() + + self.config.repository.repo_path = self.config.repository.repo_path or str(self.local_git.repo_path) + self.config.repository.repo_name = self.config.repository.repo_name or self.local_git.name + self.config.repository.full_name = self.config.repository.full_name or self.local_git.full_name + self.config.repository.user_name = self.config.repository.user_name or self.local_git.user_name + self.config.repository.user_email = self.config.repository.user_email or self.local_git.user_email + self.config.repository.language = self.config.repository.language or self.local_git.get_language(access_token=self.config.secrets.github_token).upper() + self.config.save() + + def _validate(self) -> None: + """Validates that the session configuration is correct, otherwise raises an error""" + if not self.codegen_dir.exists(): + rich.print(f"\n[bold red]Error:[/bold red] Codegen folder is missing at {self.codegen_dir}") + raise click.Abort() + + if not Path(self.config.file_path).exists(): + rich.print(f"\n[bold red]Error:[/bold red] Missing config.toml at {self.codegen_dir}") + rich.print("[white]Please remove the codegen folder and reinitialize.[/white]") + rich.print(format_command(f"rm -rf {self.codegen_dir} && codegen init")) + raise click.Abort() + + git_token = self.config.secrets.github_token + if git_token is None: + rich.print("\n[bold yellow]Warning:[/bold yellow] GitHub token not found") + rich.print("To enable full functionality, please set your GitHub token:") + rich.print(format_command("export CODEGEN_SECRETS__GITHUB_TOKEN=")) + rich.print("Or pass in as a parameter:") + rich.print(format_command("codegen init --token ")) + raise click.Abort() + + if self.local_git.origin_remote is None: + rich.print("\n[bold red]Error:[/bold red] No remote found for repository") + rich.print("[white]Please add a remote to the repository.[/white]") + rich.print("\n[dim]To add a remote to the repository:[/dim]") + rich.print(format_command("git remote add origin ")) + raise click.Abort() + + try: + Github(login_or_token=git_token).get_repo(self.local_git.full_name) + except BadCredentialsException: + rich.print(format_command(f"\n[bold red]Error:[/bold red] Invalid GitHub token={git_token} for repo={self.local_git.full_name}")) + rich.print("[white]Please provide a valid GitHub token for this repository.[/white]") + raise click.Abort() def __str__(self) -> str: return f"CodegenSession(user={self.config.repository.user_name}, repo={self.config.repository.repo_name})" diff --git a/src/codegen/cli/commands/init/main.py b/src/codegen/cli/commands/init/main.py index d79cc6261..e85de1c4e 100644 --- a/src/codegen/cli/commands/init/main.py +++ b/src/codegen/cli/commands/init/main.py @@ -3,14 +3,11 @@ import rich import rich_click as click -from github import BadCredentialsException -from github.MainClass import Github from codegen.cli.auth.session import CodegenSession from codegen.cli.commands.init.render import get_success_message from codegen.cli.rich.codeblocks import format_command from codegen.cli.workspace.initialize_workspace import initialize_codegen -from codegen.git.repo_operator.local_git_repo import LocalGitRepo from codegen.git.utils.path import get_git_root_path @@ -25,6 +22,7 @@ def init_command(path: str | None = None, token: str | None = None, language: st path = Path.cwd() if path is None else Path(path) repo_path = get_git_root_path(path) rich.print(f"Found git repository at: {repo_path}") + if repo_path is None: rich.print(f"\n[bold red]Error:[/bold red] Path={path} is not in a git repository") rich.print("[white]Please run this command from within a git repository.[/white]") @@ -33,45 +31,10 @@ def init_command(path: str | None = None, token: str | None = None, language: st rich.print(format_command("codegen init")) sys.exit(1) - local_git = LocalGitRepo(repo_path=repo_path) - if local_git.origin_remote is None: - rich.print("\n[bold red]Error:[/bold red] No remote found for repository") - rich.print("[white]Please add a remote to the repository.[/white]") - rich.print("\n[dim]To add a remote to the repository:[/dim]") - rich.print(format_command("git remote add origin ")) - sys.exit(1) - - session = CodegenSession(repo_path=repo_path) - config = session.config - - if token is None: - token = config.secrets.github_token - else: - config.secrets.github_token = token - - if not token: - rich.print("\n[bold yellow]Warning:[/bold yellow] GitHub token not found") - rich.print("To enable full functionality, please set your GitHub token:") - rich.print(format_command("export CODEGEN_SECRETS__GITHUB_TOKEN=")) - rich.print("Or pass in as a parameter:") - rich.print(format_command("codegen init --token ")) - else: - # Validate that the token is valid for the repo path. - try: - Github(login_or_token=token).get_repo(local_git.full_name) - except BadCredentialsException: - rich.print(format_command(f"\n[bold red]Error:[/bold red] Invalid GitHub token={token} for repo={local_git.full_name}")) - rich.print("[white]Please provide a valid GitHub token for this repository.[/white]") - sys.exit(1) - - # Save repo config - config.repository.repo_path = str(local_git.repo_path) - config.repository.repo_name = local_git.name - config.repository.full_name = local_git.full_name - config.repository.user_name = local_git.user_name - config.repository.user_email = local_git.user_email - config.repository.language = (language or local_git.get_language(access_token=token)).upper() - config.save() + session = CodegenSession(repo_path=repo_path, git_token=token) + if language: + session.config.repository.language = language.upper() + session.config.save() action = "Updating" if session.existing else "Initializing" codegen_dir, docs_dir, examples_dir = initialize_codegen(status=action, session=session, fetch_docs=fetch_docs) diff --git a/src/codegen/cli/commands/run/run_cloud.py b/src/codegen/cli/commands/run/run_cloud.py index d9e4e3533..5250c1543 100644 --- a/src/codegen/cli/commands/run/run_cloud.py +++ b/src/codegen/cli/commands/run/run_cloud.py @@ -2,6 +2,7 @@ import rich import rich_click as click +from pygit2._pygit2 import Repository from rich.panel import Panel from codegen.cli.api.client import RestAPI @@ -87,7 +88,8 @@ def run_cloud(session: CodegenSession, function, apply_local: bool = False, diff if apply_local and run_output.observation: try: - apply_patch(session.git_repo, f"\n{run_output.observation}\n") + git_repo = Repository(str(session.repo_path)) + apply_patch(git_repo, f"\n{run_output.observation}\n") rich.print("") rich.print("[green]✓ Changes have been applied to your local filesystem[/green]") rich.print("[yellow]→ Don't forget to commit your changes:[/yellow]") diff --git a/src/codegen/cli/git/url.py b/src/codegen/cli/git/url.py deleted file mode 100644 index b3334ed2d..000000000 --- a/src/codegen/cli/git/url.py +++ /dev/null @@ -1,11 +0,0 @@ -import giturlparse -from pygit2.repository import Repository - - -def get_git_organization_and_repo(repo: Repository) -> tuple[str, str]: - for remote in repo.remotes: - if remote.name == "origin": - parsed = giturlparse.parse(remote.url) - return parsed.owner, parsed.name - msg = "No git remote found" - raise ValueError(msg) diff --git a/src/codegen/cli/workspace/decorators.py b/src/codegen/cli/workspace/decorators.py index b707adae6..128db7d0b 100644 --- a/src/codegen/cli/workspace/decorators.py +++ b/src/codegen/cli/workspace/decorators.py @@ -2,8 +2,6 @@ import sys from collections.abc import Callable -import click - from codegen.cli.auth.session import CodegenSession from codegen.cli.rich.pretty_print import pretty_print_error @@ -19,11 +17,6 @@ def wrapper(*args, **kwargs): pretty_print_error("Codegen not initialized. Please run `codegen init` from a git repo workspace.") sys.exit(1) - # Check for valid session - if not session.is_valid(): - pretty_print_error(f"The session at path {session.repo_path} is missing or corrupt.\nPlease run 'codegen init' to re-initialize the project.") - raise click.Abort() - kwargs["session"] = session return f(*args, **kwargs) diff --git a/src/codegen/cli/workspace/initialize_workspace.py b/src/codegen/cli/workspace/initialize_workspace.py index b29c2e423..aea224fd1 100644 --- a/src/codegen/cli/workspace/initialize_workspace.py +++ b/src/codegen/cli/workspace/initialize_workspace.py @@ -4,15 +4,13 @@ import requests import rich -import toml from rich.status import Status from codegen.cli.api.client import RestAPI +from codegen.cli.api.endpoints import CODEGEN_SYSTEM_PROMPT_URL from codegen.cli.auth.constants import CODEGEN_DIR, DOCS_DIR, EXAMPLES_DIR, PROMPTS_DIR from codegen.cli.auth.session import CodegenSession from codegen.cli.auth.token_manager import get_current_token -from codegen.cli.git.repo import get_git_repo -from codegen.cli.git.url import get_git_organization_and_repo from codegen.cli.rich.spinners import create_spinner from codegen.cli.utils.notebooks import create_notebook from codegen.cli.workspace.docs_workspace import populate_api_docs @@ -31,12 +29,10 @@ def initialize_codegen(session: CodegenSession, status: Status | str = "Initiali Returns: Tuple of (codegen_folder, docs_folder, examples_folder) """ - repo = get_git_repo() CODEGEN_FOLDER = session.repo_path / CODEGEN_DIR PROMPTS_FOLDER = session.repo_path / PROMPTS_DIR DOCS_FOLDER = session.repo_path / DOCS_DIR EXAMPLES_FOLDER = session.repo_path / EXAMPLES_DIR - CONFIG_PATH = CODEGEN_FOLDER / "config.toml" JUPYTER_DIR = CODEGEN_FOLDER / "jupyter" CODEMODS_DIR = CODEGEN_FOLDER / "codemods" SYSTEM_PROMPT_PATH = CODEGEN_FOLDER / "codegen-system-prompt.txt" @@ -62,35 +58,17 @@ def initialize_codegen(session: CodegenSession, status: Status | str = "Initiali # Download system prompt try: - from codegen.cli.api.endpoints import CODEGEN_SYSTEM_PROMPT_URL - response = requests.get(CODEGEN_SYSTEM_PROMPT_URL) response.raise_for_status() SYSTEM_PROMPT_PATH.write_text(response.text) except Exception as e: rich.print(f"[yellow]Warning: Could not download system prompt: {e}[/yellow]") - if not repo: - rich.print("No git repository found. Please run this command in a git repository.") - else: - status_obj.update(f" {'Updating' if isinstance(status, Status) else status} .gitignore...") - modify_gitignore(CODEGEN_FOLDER) - - # Create or update config.toml with basic repo info - org_name, repo_name = get_git_organization_and_repo(repo) - config = {} - if CONFIG_PATH.exists(): - config = toml.load(CONFIG_PATH) - config.update( - { - "organization_name": config.get("organization_name", org_name), - "repo_name": config.get("repo_name", repo_name), - } - ) - CONFIG_PATH.write_text(toml.dumps(config)) - - # Create notebook template - create_notebook(JUPYTER_DIR) + status_obj.update(f" {'Updating' if isinstance(status, Status) else status} .gitignore...") + modify_gitignore(CODEGEN_FOLDER) + + # Create notebook template + create_notebook(JUPYTER_DIR) # Only fetch docs and examples if requested and session is provided if fetch_docs and session: diff --git a/src/codegen/git/repo_operator/local_git_repo.py b/src/codegen/git/repo_operator/local_git_repo.py index 864f4ae17..d1afb1a81 100644 --- a/src/codegen/git/repo_operator/local_git_repo.py +++ b/src/codegen/git/repo_operator/local_git_repo.py @@ -2,6 +2,7 @@ from functools import cached_property from pathlib import Path +import giturlparse from git import Repo from git.remote import Remote @@ -30,8 +31,8 @@ def full_name(self) -> str | None: if not self.origin_remote: return None - url_segments = self.origin_remote.url.split("/") - return f"{url_segments[-2]}/{url_segments[-1].replace('.git', '')}" + parsed = giturlparse.parse(self.origin_remote.url) + return f"{parsed.owner}/{parsed.name}" @cached_property def origin_remote(self) -> Remote | None: