repgen-ai 0.1.0__py3-none-any.whl

repgen/retrieval/config.py

@@ -0,0 +1,53 @@
+import os
+from pathlib import Path
+
+
+class Config:
+    def __init__(self, repo_path: str, bug_report_path: str, output_dir: str):
+
+        self.BASE_DIR = Path(os.path.dirname(os.path.abspath(__file__))).parent
+
+        self.CODE_DIR = Path(repo_path).resolve()
+        self.BUG_REPORT_FILE = Path(bug_report_path).resolve()
+        self.OUTPUT_DIR = Path(output_dir).resolve()
+
+        # Unique project identifier based on repo name and bug report name
+        self.PROJECT_ID = f"{self.CODE_DIR.name}_{self.BUG_REPORT_FILE.stem}"
+        self.PROJECT_DIR = self.OUTPUT_DIR / self.PROJECT_ID
+
+        # --- Base paths for outputs ---
+        self.REFINED_BUG_REPORT_DIR = self.PROJECT_DIR / "refined_bug_report"
+        self.CONTEXT_DIR = self.PROJECT_DIR / "context"
+        self.PLANS_DIR = self.PROJECT_DIR / "plan"
+        self.REPRODUCTION_DIR = self.PROJECT_DIR / "reproduction_code"
+
+        # Set IN/OUT paths to the same consolidated directories
+        self.REFINED_BUG_REPORT_DIR_IN = self.REFINED_BUG_REPORT_DIR
+        self.REFINED_BUG_REPORT_DIR_OUT = self.REFINED_BUG_REPORT_DIR
+        self.CONTEXT_DIR_IN = self.CONTEXT_DIR
+        self.CONTEXT_DIR_OUT = self.CONTEXT_DIR
+        self.PLANS_DIR_IN = self.PLANS_DIR
+        self.PLANS_DIR_OUT = self.PLANS_DIR
+        self.REPRODUCTION_DIR_OUT = self.REPRODUCTION_DIR
+
+        # Model configurations
+        self.EMBEDDING_MODEL = (
+            "flax-sentence-embeddings/st-codesearch-distilroberta-base"
+        )
+        self.RERANKER_MODEL = "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"
+
+        # Search parameters
+        self.SEARCH_TOP_K = 200
+        self.RERANK_TOP_K = 20
+        self.ALPHA = 0.55
+
+        # Ensure directories exist
+        self._setup_directories()
+
+    def _setup_directories(self):
+        # Ensure base directories exist
+        self.PROJECT_DIR.mkdir(parents=True, exist_ok=True)
+        self.REFINED_BUG_REPORT_DIR.mkdir(parents=True, exist_ok=True)
+        self.CONTEXT_DIR.mkdir(parents=True, exist_ok=True)
+        self.PLANS_DIR.mkdir(parents=True, exist_ok=True)
+        self.REPRODUCTION_DIR.mkdir(parents=True, exist_ok=True)

repgen/retrieval/core/code_indexer.py

@@ -0,0 +1,138 @@
+import ast
+import logging
+import time
+from multiprocessing import Pool, cpu_count
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from ..models.hybrid_search import HybridSearchIndex
+from .utils import tokenize
+
+logger = logging.getLogger(__name__)
+
+
+class CodeIndexer:
+    def __init__(self, config):
+        self.config = config
+
+    def _load_file(self, file_path: Path) -> Optional[Dict[str, Any]]:
+        """Load a single file and return as document."""
+        try:
+            with open(file_path, "r", encoding="utf-8") as f:
+                return {
+                    "page_content": f.read(),
+                    "metadata": {"source": str(file_path), "type": "file"},
+                }
+        except Exception as e:
+            logger.warning(f"Error loading {file_path}: {e}")
+            return None
+
+    def _semantic_chunking(self, doc: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """Split code into semantic chunks (functions, classes)."""
+        chunks = []
+        file_path = doc["metadata"]["source"]
+        code = doc["page_content"]
+
+        try:
+            tree = ast.parse(code)
+            current_chunk = []
+
+            for node in ast.walk(tree):
+                if isinstance(
+                    node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)
+                ):
+                    start_lineno = node.lineno - 1
+                    end_lineno = (
+                        node.end_lineno
+                        if hasattr(node, "end_lineno")
+                        else start_lineno + 1
+                    )
+                    chunk_code = "\n".join(code.splitlines()[start_lineno:end_lineno])
+                    current_chunk.append(chunk_code)
+
+                    if len(current_chunk) >= 3:
+                        chunk = {
+                            "page_content": "\n\n".join(current_chunk),
+                            "metadata": {
+                                "source": file_path,
+                                "type": "semantic_chunk",
+                                "start_line": start_lineno + 1,
+                                "end_line": end_lineno,
+                            },
+                        }
+                        chunks.append(chunk)
+                        current_chunk = []
+
+            if current_chunk:
+                chunk = {
+                    "page_content": "\n\n".join(current_chunk),
+                    "metadata": {
+                        "source": file_path,
+                        "type": "semantic_chunk",
+                        "start_line": start_lineno + 1,
+                        "end_line": end_lineno,
+                    },
+                }
+                chunks.append(chunk)
+
+        except SyntaxError as e:
+            logger.warning(f"Syntax error in {file_path}: {e}")
+
+        return chunks
+
+    def _process_codebase(
+        self, code_dir: Path
+    ) -> tuple[List[Dict[str, Any]], List[List[str]]]:
+        """Process all files in the code directory."""
+        files = [
+            f
+            for f in code_dir.glob("**/*.py")
+            if "__pycache__" not in str(f)
+            and not any(part.startswith(".") for part in f.parts)
+        ]
+
+        with Pool(cpu_count()) as pool:
+            docs = pool.map(self._load_file, files)
+
+        valid_docs = [doc for doc in docs if doc is not None]
+        chunks = []
+        for doc in valid_docs:
+            chunks.extend(self._semantic_chunking(doc))
+
+        corpus = [tokenize(chunk["page_content"]) for chunk in chunks]
+        return chunks, corpus
+
+    def index_codebase(self, code_dir: Path) -> HybridSearchIndex:
+        """Index the codebase for hybrid search."""
+        start_time = time.time()
+        index_dir = self.config.PROJECT_DIR / "hybrid_index"
+
+        hybrid_index = HybridSearchIndex(
+            embedding_model=self.config.EMBEDDING_MODEL,
+            reranker_model=self.config.RERANKER_MODEL,
+            config=self.config,
+        )
+
+        if (index_dir / "documents.json").exists():
+            logger.info("Loading existing index")
+            hybrid_index.load_index(index_dir)
+            return hybrid_index
+
+        logger.info("Building new index")
+        chunks, corpus = self._process_codebase(code_dir)
+        hybrid_index.build_index(chunks, corpus)
+        hybrid_index.save_index(index_dir)
+
+        logger.info(f"Indexing completed in {time.time()-start_time:.2f}s")
+        return hybrid_index
+
+    def find_relevant_code(
+        self, bug_report: str, hybrid_index: HybridSearchIndex
+    ) -> List[dict]:
+        """Find code relevant to the bug report."""
+        return hybrid_index.search(
+            bug_report,
+            top_k=self.config.SEARCH_TOP_K,
+            rerank_top_k=self.config.RERANK_TOP_K,
+            alpha=self.config.ALPHA,
+        )

repgen/retrieval/core/dependency_analyzer.py

@@ -0,0 +1,121 @@
+import ast
+import logging
+from pathlib import Path
+from typing import Any, Dict, List, Set, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+class DependencyAnalyzer:
+    def __init__(self, config):
+        self.config = config
+        self.code_dir = Path(config.CODE_DIR).resolve()
+        self.ignore_modules = {
+            "os",
+            "sys",
+            "re",
+            "math",
+            "tensorflow",
+            "pytorch",
+            "numpy",
+            "pandas",
+        }
+
+    def _make_relative(self, path: Path) -> str:
+        """Convert path to relative path from the code directory."""
+        try:
+            return str(path.relative_to(self.code_dir))
+        except ValueError:
+            return str(path)
+
+    def _get_function_calls(self, node: ast.AST) -> Set[Tuple[str, str]]:
+        """Extract function calls from AST nodes."""
+        calls = set()
+
+        if isinstance(node, ast.Call):
+            if isinstance(node.func, ast.Attribute):
+                # Handle method calls like x.y()
+                if isinstance(node.func.value, ast.Name):
+                    calls.add((node.func.value.id, node.func.attr))
+            elif isinstance(node.func, ast.Name):
+                # Handle direct function calls like x()
+                calls.add(("", node.func.id))
+
+        return calls
+
+    def analyze_file_dependencies(self, file_path: Path) -> List[Dict[str, str]]:
+        """Analyze a file and return its external dependencies."""
+        try:
+            with open(file_path, "r", encoding="utf-8") as f:
+                tree = ast.parse(f.read())
+        except Exception as e:
+            logger.debug(f"Could not parse {file_path}: {e}")
+            return []
+
+        imports = {}
+        dependencies = set()
+
+        # First pass: collect all imports
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Import):
+                for alias in node.names:
+                    module = alias.name.split(".")[0]
+                    if module not in self.ignore_modules:
+                        imports[alias.asname or alias.name.split(".")[0]] = module
+
+            elif isinstance(node, ast.ImportFrom):
+                if node.module and node.level == 0:  # Only absolute imports
+                    module = node.module.split(".")[0]
+                    if module not in self.ignore_modules:
+                        for alias in node.names:
+                            imports[alias.asname or alias.name] = (
+                                f"{module}.{alias.name}"
+                            )
+
+        # Second pass: find function calls
+        for node in ast.walk(tree):
+            for module, func in self._get_function_calls(node):
+                if module in imports:
+                    # Resolve the actual module from imports
+                    actual_module = imports[module].split(".")[0]
+                    dependencies.add((actual_module, func))
+                elif not module:  # Direct function call
+                    if func in imports:
+                        full_path = imports[func].split(".")
+                        if len(full_path) > 1:
+                            dependencies.add((full_path[0], full_path[1]))
+
+        # Convert to the desired output format
+        return [{"module": m, "function": f} for m, f in sorted(dependencies)]
+
+    def analyze_training_dependencies(
+        self, training_report: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Analyze dependencies for all training files."""
+        report = []
+
+        for entry in training_report["training_files"]:
+            try:
+                file_path = (self.code_dir / entry["file"]).resolve()
+                if not file_path.exists():
+                    continue
+
+                external_deps = self.analyze_file_dependencies(file_path)
+
+                report.append(
+                    {
+                        "file": self._make_relative(file_path),
+                        "score": entry["score"],
+                        "contains_training": entry["contains_training"],
+                        "external_dependencies": external_deps,
+                    }
+                )
+
+            except Exception as e:
+                logger.warning(f"Error processing {entry.get('file')}: {e}")
+                continue
+
+        return {
+            "bug_report": training_report["bug_report"],
+            "dependencies": sorted(report, key=lambda x: x["score"], reverse=True),
+        }

repgen/retrieval/core/module_analyzer.py

@@ -0,0 +1,65 @@
+import os
+from collections import defaultdict
+from pathlib import Path
+from typing import Any, Dict, List
+
+from .utils import extract_module_path
+
+
+class ModuleAnalyzer:
+    def __init__(self, config):
+        self.config = config
+        self.code_dir = Path(config.CODE_DIR)
+
+    def _make_relative(self, file_path: str) -> str:
+        """Convert absolute path to relative path from the code directory."""
+        try:
+            path = Path(file_path)
+            # Handle both str and Path inputs
+            if not isinstance(file_path, (str, Path)):
+                return file_path
+
+            # Make relative to the code directory
+            relative_path = os.path.relpath(str(path), start=str(self.code_dir))
+            return relative_path.replace("\\", "/")  # Normalize to forward slashes
+        except (TypeError, ValueError):
+            return file_path  # Return original if conversion fails
+
+    def analyze_modules(
+        self, relevant_code: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Group snippets by module and file using relative paths."""
+        module_map = defaultdict(lambda: defaultdict(list))
+
+        for snippet in relevant_code:
+            file_path = snippet["metadata"]["source"]
+            relative_path = self._make_relative(file_path)
+            module_path = extract_module_path(relative_path)
+
+            module_map[module_path][relative_path].append(
+                {
+                    "start_line": snippet["metadata"].get("start_line"),
+                    "end_line": snippet["metadata"].get("end_line"),
+                    "code": snippet["page_content"],
+                }
+            )
+
+        return [
+            {
+                "module": module,
+                "files": [
+                    {
+                        "path": file,
+                        "snippets": [
+                            {
+                                "lines": f"{s['start_line']}-{s['end_line']}",
+                                "code": s["code"],
+                            }
+                            for s in snippets
+                        ],
+                    }
+                    for file, snippets in files.items()
+                ],
+            }
+            for module, files in module_map.items()
+        ]

repgen/retrieval/core/training_code_detector.py

@@ -0,0 +1,240 @@
+import ast
+import logging
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Tuple
+
+import torch
+from sentence_transformers import CrossEncoder
+from transformers import AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+
+class TrainingLoopVisitor(ast.NodeVisitor):
+    """
+    AST visitor to detect training loops in:
+    - TensorFlow 1.x (Session.run in loops)
+    - TensorFlow 2.x (Keras APIs, GradientTape, and optimizer methods)
+    - PyTorch (optimizer steps in loops, DataLoader usage)
+    - PyTorch Lightning (training_step method)
+    """
+
+    def __init__(self):
+        # Framework flags
+        self.found_training = False
+
+        # Context tracking
+        self.loop_depth = 0  # Track nested loops
+        self.gradient_tape_depth = 0  # Track GradientTape contexts
+        self.current_class = None  # Track class context
+
+    # Context management ------------------------------------------------------
+    def visit_ClassDef(self, node):
+        """Track class context for LightningModule detection"""
+        self.current_class = node.name
+        self.generic_visit(node)
+        self.current_class = None
+
+    def visit_For(self, node):
+        self.loop_depth += 1
+        self.generic_visit(node)
+        self.loop_depth -= 1
+
+    def visit_While(self, node):
+        self.loop_depth += 1
+        self.generic_visit(node)
+        self.loop_depth -= 1
+
+    def visit_With(self, node):
+        """Handle GradientTape and PyTorch profiler contexts"""
+        # Check for GradientTape first
+        if any(self._is_gradient_tape(item.context_expr) for item in node.items):
+            self.gradient_tape_depth += 1
+            self.generic_visit(node)
+            self.gradient_tape_depth -= 1
+        else:
+            self.generic_visit(node)
+
+    # Core detection logic -----------------------------------------------------
+    def visit_Call(self, node):
+        """Analyze method calls for framework patterns"""
+        self._check_tf1_patterns(node)
+        self._check_tf2_patterns(node)
+        self._check_pytorch_patterns(node)
+        self.generic_visit(node)
+
+    def _check_tf1_patterns(self, node):
+        """TensorFlow 1.x: Session.run in loops with training ops"""
+        if self.loop_depth > 0:
+            if isinstance(node.func, ast.Attribute) and node.func.attr == "run":
+                # Basic check for Session-like objects
+                if isinstance(node.func.value, (ast.Name, ast.Attribute)):
+                    self.found_training = True
+
+    def _check_tf2_patterns(self, node):
+        """TensorFlow 2.x detection"""
+        if isinstance(node.func, ast.Attribute):
+            # High-level APIs
+            if node.func.attr in {"fit", "fit_generator", "train_on_batch"}:
+                self.found_training = True
+
+            # Optimizer methods
+            if node.func.attr in {"apply_gradients", "minimize"}:
+                self.found_training = True
+
+            # GradientTape context patterns
+            if self.gradient_tape_depth > 0:
+                if node.func.attr in {"gradient", "watch"}:
+                    self.found_training = True
+
+    def _check_pytorch_patterns(self, node):
+        """PyTorch detection: Training steps in loops"""
+        if self.loop_depth > 0 and isinstance(node.func, ast.Attribute):
+            # Core training methods
+            if node.func.attr in {"backward", "step", "zero_grad"}:
+                self.found_training = True
+
+            # Loss calculation patterns
+            if node.func.attr in {"item", "backward"} and self._is_loss_node(
+                node.func.value
+            ):
+                self.found_training = True
+
+            # DataLoader patterns
+            if node.func.attr in {"to", "cuda"} and self._is_dataloader_node(
+                node.func.value
+            ):
+                self.found_training = True
+
+    # Helper methods -----------------------------------------------------------
+    def _is_gradient_tape(self, node):
+        """Identify GradientTape usage (direct or via tf.GradientTape)"""
+        if isinstance(node, ast.Call):
+            func = node.func
+            return (
+                isinstance(func, ast.Attribute) and func.attr == "GradientTape"
+            ) or (isinstance(func, ast.Name) and func.id == "GradientTape")
+        return False
+
+    def _is_loss_node(self, node):
+        """Heuristic to identify loss nodes (e.g., 'loss' in names)"""
+        if isinstance(node, ast.Name):
+            return "loss" in node.id.lower()
+        elif isinstance(node, ast.Attribute):
+            return "loss" in node.attr.lower()
+        return False
+
+    def _is_dataloader_node(self, node):
+        """Heuristic to identify DataLoader instances"""
+        if isinstance(node, ast.Name):
+            return "loader" in node.id.lower() or "dataloader" in node.id.lower()
+        elif isinstance(node, ast.Attribute):
+            return "loader" in node.attr.lower() or "dataloader" in node.attr.lower()
+        return False
+
+    # PyTorch Lightning support -----------------------------------------------
+    def visit_FunctionDef(self, node):
+        """Check for LightningModule training_step hooks"""
+        if self.current_class and node.name == "training_step":
+            self.found_training = True
+        self.generic_visit(node)
+
+
+class TrainingCodeDetector:
+    def __init__(self, config):
+        self.config = config
+        self.reranker = CrossEncoder(config.RERANKER_MODEL)
+        self.tokenizer = AutoTokenizer.from_pretrained(config.RERANKER_MODEL)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.reranker.to(self.device)
+
+    def contains_training(self, file_path: Path) -> bool:
+        """Check if a file contains training-related code using AST analysis."""
+        try:
+            with open(file_path, "r", encoding="utf-8") as f:
+                code = f.read()
+            tree = ast.parse(code)
+        except Exception as e:
+            logger.error(f"Error parsing {file_path}: {str(e)}")
+            return False
+
+        visitor = TrainingLoopVisitor()
+        visitor.visit(tree)
+        return visitor.found_training
+
+    def _get_all_files(self, module_report: Dict[str, Any]) -> List[Path]:
+        """Get all Python files from modules in the report."""
+        # --- MODIFICATION: Ablation: No Module-Centric Partitioning ---
+        files = set()  # Use set to avoid duplicates
+
+        files = set()  # Use set to avoid duplicates
+
+        # Original: Walk all files in the retrieved modules' directories
+        for module in module_report["modules"]:
+            module_path = self.config.CODE_DIR / module["module"]
+            if not module_path.exists():
+                continue
+
+            for root, _, filenames in os.walk(module_path):
+                for filename in filenames:
+                    if filename.endswith(".py"):
+                        files.add(Path(root) / filename)
+
+        return list(files)
+
+    def _rank_files(
+        self, files: List[Path], bug_report: str
+    ) -> List[Tuple[Path, float]]:
+        """Rank files by relevance to the bug report."""
+        if not files:
+            return []
+
+        # Tokenize with proper truncation
+        features = self.tokenizer(
+            [bug_report[:100000]] * len(files),
+            [file.read_text(encoding="utf-8")[:100000] for file in files],
+            padding=True,
+            truncation="longest_first",
+            max_length=512,
+            return_tensors="pt",
+        ).to(self.device)
+
+        # Run cross-encoder
+        with torch.no_grad():
+            scores = self.reranker(**features).logits.squeeze().cpu().numpy()
+
+        # Normalize scores
+        if len(scores) > 1:
+            scores = (scores - scores.min()) / (scores.max() - scores.min() + 1e-8)
+
+        return sorted(zip(files, scores), key=lambda x: x[1], reverse=True)
+
+    def detect_training_code(
+        self, module_report: Dict[str, Any], bug_report_path: Path
+    ) -> Dict[str, Any]:
+        """Detect and rank training-related code."""
+        bug_report = bug_report_path.read_text(encoding="utf-8")
+        all_files = self._get_all_files(module_report)
+
+        training_files = []
+        training_files = []
+        for file in all_files:
+            if self.contains_training(file):
+                training_files.append(file)
+
+        # Original ranking logic
+        ranked_files = self._rank_files(training_files, bug_report)
+        # --- END MODIFICATION ---
+
+        return {
+            "bug_report": module_report["bug_report"],
+            "training_files": [
+                {
+                    "file": str(file.relative_to(self.config.CODE_DIR)),
+                    "score": float(score),
+                    "contains_training": True,
+                }
+                for file, score in ranked_files
+            ],
+        }

repgen/retrieval/core/utils.py

@@ -0,0 +1,52 @@
+import json
+import logging
+import os
+import re
+from pathlib import Path
+from typing import Any, List
+
+
+def extract_module_path(file_path: str, dataset_root: str = "dataset/") -> str:
+    """Extract full module path relative to dataset root."""
+    try:
+        file_path = os.path.normpath(file_path)
+        dataset_root = os.path.normpath(dataset_root)
+        parts = file_path.split(os.sep)
+
+        try:
+            root_idx = parts.index(dataset_root)
+        except ValueError:
+            return os.path.dirname(file_path)
+
+        module_parts = parts[root_idx + 1 : -1]  # Exclude filename
+        return os.path.join(*module_parts) if module_parts else "root"
+    except Exception as e:
+        logging.warning(f"Path parsing error {file_path}: {e}")
+        return "unknown"
+
+
+def load_bug_report(bug_report_path: Path) -> str:
+    """Load bug report content."""
+    try:
+        with open(bug_report_path, "r", encoding="utf-8") as f:
+            return f.read()
+    except Exception as e:
+        raise ValueError(f"Error loading bug report: {str(e)}")
+
+
+def save_json(data: Any, file_path: Path) -> None:
+    """Save data to JSON file."""
+    file_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(file_path, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2)
+
+
+def load_json(file_path: Path) -> Any:
+    """Load data from JSON file."""
+    with open(file_path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def tokenize(text: str) -> List[str]:
+    """Basic tokenizer for code search."""
+    return re.findall(r"\b\w+[\w\d_]*\b", text.lower())