repgen-ai 0.1.0__py3-none-any.whl

repgen/__init__.py

@@ -0,0 +1,51 @@
+from typing import Optional
+
+from .core import RepGenService
+
+
+def reproduce(
+    bug_report_source: str,
+    repo_source: str,
+    backend: str = "openai",
+    model: str = "gpt-4o",
+    api_key: Optional[str] = None,
+    commit: Optional[str] = None,
+    output_dir: str = "./repgen_results",
+) -> str:
+    """
+    Generates a reproduction script for a given bug report and repository.
+
+    Args:
+        bug_report_source: URL or path to the bug report.
+        repo_source: URL or path to the repository.
+        backend: The LLM backend to use (default: "openai").
+        model: The model name to use (default: "gpt-4o").
+        api_key: API key for the backend (optional).
+        commit: Specific commit hash to checkout (optional).
+        output_dir: Directory to store intermediate artifacts (default: "./repgen_results").
+
+    Returns:
+        The generated reproduction script as a string.
+
+    Raises:
+        RuntimeError: If reproduction fails.
+    """
+    service = RepGenService(output_dir=output_dir)
+    result = service.run_reproduction(
+        bug_report_source=bug_report_source,
+        repo_source=repo_source,
+        backend=backend,
+        model=model,
+        commit=commit,
+        api_key=api_key,
+    )
+
+    if result["success"] and result["files"]:
+        # Return the content of the first generated file (usually the reproduction script)
+        return result["files"][0]["content"]
+    else:
+        error_msg = result.get("error", "Unknown error during reproduction")
+        raise RuntimeError(f"Reproduction failed: {error_msg}")
+
+
+__all__ = ["reproduce", "RepGenService"]

repgen/cli.py

@@ -0,0 +1,375 @@
+import argparse
+import os
+import subprocess
+import sys
+import time
+
+# Rich Imports
+from rich.panel import Panel
+from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
+from rich.prompt import Confirm, IntPrompt, Prompt
+from rich.table import Table
+from rich_argparse import RichHelpFormatter
+
+from .core import RepGenService
+from .utils import console
+
+# Import from utils
+
+
+# ==========================================
+# INTERACTIVE UI HELPERS
+# ==========================================
+
+
+def check_ollama_status():
+    try:
+        subprocess.run(["ollama", "--version"], capture_output=True, check=True)
+        return True
+    except Exception:
+        return False
+
+
+def ensure_ollama_model(model: str):
+    try:
+        result = subprocess.run(["ollama", "list"], capture_output=True, text=True)
+        if model not in result.stdout:
+            console.print(f"[yellow]Model {model} not found. Pulling...[/yellow]")
+            subprocess.run(["ollama", "pull", model], check=True)
+            console.print(f"[green]Successfully pulled {model}[/green]")
+    except Exception as e:
+        console.print(f"[red]Error checking/pulling model: {e}[/red]")
+
+
+def get_interactive_config():
+    """Interactive prompt to gather missing configurations via Rich."""
+    console.print(
+        Panel.fit(
+            "[bold cyan]RepGen Interactive Setup[/bold cyan]", border_style="cyan"
+        )
+    )
+
+    # 1. Backend Selection
+    backends = ["ollama", "openai", "gemini", "claude"]
+    console.print("\n[bold]Select LLM Backend:[/bold]")
+    for i, b in enumerate(backends, 1):
+        console.print(f" {i}) [green]{b.capitalize()}[/green]")
+
+    choice = IntPrompt.ask(
+        "Choice",
+        choices=[str(i) for i in range(1, len(backends) + 1)],
+        show_choices=False,
+    )
+    backend = backends[choice - 1]
+
+    # 2. Environment & Key Setup
+    model = ""
+    if backend == "ollama":
+        with console.status("[bold green]Checking Ollama status...[/bold green]"):
+            is_installed = check_ollama_status()
+
+        if not is_installed:
+            console.print(
+                "[bold red]Error: Ollama is not installed or not running.[/bold red]"
+            )
+            if Confirm.ask("Would you like me to try installing it via Homebrew?"):
+                try:
+                    subprocess.run(["brew", "install", "ollama"], check=True)
+                    console.print(
+                        "[green]Ollama installed. Please start it logic and try again.[/green]"
+                    )
+                except Exception:
+                    console.print(
+                        "[red]Installation failed. Please install manually.[/red]"
+                    )
+                sys.exit(1)
+            else:
+                sys.exit(1)
+
+        model = Prompt.ask("Ollama Model", default="qwen2.5-coder:7b")
+        with console.status(
+            f"[bold green]Ensuring model {model} is available...[/bold green]"
+        ):
+            ensure_ollama_model(model)
+
+    else:
+        env_keys = {
+            "openai": "OPENAI_API_KEY",
+            "gemini": "GEMINI_API_KEY",
+            "claude": "ANTHROPIC_API_KEY",
+        }
+        env_var = env_keys[backend]
+
+        if not os.getenv(env_var):
+            console.print(f"\n[yellow]{env_var} not found.[/yellow]")
+            key = Prompt.ask(
+                f"Enter your {backend.capitalize()} API Key", password=True
+            )
+            if key:
+                os.environ[env_var] = key
+            else:
+                console.print("[red]API Key required. Terminating.[/red]")
+                sys.exit(1)
+
+        # Model Selection Suggestions
+        suggestions = {
+            "openai": ["gpt-4o", "gpt-4-turbo", "gpt-3.5-turbo"],
+            "gemini": ["gemini-1.5-pro", "gemini-1.5-flash"],
+            "claude": ["claude-3-5-sonnet-20240620", "claude-3-opus-20240229"],
+        }
+
+        model = Prompt.ask(
+            f"Select {backend.capitalize()} Model",
+            choices=suggestions[backend],
+            default=suggestions[backend][0],
+        )
+
+    return backend, model
+
+
+# ==========================================
+# MAIN EXECUTION
+# ==========================================
+
+
+def main():
+    # Silence internal loggers to keep the UI clean (Agent-style)
+    import logging
+
+    logging.getLogger("repgen").setLevel(logging.ERROR)
+
+    # Styled Banner
+    console.print(
+        Panel(
+            "[bold white]RepGen CLI: Automated Deep Learning Bug Reproduction[/bold white]\n"
+            "[italic grey70]Generate reproduction scripts from bug reports automatically.[/italic grey70]",
+            border_style="bold blue",
+            title="[bold cyan]RepGen[/bold cyan]",
+            subtitle="[dim]v1.0.0[/dim]",
+            expand=False,
+            padding=(1, 2),
+        )
+    )
+
+    parser = argparse.ArgumentParser(
+        description="RepGen CLI: Automated Bug Reproduction",
+        formatter_class=RichHelpFormatter,
+    )
+    parser.add_argument("--bug-report", help="Path or URL to the bug report")
+    parser.add_argument("--repo-path", help="Path or URL to the repository")
+    parser.add_argument("--commit", help="Specific git commit ID to checkout")
+    parser.add_argument(
+        "--output-dir", default="./repgen_results", help="Directory for results"
+    )
+    parser.add_argument(
+        "--backend",
+        choices=["ollama", "openai", "gemini", "claude"],
+        help="LLM backend",
+    )
+    parser.add_argument("--model", help="Model name")
+    parser.add_argument(
+        "--max-attempts", type=int, default=3, help="Max attempts per context"
+    )
+
+    args = parser.parse_args()
+
+    try:
+        # Interactive setup if arguments are missing
+        if not args.bug_report:
+            args.bug_report = Prompt.ask(
+                "[bold cyan]Enter URI to Bug Report file[/bold cyan] (Path or URL)"
+            )
+        if not args.repo_path:
+            args.repo_path = Prompt.ask(
+                "[bold cyan]Enter URI to Repository[/bold cyan] (Path or URL)"
+            )
+
+        if not args.backend or not args.model:
+            args.backend, args.model = get_interactive_config()
+
+        service = RepGenService(output_dir=args.output_dir)
+
+        console.print("\n[bold]Starting Reproduction Pipeline...[/bold]")
+        start_time = time.time()
+
+        # Define UI Helper for Agent Panels
+        def print_agent_panel(
+            content: str, agent_name: str, color: str, title_desc: str = ""
+        ):
+            from rich.markdown import Markdown
+            from rich.panel import Panel
+
+            width = min(120, console.size.width - 4)
+
+            md = Markdown(content)
+            title = f"{agent_name} ({title_desc})" if title_desc else agent_name
+
+            panel = Panel(
+                md,
+                title=f"[bold {color}]{title}[/bold {color}]",
+                title_align="left",
+                border_style=color,
+                width=width,
+                padding=(0, 1),
+            )
+            console.print(panel)
+
+        with Progress(
+            SpinnerColumn("dots", style="cyan"),
+            TextColumn("[bold blue]{task.description}"),
+            TimeElapsedColumn(),
+            console=console,
+            transient=False,
+        ) as progress:
+            task = progress.add_task("Initializing...", total=None)
+
+            def progress_cb(stage, msg, data=None):
+                # Update description based on stage
+                stage_colors = {
+                    "setup": "cyan",
+                    "retrieval": "magenta",
+                    "refinement": "yellow",
+                    "generation": "green",
+                }
+                color = stage_colors.get(stage, "white")
+                progress.update(
+                    task, description=f"[{color}][{stage.upper()}][/] {msg}"
+                )
+
+                # Display intermediate artifacts
+                if data:
+                    if data.get("type") == "refined_report":
+                        print_agent_panel(
+                            data["content"],
+                            "Refinement Agent",
+                            "yellow",
+                            "Bug Report processed",
+                        )
+                    elif data.get("type") == "plan":
+                        import json
+
+                        try:
+                            plan_obj = json.loads(data["content"])
+                            # Format list as a bulleted list for better readability
+                            plan_text = ""
+                            for step in plan_obj:
+                                plan_text += f"- {step}\n"
+                            print_agent_panel(
+                                plan_text,
+                                "Code Generation Agent",
+                                "green",
+                                f"Plan for Context {data.get('index')}",
+                            )
+                        except Exception:
+                            print_agent_panel(
+                                data["content"],
+                                "Code Generation Agent",
+                                "green",
+                                f"Plan for Context {data.get('index')}",
+                            )
+                    elif data.get("type") == "context":
+                        # Summarize context
+                        summary = f"**Loaded Context {data.get('index')}**\n\n"
+                        summary += f"Contains {len(data['content'].splitlines())} lines of code context."
+                        print_agent_panel(
+                            summary,
+                            "Context Retrieval Agent",
+                            "blue",
+                            f"Context {data.get('index')}",
+                        )
+                    elif data.get("type") == "feedback":
+                        print_agent_panel(
+                            f"**Feedback received:** {data['content']}",
+                            "Self-Correction",
+                            "red",
+                            "Error detected",
+                        )
+
+            result = service.run_reproduction(
+                bug_report_source=args.bug_report,
+                repo_source=args.repo_path,
+                backend=args.backend,
+                model=args.model,
+                commit=args.commit,
+                max_attempts=args.max_attempts,
+                progress_callback=progress_cb,
+            )
+
+        end_time = time.time()
+        duration = end_time - start_time
+
+        console.print(f"\n[bold]Total Duration:[/bold] {duration:.2f} seconds\n")
+
+        if result["success"]:
+            console.print(
+                Panel(
+                    "[bold green]Reproduction Success![/bold green]",
+                    title="System",
+                    title_align="left",
+                    border_style="green",
+                    width=min(120, console.size.width - 4),
+                )
+            )
+
+            # Display Results Table
+            table = Table(
+                title="Generated Files",
+                show_header=True,
+                header_style="bold magenta",
+                width=min(120, console.size.width - 4),
+            )
+            table.add_column("File Name", style="cyan")
+            table.add_column("Path", style="dim")
+            table.add_column("Size", justify="right")
+
+            for file_info in result["files"]:
+                path = file_info["path"]
+                content = file_info["content"]
+                name = os.path.basename(path)
+                size = f"{len(content.encode('utf-8'))} bytes"
+                table.add_row(name, path, size)
+
+            console.print(table)
+
+            # Display File Contents
+            from rich.syntax import Syntax
+
+            for file_info in result["files"]:
+                syntax = Syntax(
+                    file_info["content"], "python", theme="monokai", line_numbers=True
+                )
+
+                panel = Panel(
+                    syntax,
+                    title=f"[bold cyan]Reproducer Code ({os.path.basename(file_info['path'])})[/bold cyan]",
+                    title_align="left",
+                    border_style="cyan",
+                    width=min(120, console.size.width - 4),
+                )
+                console.print(panel)
+
+            console.print("\n[bold green]Next Steps:[/bold green]")
+            for file_info in result["files"]:
+                console.print(f"  Run: [code]python {file_info['path']}[/code]")
+
+        else:
+            console.print(
+                Panel(
+                    f"Error: {result['error']}",
+                    title="[bold red]Reproduction Failed[/bold red]",
+                    title_align="left",
+                    border_style="red",
+                    width=min(120, console.size.width - 4),
+                )
+            )
+
+    except KeyboardInterrupt:
+        console.print("\n[bold red]⚠ Operation cancelled by user.[/bold red]")
+        sys.exit(130)
+    except Exception as e:
+        console.print(f"\n[bold red]Unexpected Error:[/bold red] {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

repgen/core.py

@@ -0,0 +1,239 @@
+import json
+import logging
+import time
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from .retrieval.pipeline import RetrievalPipeline
+from .utils import (
+    _build_prompt,
+    check_relevance,
+    check_structural_correctness,
+    checkout_commit,
+    clean_context,
+    create_prompt_plan,
+    create_prompt_refinement,
+    extract_json_content,
+    fetch_content,
+    prepare_repository,
+    query_llm,
+)
+
+logger = logging.getLogger("repgen.core")
+
+
+class RepGenService:
+    def __init__(self, output_dir: str = "./repgen_results"):
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+    def run_reproduction(
+        self,
+        bug_report_source: str,
+        repo_source: str,
+        backend: str,
+        model: str,
+        commit: Optional[str] = None,
+        max_attempts: int = 3,
+        api_key: Optional[str] = None,
+        progress_callback=None,
+    ) -> Dict[str, Any]:
+        """
+        Executes the full reproduction pipeline.
+
+        Args:
+            bug_report_source: URL or path to bug report
+            repo_source: URL or path to repository
+            backend: LLM backend (ollama, openai, etc)
+            model: Model name
+            commit: Optional commit hash
+            max_attempts: Max retry attempts per context
+            api_key: Optional API key for backend
+            progress_callback: Optional function(stage, message)
+
+        Returns:
+            Dict containing results and paths
+        """
+        result = {"success": False, "project_id": "", "files": [], "error": None}
+
+        def notify(stage, msg, data=None):
+            if progress_callback:
+                progress_callback(stage, msg, data)
+            logger.info(f"[{stage}] {msg}")
+
+        try:
+            # 1. Prepare Inputs
+            notify("setup", "Fetching inputs...")
+            bug_report_content = fetch_content(bug_report_source)
+            local_repo_path = prepare_repository(repo_source)
+
+            if commit:
+                notify("setup", f"Checking out commit {commit}...")
+                if not checkout_commit(local_repo_path, commit):
+                    raise RuntimeError("Failed to checkout commit")
+
+            # 2. Initialize Pipeline
+            # Create a unique project ID/folder
+            project_id = f"RepGen_Run_{int(time.time())}"
+            run_dir = self.output_dir / project_id
+            run_dir.mkdir(parents=True, exist_ok=True)
+
+            local_bug_report_path = run_dir / "input_bug_report.txt"
+            local_bug_report_path.write_text(bug_report_content)
+
+            notify("retrieval", "Initializing Retrieval Pipeline...")
+            pipeline = RetrievalPipeline(
+                repo_path=local_repo_path,
+                bug_report_path=str(local_bug_report_path),
+                output_dir=str(self.output_dir),
+            )
+
+            # 3. Retrieval
+            notify("retrieval", "Running Code Retrieval...")
+            ret_result = pipeline.run_pipeline()
+            if not ret_result:
+                raise RuntimeError("Retrieval failed to find relevant code")
+
+            # 4. Refinement
+            notify("refinement", "Refining Bug Report...")
+            refined_report = query_llm(
+                create_prompt_refinement(bug_report_content),
+                backend,
+                model,
+                api_key=api_key,
+            )
+            if not refined_report:
+                refined_report = bug_report_content
+
+            refined_path = (
+                pipeline.config.REFINED_BUG_REPORT_DIR_OUT / "refined_report.txt"
+            )
+            refined_path.write_text(refined_report)
+            notify(
+                "refinement",
+                "Bug Report Refined",
+                data={"type": "refined_report", "content": refined_report},
+            )
+
+            # 5. Generation
+            context_files = sorted(list(pipeline.config.CONTEXT_DIR_OUT.glob("*.json")))
+            if not context_files:
+                raise RuntimeError("No context files generated")
+
+            notify(
+                "generation",
+                f"Starting generation for {len(context_files)} contexts...",
+            )
+
+            successful_gen = False
+            generated_script_path = None
+
+            for idx, ctx_file in enumerate(context_files, 1):
+                ctx_num = ctx_file.stem.split("_")[-1]
+                notify("generation", f"Processing Context {idx}/{len(context_files)}")
+
+                with open(ctx_file, "r") as f:
+                    ctx_content = f.read()
+
+                notify(
+                    "generation",
+                    f"Context {ctx_num} Loaded",
+                    data={
+                        "type": "context",
+                        "content": clean_context(ctx_content),
+                        "index": idx,
+                    },
+                )
+
+                # Plan
+                plan_raw = query_llm(
+                    create_prompt_plan(refined_report, json.loads(ctx_content)),
+                    backend,
+                    model,
+                    api_key=api_key,
+                )
+                plan = extract_json_content(plan_raw) or "[]"
+                notify(
+                    "generation",
+                    f"Plan Generated for Context {ctx_num}",
+                    data={"type": "plan", "content": plan, "index": idx},
+                )
+
+                prompt_code = _build_prompt(refined_report, ctx_content, plan)
+
+                # Attempts
+                for attempt in range(1, max_attempts + 1):
+                    notify("generation", f"Context {ctx_num} - Attempt {attempt}")
+
+                    raw_code = query_llm(prompt_code, backend, model, api_key=api_key)
+                    code = extract_json_content(raw_code) or raw_code
+
+                    is_valid, err = check_structural_correctness(code)
+                    if err:
+                        notify(
+                            "generation",
+                            f"Context {ctx_num} - Attempt {attempt} - Syntax Error",
+                            data={"type": "feedback", "content": err, "code": code},
+                        )
+                        prompt_code = _build_prompt(
+                            refined_report,
+                            ctx_content,
+                            plan,
+                            feedback=f"Syntax Error: {err}",
+                        )
+                        continue
+
+                    if not check_relevance(
+                        refined_report, code, backend, model, api_key=api_key
+                    ):
+                        notify(
+                            "generation",
+                            f"Context {ctx_num} - Attempt {attempt} - Irrelevant",
+                            data={
+                                "type": "feedback",
+                                "content": "Code deemed irrelevant",
+                                "code": code,
+                            },
+                        )
+                        continue
+
+                    # Success
+                    generated_script_path = (
+                        pipeline.config.REPRODUCTION_DIR_OUT / f"reproduce_{ctx_num}.py"
+                    )
+                    generated_script_path.write_text(code)
+
+                    successful_gen = True
+                    result["success"] = True
+
+                    result["files"].append(
+                        {"path": str(generated_script_path), "content": code}
+                    )
+                    notify(
+                        "generation",
+                        "Reproduction Successful",
+                        data={
+                            "type": "code",
+                            "content": code,
+                            "path": str(generated_script_path),
+                        },
+                    )
+                    break
+
+                if successful_gen:
+                    break
+
+            if not successful_gen:
+                result["error"] = "Failed to generate valid script after all attempts"
+
+            result["project_id"] = pipeline.config.PROJECT_ID
+            return result
+
+        except RuntimeError as e:
+            logger.error(f"Pipeline flow error: {e}")
+            result["error"] = str(e)
+            return result
+        except Exception as e:
+            logger.exception("Unexpected reproduction failure")
+            result["error"] = f"Unexpected error: {str(e)}"
+            return result

repgen/retrieval/__init__.py

@@ -0,0 +1,4 @@
+from .config import Config
+from .pipeline import RetrievalPipeline
+
+__all__ = ["RetrievalPipeline", "Config"]