From 7ffdb6770cba8c8e5f3da08a918b35634598879e Mon Sep 17 00:00:00 2001 From: Carol Jung Date: Wed, 12 Feb 2025 16:09:41 -0800 Subject: [PATCH 1/4] wip --- src/codegen/cli/auth/session.py | 46 ++++++++++++++++------- src/codegen/cli/commands/init/main.py | 3 +- src/codegen/cli/commands/run/run_cloud.py | 4 +- 3 files changed, 37 insertions(+), 16 deletions(-) diff --git a/src/codegen/cli/auth/session.py b/src/codegen/cli/auth/session.py index 421f08181..415252056 100644 --- a/src/codegen/cli/auth/session.py +++ b/src/codegen/cli/auth/session.py @@ -1,8 +1,7 @@ from pathlib import Path -from pygit2.repository import Repository - from codegen.cli.git.repo import get_git_repo +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,7 +10,7 @@ class CodegenSession: """Represents an authenticated codegen session with user and repository context""" - repo_path: Path # TODO: rename to root_path + repo_path: Path codegen_dir: Path config: SessionConfig existing: bool @@ -23,6 +22,9 @@ def __init__(self, repo_path: Path): self.config = load_session_config(self.codegen_dir / CONFIG_FILENAME) global_config.set_active_session(repo_path) + if not self.existing: + self._initialize_repo_config() + @classmethod def from_active_session(cls) -> "CodegenSession | None": active_session = global_config.get_active_session() @@ -31,19 +33,35 @@ def from_active_session(cls) -> "CodegenSession | None": return cls(active_session) + @classmethod + def from_local_git(cls, git_repo: LocalGitRepo, token: str | None = None) -> "CodegenSession": + session = cls(git_repo.repo_path) + session._initialize_repo_config(git_repo, token) + return session + + def _initialize_repo_config(self, git_repo: LocalGitRepo, token: str | None = None): + """Initialize the codegen session""" + self.config.repository.repo_path = str(git_repo.repo_path) + self.config.repository.repo_name = git_repo.name + self.config.repository.full_name = git_repo.full_name + self.config.repository.user_name = git_repo.user_name + self.config.repository.user_email = git_repo.user_email + self.config.repository.language = git_repo.get_language(access_token=token) + self.config.save() + 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 + if not self.repo_path.exists(): + return False + + if not self.codegen_dir.exists(): + return False + + if not Path(self.config.file_path).exists(): + return False + + if get_git_repo(self.repo_path) is None: + return False 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..7ad16914f 100644 --- a/src/codegen/cli/commands/init/main.py +++ b/src/codegen/cli/commands/init/main.py @@ -25,6 +25,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]") @@ -41,7 +42,7 @@ def init_command(path: str | None = None, token: str | None = None, language: st rich.print(format_command("git remote add origin ")) sys.exit(1) - session = CodegenSession(repo_path=repo_path) + session = CodegenSession.from_local_git(local_git) config = session.config if token is None: 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]") From a72ef8ee50afa7d9f53c33a0852393ddefd5b9ef Mon Sep 17 00:00:00 2001 From: Carol Jung Date: Wed, 12 Feb 2025 17:33:31 -0800 Subject: [PATCH 2/4] feat: Validate session on create and fetch --- src/codegen/cli/auth/decorators.py | 2 +- src/codegen/cli/auth/session.py | 80 +++++++++++++------ src/codegen/cli/commands/init/main.py | 46 +---------- src/codegen/cli/git/url.py | 11 --- src/codegen/cli/workspace/decorators.py | 7 -- .../cli/workspace/initialize_workspace.py | 34 ++------ .../git/repo_operator/local_git_repo.py | 5 +- 7 files changed, 69 insertions(+), 116 deletions(-) delete mode 100644 src/codegen/cli/git/url.py 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 415252056..ec149f7d3 100644 --- a/src/codegen/cli/auth/session.py +++ b/src/codegen/cli/auth/session.py @@ -1,6 +1,12 @@ from pathlib import Path +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 @@ -11,19 +17,28 @@ class CodegenSession: """Represents an authenticated codegen session with user and repository context""" 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) - global_config.set_active_session(repo_path) + self.config.secrets.github_token = git_token or self.config.secrets.github_token + global_config.set_active_session(repo_path) if not self.existing: - self._initialize_repo_config() + self.initialize() + + self.validate() @classmethod def from_active_session(cls) -> "CodegenSession | None": @@ -33,35 +48,50 @@ def from_active_session(cls) -> "CodegenSession | None": return cls(active_session) - @classmethod - def from_local_git(cls, git_repo: LocalGitRepo, token: str | None = None) -> "CodegenSession": - session = cls(git_repo.repo_path) - session._initialize_repo_config(git_repo, token) - return session - - def _initialize_repo_config(self, git_repo: LocalGitRepo, token: str | None = None): + def initialize(self) -> None: """Initialize the codegen session""" - self.config.repository.repo_path = str(git_repo.repo_path) - self.config.repository.repo_name = git_repo.name - self.config.repository.full_name = git_repo.full_name - self.config.repository.user_name = git_repo.user_name - self.config.repository.user_email = git_repo.user_email - self.config.repository.language = git_repo.get_language(access_token=token) + self.config.repository.repo_path = str(self.local_git.repo_path) + self.config.repository.repo_name = self.local_git.name + self.config.repository.full_name = self.local_git.full_name + self.config.repository.user_name = self.local_git.user_name + self.config.repository.user_email = self.local_git.user_email + self.config.repository.language = self.local_git.get_language(access_token=self.config.secrets.github_token).upper() self.config.save() - def is_valid(self) -> bool: - """Validates that the session configuration is correct""" - if not self.repo_path.exists(): - return False - + def validate(self) -> None: + """Validates that the session configuration is correct, otherwise raises an error""" if not self.codegen_dir.exists(): - return False + 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(): - return False + 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() - if get_git_repo(self.repo_path) is None: - return False + 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 7ad16914f..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 @@ -34,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.from_local_git(local_git) - 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/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: From ec8a1b77776d89a8e3be4e626a7de20c462d99e5 Mon Sep 17 00:00:00 2001 From: Carol Jung Date: Wed, 12 Feb 2025 17:48:34 -0800 Subject: [PATCH 3/4] always initialize --- src/codegen/cli/auth/session.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/codegen/cli/auth/session.py b/src/codegen/cli/auth/session.py index ec149f7d3..6ac2ff915 100644 --- a/src/codegen/cli/auth/session.py +++ b/src/codegen/cli/auth/session.py @@ -30,15 +30,13 @@ def __init__(self, repo_path: Path, git_token: str | None = None) -> None: 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 - - global_config.set_active_session(repo_path) - if not self.existing: - self.initialize() + self.existing = global_config.get_session(repo_path) is not None self.validate() + self.initialize() + global_config.set_active_session(repo_path) @classmethod def from_active_session(cls) -> "CodegenSession | None": From 2ca05256d4a29fc4a9f139a528c9891554159c01 Mon Sep 17 00:00:00 2001 From: Carol Jung Date: Wed, 12 Feb 2025 17:55:24 -0800 Subject: [PATCH 4/4] dont override existing config --- src/codegen/cli/auth/session.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/src/codegen/cli/auth/session.py b/src/codegen/cli/auth/session.py index 6ac2ff915..e34b4b982 100644 --- a/src/codegen/cli/auth/session.py +++ b/src/codegen/cli/auth/session.py @@ -34,8 +34,7 @@ def __init__(self, repo_path: Path, git_token: str | None = None) -> None: 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.validate() - self.initialize() + self._initialize() global_config.set_active_session(repo_path) @classmethod @@ -46,17 +45,19 @@ def from_active_session(cls) -> "CodegenSession | None": return cls(active_session) - def initialize(self) -> None: + def _initialize(self) -> None: """Initialize the codegen session""" - self.config.repository.repo_path = str(self.local_git.repo_path) - self.config.repository.repo_name = self.local_git.name - self.config.repository.full_name = self.local_git.full_name - self.config.repository.user_name = self.local_git.user_name - self.config.repository.user_email = self.local_git.user_email - self.config.repository.language = self.local_git.get_language(access_token=self.config.secrets.github_token).upper() + 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: + 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}")