From 08acc44395e98297ca89ddbfb2fdf6a347dd9786 Mon Sep 17 00:00:00 2001 From: shrutu0929 Date: Thu, 16 Apr 2026 20:53:19 +0530 Subject: [PATCH 1/2] feat: complete Refactron optimization and LLM integration - Implemented high-performance TaintAnalyzer with sink/source indexing and memoization - Resolved symbol table cache invalidation and incremental analysis bugs - Integrated optimized AI refactoring with batching and RAG keyword fallback - Stabilized CLI for Windows and experimental Python 3.14 environments --- bad_code.py | 19 +++ refactron/analysis/symbol_table.py | 159 ++++++++++++++---- refactron/analysis/taint.py | 103 ++++++++---- refactron/cli/analysis.py | 12 +- refactron/cli/main.py | 11 +- refactron/cli/rag.py | 17 +- refactron/cli/refactor.py | 141 +++++++++++++++- refactron/cli/ui.py | 79 +++++---- refactron/cli/utils.py | 2 +- refactron/core/config.py | 4 + refactron/core/inference.py | 67 +++++++- refactron/core/parallel.py | 63 +++++-- refactron/core/refactron.py | 76 +++++++-- refactron/core/workspace.py | 16 +- refactron/llm/backend_client.py | 2 +- refactron/llm/client.py | 2 +- refactron/llm/orchestrator.py | 219 ++++++++++++++++++++++++- refactron/llm/prompts.py | 111 ++++++++----- refactron/rag/indexer.py | 173 +++++++++++-------- refactron/rag/retriever.py | 153 +++++++++++++---- tests/test_config_management.py | 35 ++-- tests/test_performance_optimization.py | 5 +- tests/test_symbol_table_incremental.py | 122 ++++++++++++++ 23 files changed, 1272 insertions(+), 319 deletions(-) create mode 100644 bad_code.py create mode 100644 tests/test_symbol_table_incremental.py diff --git a/bad_code.py b/bad_code.py new file mode 100644 index 0000000..a946b0b --- /dev/null +++ b/bad_code.py @@ -0,0 +1,19 @@ +THRESHOLD_VALUE = 5 +MIN_X_VALUE = 5 +MAX_Y_VALUE = 10 +ITERATION_LIMIT = 100 +MIN_ITERATION_VALUE = 10 +MAX_ITERATION_VALUE = 5 +def do_something_crazy(x: int, y: int) -> int: + """ + This function performs a series of operations based on the input values x and y. + It checks if x is greater than the threshold value and y is less than the max y value. + If the conditions are met, it iterates over a range of numbers and prints a message. + Finally, it returns the sum of x and y. + """ + if x > THRESHOLD_VALUE: + if y < MAX_Y_VALUE: + for i in range(ITERATION_LIMIT): + print("doing something", x) + return x + y +do_something_crazy(10, 5) \ No newline at end of file diff --git a/refactron/analysis/symbol_table.py b/refactron/analysis/symbol_table.py index de9bbf1..dcb3b5d 100644 --- a/refactron/analysis/symbol_table.py +++ b/refactron/analysis/symbol_table.py @@ -3,12 +3,13 @@ Maps classes, functions, variables, and their relationships across the codebase. """ +import hashlib import json import logging from dataclasses import dataclass, field from enum import Enum from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Set from refactron.core.inference import InferenceEngine @@ -61,16 +62,27 @@ class SymbolTable: symbols: Dict[str, Dict[str, Dict[str, Symbol]]] = field(default_factory=dict) # Map: global_name -> Symbol (for easy cross-file lookup of exports) exports: Dict[str, Symbol] = field(default_factory=dict) + # Map: file_path -> { "mtime": float, "size": int, "sha256": str } + file_metadata: Dict[str, Dict[str, Any]] = field(default_factory=dict) + + @staticmethod + def _normalize_path(path: str) -> str: + """Standardize path format for consistent keys/storage.""" + return Path(path).resolve().as_posix() def add_symbol(self, symbol: Symbol) -> None: """Add a symbol to the table.""" - if symbol.file_path not in self.symbols: - self.symbols[symbol.file_path] = {} + path = self._normalize_path(symbol.file_path) + # Ensure the symbol itself stores the normalized path + symbol.file_path = path + + if path not in self.symbols: + self.symbols[path] = {} - if symbol.scope not in self.symbols[symbol.file_path]: - self.symbols[symbol.file_path][symbol.scope] = {} + if symbol.scope not in self.symbols[path]: + self.symbols[path][symbol.scope] = {} - self.symbols[symbol.file_path][symbol.scope][symbol.name] = symbol + self.symbols[path][symbol.scope][symbol.name] = symbol # Track global exports (top-level functions and classes) if symbol.scope == "global" and symbol.type in ( @@ -78,13 +90,31 @@ def add_symbol(self, symbol: Symbol) -> None: SymbolType.FUNCTION, SymbolType.VARIABLE, ): - # Key by module path + name? Or just name for now? - # Using simple name collision strategy for MVP self.exports[symbol.name] = symbol + def remove_file(self, file_path: str) -> None: + """Remove all symbols and metadata associated with a file.""" + norm_path = self._normalize_path(file_path) + + if norm_path in self.symbols: + del self.symbols[norm_path] + + # Remove from exports + names_to_remove = [ + name + for name, sym in self.exports.items() + if self._normalize_path(sym.file_path) == norm_path + ] + for name in names_to_remove: + self.exports.pop(name, None) + + if norm_path in self.file_metadata: + del self.file_metadata[norm_path] + def get_symbol(self, file_path: str, name: str, scope: str = "global") -> Optional[Symbol]: """Retrieve a symbol.""" - return self.symbols.get(file_path, {}).get(scope, {}).get(name) + norm_path = self._normalize_path(file_path) + return self.symbols.get(norm_path, {}).get(scope, {}).get(name) def resolve_reference( self, name: str, current_file: str, current_scope: str @@ -106,8 +136,7 @@ def resolve_reference( if file_global: return file_global - # 3. Cross-file exports (Naive implementation) - # TODO: Enhance this with proper import resolution + # 3. Cross-file exports return self.exports.get(name) @@ -120,33 +149,85 @@ def __init__(self, cache_dir: Optional[Path] = None): self.inference_engine = InferenceEngine() def build_for_project(self, project_root: Path) -> SymbolTable: - """Scan project and build symbol table.""" + """Scan project and build symbol table incrementally.""" if self.cache_dir: - cached = self._load_cache() - if cached: - # TODO: Implement incremental update logic here - return cached + cached_table = self._load_cache() + if cached_table: + self.symbol_table = cached_table python_files = list(project_root.rglob("*.py")) + current_file_paths = {fp.resolve().as_posix() for fp in python_files} + + # 1. Remove deleted files + cached_files = list(self.symbol_table.file_metadata.keys()) + for cached_path in cached_files: + if cached_path not in current_file_paths: + logger.debug(f"Removing deleted file from symbol table: {cached_path}") + self.symbol_table.remove_file(cached_path) + + # 2. Analyze new or modified files for file_path in python_files: - self._analyze_file(file_path) + abs_path = file_path.resolve() + path_str = abs_path.as_posix() + if self._has_file_changed(abs_path, path_str): + logger.debug(f"Analyzing changed file: {path_str}") + self.symbol_table.remove_file(path_str) + self._analyze_file(abs_path) + self._update_file_metadata(abs_path, path_str) if self.cache_dir: self._save_cache() return self.symbol_table - def _analyze_file(self, file_path: Path) -> None: - """Analyze a single file and populate symbols.""" + def _has_file_changed(self, file_path: Path, file_path_str: str) -> bool: + """Check if file has changed since last analysis.""" + if file_path_str not in self.symbol_table.file_metadata: + return True + + metadata = self.symbol_table.file_metadata[file_path_str] + try: + stat = file_path.stat() + if stat.st_size != metadata.get("size"): + return True + + # Authoritative check: compare SHA-256 hashes + stored_hash = metadata.get("sha256") + if stored_hash: + current_hash = self._calculate_hash(file_path) + return current_hash != stored_hash + + return stat.st_mtime != metadata.get("mtime") + except Exception: + return True + + def _calculate_hash(self, file_path: Path) -> str: + """Calculate SHA-256 hash of file content.""" try: - # We use astroid for better inference capabilities later - tree = self.inference_engine.parse_file(str(file_path)) + return hashlib.sha256(file_path.read_bytes()).hexdigest() + except Exception: + return "" - # Walk the tree - self._visit_node(tree, str(file_path), "global") + def _update_file_metadata(self, file_path: Path, path_str: str) -> None: + """Update file metadata in symbol table.""" + try: + stat = file_path.stat() + self.symbol_table.file_metadata[path_str] = { + "mtime": stat.st_mtime, + "size": stat.st_size, + "sha256": self._calculate_hash(file_path), + } + except Exception as e: + logger.warning(f"Failed to update metadata for {path_str}: {e}") + def _analyze_file(self, file_path: Path) -> None: + """Analyze a single file and populate symbols.""" + path_str = file_path.resolve().as_posix() + try: + tree = self.inference_engine.parse_file(path_str) + self._visit_node(tree, path_str, "global") except Exception as e: - logger.warning(f"Failed to build symbol table for {file_path}: {e}") + logger.warning(f"Failed to build symbol table for {path_str}: {e}") def _visit_node(self, node: Any, file_path: str, scope: str) -> None: """Recursive node visitor.""" @@ -154,8 +235,8 @@ def _visit_node(self, node: Any, file_path: str, scope: str) -> None: new_scope = scope - if isinstance(node, (nodes.ClassDef, nodes.FunctionDef)): - # Register the definition itself in the CURRENT scope + # Recognize both FunctionDef and AsyncFunctionDef + if isinstance(node, (nodes.ClassDef, nodes.FunctionDef, nodes.AsyncFunctionDef)): symbol_type = ( SymbolType.CLASS if isinstance(node, nodes.ClassDef) else SymbolType.FUNCTION ) @@ -192,9 +273,8 @@ def _visit_node(self, node: Any, file_path: str, scope: str) -> None: self.symbol_table.add_symbol(symbol) # Recurse children - if hasattr(node, "get_children"): - for child in node.get_children(): - self._visit_node(child, file_path, new_scope) + for child in node.get_children(): + self._visit_node(child, file_path, new_scope) def _save_cache(self) -> None: """Save symbol table to cache.""" @@ -214,6 +294,7 @@ def _save_cache(self) -> None: for f, scopes in self.symbol_table.symbols.items() }, "exports": {n: sym.to_dict() for n, sym in self.symbol_table.exports.items()}, + "file_metadata": self.symbol_table.file_metadata, } with open(cache_file, "w") as f: @@ -238,15 +319,27 @@ def _load_cache(self) -> Optional[SymbolTable]: # Reconstruct symbols for f_path, scopes in data.get("symbols", {}).items(): - table.symbols[f_path] = {} + # Normalize path on load just in case + norm_f_path = SymbolTable._normalize_path(f_path) + table.symbols[norm_f_path] = {} for scope_name, names in scopes.items(): - table.symbols[f_path][scope_name] = {} + table.symbols[norm_f_path][scope_name] = {} for name, sym_data in names.items(): - table.symbols[f_path][scope_name][name] = Symbol.from_dict(sym_data) + sym = Symbol.from_dict(sym_data) + sym.file_path = norm_f_path + table.symbols[norm_f_path][scope_name][name] = sym # Reconstruct exports for name, sym_data in data.get("exports", {}).items(): - table.exports[name] = Symbol.from_dict(sym_data) + sym = Symbol.from_dict(sym_data) + sym.file_path = SymbolTable._normalize_path(sym.file_path) + table.exports[name] = sym + + # Reconstruct metadata + file_metadata = data.get("file_metadata", {}) + table.file_metadata = { + SymbolTable._normalize_path(k): v for k, v in file_metadata.items() + } return table diff --git a/refactron/analysis/taint.py b/refactron/analysis/taint.py index 92da4f1..45b7661 100644 --- a/refactron/analysis/taint.py +++ b/refactron/analysis/taint.py @@ -76,6 +76,19 @@ def __init__(self, cfg_entry: CFGNode, config: TaintConfig = DEFAULT_TAINT_CONFI self._sources = {s.name for s in config.sources} self._sinks = {s.name: s for s in config.sinks} self._sanitizers = set(config.sanitizers) + self._statement_meta: Dict[ast.AST, List[ast.AST]] = {} + self._index_sensitive_nodes() + + def _index_sensitive_nodes(self): + """Pre-index statements that contain potential sinks or sources.""" + for node in self.data_flow.nodes: + for stmt in node.statements: + sensitive = [] + for child in ast.walk(stmt): + if isinstance(child, (ast.Call, ast.Attribute)): + sensitive.append(child) + if sensitive: + self._statement_meta[stmt] = sensitive def analyze(self) -> List[TaintVulnerability]: """ @@ -109,20 +122,22 @@ def analyze(self) -> List[TaintVulnerability]: for pred in node.predecessors: incoming_taint.update(tainted_vars[pred.id]) + # 2. Iterative Taint Propagation # Process block statements current_taint = incoming_taint.copy() for stmt in node.statements: + # Shared memo for the entire statement processing + memo: Dict[ast.AST, bool] = {} + # Check for Sink usage - vuls = self._check_sink(stmt, current_taint, node.id) - # We accumulate vulnerabilities but continue analysis - # De-duplicate vuls? + vuls = self._check_sink(stmt, current_taint, node.id, memo) for v in vuls: if v not in vulnerabilities: vulnerabilities.append(v) # Update Taint (Sources & Propagation) - new_taints, cleansed = self._propagate_taint(stmt, current_taint) + new_taints, cleansed = self._propagate_taint(stmt, current_taint, memo) current_taint.update(new_taints) current_taint.difference_update(cleansed) @@ -137,7 +152,9 @@ def analyze(self) -> List[TaintVulnerability]: return vulnerabilities - def _propagate_taint(self, stmt: ast.AST, current_taint: Set[str]) -> Tuple[Set[str], Set[str]]: + def _propagate_taint( + self, stmt: ast.AST, current_taint: Set[str], memo: Dict[ast.AST, bool] + ) -> Tuple[Set[str], Set[str]]: """ Analyze a statement and return (newly_tainted_vars, cleansed_vars). """ @@ -152,7 +169,7 @@ def _propagate_taint(self, stmt: ast.AST, current_taint: Set[str]) -> Tuple[Set[ targets = [stmt.target] value = stmt.value # type: ignore[union-attr, attr-defined] - is_tainted = self._is_expression_tainted(value, current_taint) # type: ignore[arg-type] + is_tainted = self._is_expression_tainted(value, current_taint, memo) # type: ignore[arg-type] for target in targets: if isinstance(target, ast.Name): @@ -164,70 +181,84 @@ def _propagate_taint(self, stmt: ast.AST, current_taint: Set[str]) -> Tuple[Set[ return generated, killed - def _is_expression_tainted(self, expr: ast.AST, current_taint: Set[str]) -> bool: + def _is_expression_tainted( + self, expr: ast.AST, current_taint: Set[str], memo: Dict[ast.AST, bool] + ) -> bool: """Check if an expression evaluates to a tainted value.""" + if expr in memo: + return memo[expr] + + result = False if isinstance(expr, ast.Name): # Check if variable is already tainted if expr.id in current_taint: - return True + result = True # Check if it's a direct source (e.g. 'request') - if expr.id in self._sources: - return True + elif expr.id in self._sources: + result = True elif isinstance(expr, ast.Call): # Check if function call returns taint (Source) func_name = self._get_call_name(expr) if func_name in self._sources: - return True - + result = True # Check if function call propagates taint (Sanitizer check) - if func_name in self._sanitizers: - return False - - # Default: propagate if any arg is tainted - for arg in expr.args: - if self._is_expression_tainted(arg, current_taint): - return True + elif func_name in self._sanitizers: + result = False + else: + # Default: propagate if any arg is tainted + for arg in expr.args: + if self._is_expression_tainted(arg, current_taint, memo): + result = True + break elif isinstance(expr, ast.BinOp): # Binary op is tainted if either side is tainted - return self._is_expression_tainted( - expr.left, current_taint - ) or self._is_expression_tainted(expr.right, current_taint) + result = self._is_expression_tainted( + expr.left, current_taint, memo + ) or self._is_expression_tainted(expr.right, current_taint, memo) elif isinstance(expr, ast.JoinedStr): # f-string: tainted if any value in it is tainted for value in expr.values: if isinstance(value, ast.FormattedValue): - if self._is_expression_tainted(value.value, current_taint): - return True - elif self._is_expression_tainted(value, current_taint): - return True + if self._is_expression_tainted(value.value, current_taint, memo): + result = True + break + elif self._is_expression_tainted(value, current_taint, memo): + result = True + break elif isinstance(expr, ast.Subscript): # Propagate from value (e.g. args.input) - if self._is_expression_tainted(expr.value, current_taint): - return True + result = self._is_expression_tainted(expr.value, current_taint, memo) elif isinstance(expr, ast.Attribute): # Check specific attributes (e.g. os.environ) full_name = self._get_attribute_name(expr) if full_name in self._sources: - return True + result = True # Propagate from object (simple object taint) - if self._is_expression_tainted(expr.value, current_taint): - return True + elif self._is_expression_tainted(expr.value, current_taint, memo): + result = True - return False + memo[expr] = result + return result def _check_sink( - self, stmt: ast.AST, current_taint: Set[str], node_id: int + self, + stmt: ast.AST, + current_taint: Set[str], + node_id: int, + memo: Dict[ast.AST, bool], ) -> List[TaintVulnerability]: """Check if a statement uses a tainted variable in a sink.""" vuls = [] - # We need to traverse the statement to find Call nodes - for node in ast.walk(stmt): + # Use pre-indexed potentially sensitive nodes instead of ast.walk + sensitive_nodes = self._statement_meta.get(stmt, []) + + for node in sensitive_nodes: if isinstance(node, ast.Call): func_name = self._get_call_name(node) if func_name in self._sinks: @@ -235,7 +266,7 @@ def _check_sink( # Check the sensitive argument if len(node.args) > sink_def.arg_index: arg = node.args[sink_def.arg_index] - if self._is_expression_tainted(arg, current_taint): + if self._is_expression_tainted(arg, current_taint, memo): # Identify which variable caused it for reporting var_name = "expression" if isinstance(arg, ast.Name): diff --git a/refactron/cli/analysis.py b/refactron/cli/analysis.py index 95877f5..692974b 100644 --- a/refactron/cli/analysis.py +++ b/refactron/cli/analysis.py @@ -524,17 +524,7 @@ def suggest(target: Optional[str], line: Optional[int], interactive: bool, apply console.print(f"[bold]Line:[/bold] {line}") # 2. Initialize Components - try: - retriever = ContextRetriever(workspace_path) - console.print("[dim]RAG Index loaded.[/dim]") - except Exception: - console.print( - "[yellow]Warning: RAG index not found. Context retrieval will be limited.[/yellow]" - ) - console.print("[dim]Run 'refactron rag index' to enable full context.[/dim]") - retriever = None - - orchestrator = LLMOrchestrator(retriever=retriever) + orchestrator = LLMOrchestrator(workspace_path=workspace_path) # 3. Read Code start_line_idx = 0 diff --git a/refactron/cli/main.py b/refactron/cli/main.py index 93a777a..a506d74 100644 --- a/refactron/cli/main.py +++ b/refactron/cli/main.py @@ -34,9 +34,15 @@ def main(ctx: click.Context) -> None: exempt_commands = ["login", "logout", "auth"] # 1. Pre-check authentication status + import os + creds = load_credentials() is_authenticated = False - if creds and creds.access_token: + + # If using local GROQ, we bypass cloud authentication checks + if os.environ.get("GROQ_API_KEY"): + is_authenticated = True + elif creds and creds.access_token: now = datetime.now(timezone.utc) if not creds.expires_at or creds.expires_at > now: is_authenticated = True @@ -115,10 +121,11 @@ def main(ctx: click.Context) -> None: pass try: - from refactron.cli.refactor import autofix, document, refactor, rollback + from refactron.cli.refactor import ai_fix, autofix, document, refactor, rollback main.add_command(refactor) main.add_command(autofix) + main.add_command(ai_fix) main.add_command(rollback) main.add_command(document) except ImportError: diff --git a/refactron/cli/rag.py b/refactron/cli/rag.py index e34345b..cebe9da 100644 --- a/refactron/cli/rag.py +++ b/refactron/cli/rag.py @@ -16,6 +16,7 @@ from refactron.cli.ui import _auth_banner, console from refactron.cli.utils import _setup_logging from refactron.core.workspace import WorkspaceManager +from refactron.llm.orchestrator import LLMOrchestrator from refactron.rag.indexer import RAGIndexer from refactron.rag.retriever import ContextRetriever @@ -68,14 +69,18 @@ def rag_index(background: bool, summarize: bool) -> None: console.print(f"[primary]Indexing:[/primary] {current_workspace.repo_full_name}\n") try: + orchestrator = LLMOrchestrator(workspace_path=local_path) + if background: # Run without visual feedback - indexer = RAGIndexer(local_path) - indexer.index_repository(local_path, summarize=summarize) + orchestrator.build_vector_index(local_path, summarize=summarize) else: with console.status("[primary]Parsing and indexing code...[/primary]"): + orchestrator.build_vector_index(local_path, summarize=summarize) + + # We need stats for the panel, get them from indexer indexer = RAGIndexer(local_path) - stats = indexer.index_repository(local_path, summarize=summarize) + stats = indexer.get_stats() console.print( Panel( @@ -141,16 +146,14 @@ def rag_search(query: str, top_k: int, chunk_type: Optional[str], rerank: bool) # AI Reranking if enabled if rerank: try: - from refactron.llm.client import GroqClient - - client = GroqClient() + orchestrator = LLMOrchestrator(workspace_path=local_path) prompt = ( # noqa: E501 f"Rate the relevance of the following code snippet to the user query: '{query}'\n\n" # noqa: E501 f"Code:\n{result.content[:500]}\n\n" "Provide only a percentage number (e.g. 85%) representing how well this code matches " # noqa: E501 "the semantic intent of the query." ) - ai_response = client.generate( + ai_response = orchestrator.client.generate( prompt=prompt, system="You are a code relevance evaluator. Output only the percentage.", max_tokens=10, diff --git a/refactron/cli/refactor.py b/refactron/cli/refactor.py index 5fd7c83..66b3213 100644 --- a/refactron/cli/refactor.py +++ b/refactron/cli/refactor.py @@ -527,15 +527,7 @@ def document(target: str, apply: bool, interactive: bool) -> None: console.print(f"[bold]Documenting:[/bold] {target_path}") # Initialize components - try: - retriever = ContextRetriever(workspace_path) - except Exception: - console.print( - "[yellow]Warning: RAG index not found. Context retrieval will be limited.[/yellow]" - ) - retriever = None - - orchestrator = LLMOrchestrator(retriever=retriever) + orchestrator = LLMOrchestrator(workspace_path=workspace_path) # Generate code = target_path.read_text(encoding="utf-8") @@ -587,3 +579,134 @@ def document(target: str, apply: bool, interactive: bool) -> None: except Exception as e: console.print(f"[red]Failed to create documentation: {e}[/red]") + + +@click.command() +@click.argument("target", type=click.Path(exists=True)) +@click.option( + "--config", + "-c", + type=click.Path(exists=True), + help="Path to configuration file", +) +@click.option( + "--apply/--no-apply", + default=False, + help="Apply the suggested changes to the file without prompt (if interactive is off)", +) +@click.option( + "--interactive/--no-interactive", + default=True, + help="Stop and prompt before applying each AI fix", +) +def ai_fix(target: str, config: Optional[str], apply: bool, interactive: bool) -> None: + """ + Orchestrate automated AI fixes for code issues. + + Analyzes the specified TARGET file, finds critical code issues, + and runs the LLM Orchestrator to suggest and apply fixes. + """ + console.print() + _auth_banner("AI Auto-Fix Orchestrator") + console.print() + + # 1. Setup + cfg = _load_config(config) + cfg.enable_incremental_analysis = False + _setup_logging() + + target_path = Path(target).resolve() + + if not target_path.is_file(): + console.print( + "[red]Error: Please specify a single file. Directory analysis for ai-fix is coming soon.[/red]" + ) + return + + refactron_instance = Refactron(cfg) + workspace_path = refactron_instance.detect_project_root(target_path) + + console.print(f"[bold]Step 1: Analyzing[/bold] {target_path}") + + # Run analysis + try: + with console.status("[primary]Running static analyzers...[/primary]"): + result = refactron_instance.analyze(str(target_path)) + except Exception as e: + console.print(f"[red]Analysis failed: {e}[/red]") + raise SystemExit(1) + + # Get issues for the file + issues = [i for i in result.all_issues if str(target_path) in str(i.file_path)] + issues_to_fix = issues + + if not issues_to_fix: + console.print("\n[green]Excellent! No issues found that require LLM fixing![/green]") + return + + console.print(f"\n[bold]Found {len(issues_to_fix)} issues to fix.[/bold]") + + # 2. Init AI components + orchestrator = LLMOrchestrator(workspace_path=workspace_path) + + # 3. Process issues in batch + file_content = target_path.read_text(encoding="utf-8") + + console.print( + f"\n[bold cyan]--- Orchestrating Batch Fix for {len(issues_to_fix)} issues ---[/bold cyan]" + ) + for idx, issue in enumerate(issues_to_fix, 1): + console.print( + f" {idx}. [yellow]{issue.category.value}[/yellow]: {issue.message} (Line {issue.line_number})" + ) + + with console.status("[bold cyan]Asking LLM for a comprehensive solution...[/bold cyan]"): + suggestion = orchestrator.generate_batch_suggestion(issues_to_fix, file_content) + + if suggestion.status == SuggestionStatus.FAILED: + console.print(f"\n[red]AI Failed to generate a batch fix:[/red] {suggestion.explanation}") + return + + console.print() + console.print( + Panel( + Markdown(suggestion.explanation), + title=f"AI Unified Fix Proposal ({suggestion.model_name})", + border_style="green", + ) + ) + + console.print( + Panel(suggestion.proposed_code, title="Proposed Unified Code", style="on #1e1e1e") + ) + console.print( + f"[dim]AI Confidence: {suggestion.llm_confidence:.2f}, Safety Score: {suggestion.confidence_score:.2f}[/dim]" + ) + + if suggestion.safety_result and not suggestion.safety_result.passed: + console.print( + f"[red]Warning: Fix failed basic safety checks: {', '.join(suggestion.safety_result.issues)}[/red]" + ) + + do_apply = apply + if interactive: + do_apply = click.confirm("\nDo you want to apply this unified AI fix to the file?") + + if do_apply: + try: + # Backup + backup_sys = BackupRollbackSystem(workspace_path) + session_id, _ = backup_sys.prepare_for_refactoring( + [target_path], description=f"AI Batch Fix for {len(issues_to_fix)} issues" + ) + console.print(f"[dim]Backup created: {session_id}[/dim]") + + # Apply changes + target_path.write_text(suggestion.proposed_code, encoding="utf-8") + console.print("[green bold]Applied unified AI fix successfully![/green bold]") + except Exception as e: + console.print(f"[red]Failed to apply fix: {e}[/red]") + else: + console.print("[yellow]Skipping fix application.[/yellow]") + + console.print("\n[bold]AI Auto-Fix Complete.[/bold] Verified and applied unified refactoring.") diff --git a/refactron/cli/ui.py b/refactron/cli/ui.py index e121c40..10ca289 100644 --- a/refactron/cli/ui.py +++ b/refactron/cli/ui.py @@ -347,25 +347,46 @@ def _handle_group_key(state: TuiState, key: str) -> TuiState: def _read_key() -> str: """Read a single keypress from stdin, handling escape sequences for arrow keys.""" - import termios - import tty - - fd = sys.stdin.fileno() - old_settings = termios.tcgetattr(fd) - try: - tty.setraw(fd) - ch = sys.stdin.read(1) - if ch == "\x1b": - ch2 = sys.stdin.read(1) - if ch2 == "[": - ch3 = sys.stdin.read(1) - return "\x1b[" + ch3 - return ch + ch2 - if ch == "\n": + import sys + import platform + + if platform.system() == "Windows": + import msvcrt + + ch = msvcrt.getch() + + # Handle special keys in Windows (arrows start with b'\xe0' or b'\x00') + if ch in (b"\xe0", b"\x00"): + ch2 = msvcrt.getch() + if ch2 == b"H": # Up arrow + return KEY_UP + elif ch2 == b"P": # Down arrow + return KEY_DOWN + return ch.decode("utf-8", "ignore") + ch2.decode("utf-8", "ignore") + elif ch in (b"\r", b"\n"): return KEY_ENTER - return ch - finally: - termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) + else: + return ch.decode("utf-8", "ignore") + else: + import termios + import tty + + fd = sys.stdin.fileno() + old_settings = termios.tcgetattr(fd) + try: + tty.setraw(fd) + ch = sys.stdin.read(1) + if ch == "\x1b": + ch2 = sys.stdin.read(1) + if ch2 == "[": + ch3 = sys.stdin.read(1) + return "\x1b[" + ch3 + return ch + ch2 + if ch in ("\n", "\r"): + return KEY_ENTER + return ch + finally: + termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) def _render_tui_summary(state: TuiState, target_path: Any) -> Text: @@ -747,12 +768,12 @@ def _run_startup_animation() -> None: console.clear() LOGO_LINES = [ - r"██████╗ ███████╗███████╗ █████╗ ██████╗████████╗██████╗ ██████╗ ███╗ ██╗", - r"██╔══██╗██╔════╝██╔════╝██╔══██╗██╔════╝╚══██╔══╝██╔══██╗██╔═══██╗████╗ ██║", - r"██████╔╝█████╗ █████╗ ███████║██║ ██║ ██████╔╝██║ ██║██╔██╗ ██║", - r"██╔══██╗██╔══╝ ██╔══╝ ██╔══██║██║ ██║ ██╔══██╗██║ ██║██║╚██╗██║", - r"██║ ██║███████╗██║ ██║ ██║╚██████╗ ██║ ██║ ██║╚██████╔╝██║ ╚████║", - r"╚═╝ ╚═╝╚══════╝╚═╝ ╚═╝ ╚═╝ ╚═════╝ ╚═╝ ╚═╝ ╚═╝ ╚═════╝ ╚═╝ ╚═══╝", + r" ____ _ _ ", + r" | _ \ ___ / _| __ _ ___| |_ _ __ ___ _ __ ", + r" | |_) / _ \ |_ / _` |/ __| __| '__/ _ \| '_ \ ", + r" | _ < __/ _| (_| | (__| |_| | | (_) | | | | ", + r" |_| \_\___|_| \__,_|\___|\__|_| \___/|_| |_| ", + r" ", ] subtitle_text = "The Intelligent Code Refactoring Transformer" @@ -924,11 +945,11 @@ def print_header() -> None: # Avatar (Simple ASCII) avatar = """ - ▄▄▄ - █████ - ███████ - █████████ - ███ █ █ ███ + _ + (_) + / \\ + | | + \\___/ """ info = Table.grid(padding=(0, 1)) diff --git a/refactron/cli/utils.py b/refactron/cli/utils.py index 025ff32..4933e84 100644 --- a/refactron/cli/utils.py +++ b/refactron/cli/utils.py @@ -114,7 +114,7 @@ def _setup_logging(verbose: bool = False) -> None: from transformers import logging as tf_logging tf_logging.set_verbosity_error() - except ImportError: + except Exception: pass diff --git a/refactron/core/config.py b/refactron/core/config.py index 4f0c830..0c296d6 100644 --- a/refactron/core/config.py +++ b/refactron/core/config.py @@ -130,6 +130,9 @@ class RefactronConfig: pattern_learning_enabled: bool = True # Enable learning from feedback pattern_ranking_enabled: bool = True # Enable ranking based on learned patterns + # LLM settings + enable_llm_triage: bool = False # Analyze potential false positives with the LLM + @classmethod def from_file( cls, @@ -243,6 +246,7 @@ def to_file(self, config_path: Path) -> None: ), "pattern_learning_enabled": self.pattern_learning_enabled, "pattern_ranking_enabled": self.pattern_ranking_enabled, + "enable_llm_triage": getattr(self, "enable_llm_triage", False), } try: diff --git a/refactron/core/inference.py b/refactron/core/inference.py index 99c42f9..7991032 100644 --- a/refactron/core/inference.py +++ b/refactron/core/inference.py @@ -3,6 +3,8 @@ Provides capabilities to infer types, values, and resolve symbols. """ +import os +from pathlib import Path from typing import Any, List, Optional import astroid @@ -28,10 +30,67 @@ def parse_string(code: str, module_name: str = "") -> nodes.Module: @staticmethod def parse_file(file_path: str) -> nodes.Module: """Parse a file into an astroid node tree.""" - builder = astroid.builder.AstroidBuilder(astroid.MANAGER) - with open(file_path, "r", encoding="utf-8") as f: - code = f.read() - return builder.string_build(code, modname=file_path) + # Use canonical path (resolved and posix-style for consistency) + abs_path = Path(file_path).resolve().as_posix() + manager = astroid.MANAGER + + # Aggressively clear cache for this file to ensure fresh AST + # Try both resolved and absolute paths to handle symlinks and normalization differences + raw_abs = os.path.abspath(file_path) + manager.astroid_cache.pop(abs_path, None) + manager.astroid_cache.pop(raw_abs, None) + manager.astroid_cache.pop(file_path, None) + + # 2. Find and clear by module name if it exists in caches + file_to_mod = getattr(manager, "file_to_module_cache", {}) + # Some versions use _mod_file_cache + if not file_to_mod: + file_to_mod = getattr(manager, "_mod_file_cache", {}) + + modname = ( + file_to_mod.get(abs_path) or file_to_mod.get(raw_abs) or file_to_mod.get(file_path) + ) + if modname: + manager.astroid_cache.pop(modname, None) + + # 3. Exhaustive search in astroid_cache for any module pointing to this file + for key, val in list(manager.astroid_cache.items()): + if hasattr(val, "file") and val.file: + val_path = Path(val.file).resolve().as_posix() + if val_path == abs_path or val_path == raw_abs.replace("\\", "/"): + manager.astroid_cache.pop(key, None) + + # 4. Clear the mappings themselves + for attr in ("file_to_module_cache", "_mod_file_cache"): + cache = getattr(manager, attr, None) + if isinstance(cache, dict): + cache.pop(abs_path, None) + cache.pop(raw_abs, None) + cache.pop(file_path, None) + + # 5. Read file and parse directly to bypass astroid's file cache + try: + with open(abs_path, "r", encoding="utf-8") as f: + code = f.read() + + # Resolve module name to keep astroid's state consistent + modname = "" + try: + from astroid import modutils + + modname = modutils.modname_from_path(abs_path) + except Exception: + pass + + # Use string_build via parse to avoid manager.ast_from_file's internal caching + return astroid.parse(code, module_name=modname, path=abs_path) + except (OSError, UnicodeDecodeError): + # Fallback to manager if manual read fails + try: + return manager.ast_from_file(abs_path) + except Exception as e: + # Fallback for virtual/non-existent files if needed + raise ValueError(f"Failed to parse {abs_path}: {e}") @staticmethod def infer_node(node: nodes.NodeNG, context: Optional[InferenceContext] = None) -> List[Any]: diff --git a/refactron/core/parallel.py b/refactron/core/parallel.py index 2fea0c0..808b6a6 100644 --- a/refactron/core/parallel.py +++ b/refactron/core/parallel.py @@ -7,7 +7,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple from refactron.core.analysis_result import FileAnalysisError -from refactron.core.models import FileMetrics +from refactron.core.models import FileMetrics, AnalysisSkipWarning logger = logging.getLogger(__name__) @@ -58,20 +58,25 @@ def __init__( def process_files( self, files: List[Path], - process_func: Callable[[Path], Tuple[Optional[FileMetrics], Optional[FileAnalysisError]]], + process_func: Callable[ + [Path], + Tuple[ + Optional[FileMetrics], Optional[FileAnalysisError], Optional[AnalysisSkipWarning] + ], + ], progress_callback: Optional[Callable[[int, int], None]] = None, - ) -> Tuple[List[FileMetrics], List[FileAnalysisError]]: + ) -> Tuple[List[FileMetrics], List[FileAnalysisError], List[AnalysisSkipWarning]]: """ Process multiple files in parallel. Args: files: List of file paths to process. process_func: Function to process a single file. Should return - (FileMetrics, None) on success or (None, FileAnalysisError) on error. + (FileMetrics, None, skip_warn) on success or (None, FileAnalysisError, None) on error. progress_callback: Optional callback for progress updates (completed, total). Returns: - Tuple of (successful results, failed files). + Tuple of (successful results, failed files, skip warnings). """ if not self.enabled or len(files) <= 1: # Process sequentially if disabled or only one file @@ -86,20 +91,28 @@ def process_files( def _process_sequential( self, files: List[Path], - process_func: Callable[[Path], Tuple[Optional[FileMetrics], Optional[FileAnalysisError]]], + process_func: Callable[ + [Path], + Tuple[ + Optional[FileMetrics], Optional[FileAnalysisError], Optional[AnalysisSkipWarning] + ], + ], progress_callback: Optional[Callable[[int, int], None]] = None, - ) -> Tuple[List[FileMetrics], List[FileAnalysisError]]: + ) -> Tuple[List[FileMetrics], List[FileAnalysisError], List[AnalysisSkipWarning]]: """Process files sequentially.""" results: List[FileMetrics] = [] errors: List[FileAnalysisError] = [] + skips: List[AnalysisSkipWarning] = [] for i, file_path in enumerate(files): try: - result, error = process_func(file_path) + result, error, skip = process_func(file_path) if result is not None: results.append(result) if error is not None: errors.append(error) + if skip is not None: + skips.append(skip) except Exception as e: logger.error(f"Unexpected error processing {file_path}: {e}", exc_info=True) errors.append( @@ -114,17 +127,23 @@ def _process_sequential( if progress_callback: progress_callback(i + 1, len(files)) - return results, errors + return results, errors, skips def _process_parallel_threads( self, files: List[Path], - process_func: Callable[[Path], Tuple[Optional[FileMetrics], Optional[FileAnalysisError]]], + process_func: Callable[ + [Path], + Tuple[ + Optional[FileMetrics], Optional[FileAnalysisError], Optional[AnalysisSkipWarning] + ], + ], progress_callback: Optional[Callable[[int, int], None]] = None, - ) -> Tuple[List[FileMetrics], List[FileAnalysisError]]: + ) -> Tuple[List[FileMetrics], List[FileAnalysisError], List[AnalysisSkipWarning]]: """Process files in parallel using threads.""" results: List[FileMetrics] = [] errors: List[FileAnalysisError] = [] + skips: List[AnalysisSkipWarning] = [] completed = 0 with ThreadPoolExecutor(max_workers=self.max_workers) as executor: @@ -139,11 +158,13 @@ def _process_parallel_threads( completed += 1 try: - result, error = future.result() + result, error, skip = future.result() if result is not None: results.append(result) if error is not None: errors.append(error) + if skip is not None: + skips.append(skip) except Exception as e: logger.error(f"Unexpected error processing {file_path}: {e}", exc_info=True) recovery_msg = "Check the file for syntax errors or encoding issues" @@ -159,14 +180,19 @@ def _process_parallel_threads( if progress_callback: progress_callback(completed, len(files)) - return results, errors + return results, errors, skips def _process_parallel_processes( self, files: List[Path], - process_func: Callable[[Path], Tuple[Optional[FileMetrics], Optional[FileAnalysisError]]], + process_func: Callable[ + [Path], + Tuple[ + Optional[FileMetrics], Optional[FileAnalysisError], Optional[AnalysisSkipWarning] + ], + ], progress_callback: Optional[Callable[[int, int], None]] = None, - ) -> Tuple[List[FileMetrics], List[FileAnalysisError]]: + ) -> Tuple[List[FileMetrics], List[FileAnalysisError], List[AnalysisSkipWarning]]: """ Process files in parallel using processes. @@ -176,6 +202,7 @@ def _process_parallel_processes( """ results: List[FileMetrics] = [] errors: List[FileAnalysisError] = [] + skips: List[AnalysisSkipWarning] = [] completed = 0 try: @@ -191,11 +218,13 @@ def _process_parallel_processes( completed += 1 try: - result, error = future.result() + result, error, skip = future.result() if result is not None: results.append(result) if error is not None: errors.append(error) + if skip is not None: + skips.append(skip) except Exception as e: logger.error(f"Unexpected error processing {file_path}: {e}", exc_info=True) recovery_msg = "Check the file for syntax errors or encoding issues" @@ -216,7 +245,7 @@ def _process_parallel_processes( logger.info("Falling back to sequential processing") return self._process_sequential(files, process_func, progress_callback) - return results, errors + return results, errors, skips def get_config(self) -> Dict[str, Any]: """ diff --git a/refactron/core/refactron.py b/refactron/core/refactron.py index 573004b..955162a 100644 --- a/refactron/core/refactron.py +++ b/refactron/core/refactron.py @@ -30,6 +30,7 @@ from refactron.core.refactor_result import RefactorResult from refactron.core.telemetry import get_telemetry_collector from refactron.patterns import PatternFingerprinter, PatternStorage +from refactron.llm.orchestrator import LLMOrchestrator from refactron.refactorers.add_docstring_refactorer import AddDocstringRefactorer from refactron.refactorers.base_refactorer import BaseRefactorer from refactron.refactorers.extract_method_refactorer import ExtractMethodRefactorer @@ -131,6 +132,15 @@ def __init__(self, config: Optional[RefactronConfig] = None): self.pattern_matcher = None self.pattern_ranker = None + # Initialize LLM Triage Orchestrator + self.llm_orchestrator = None + if getattr(self.config, "enable_llm_triage", False): + try: + self.llm_orchestrator = LLMOrchestrator() + logger.debug("LLM Orchestrator initialized for triage.") + except Exception as e: + logger.warning(f"Failed to initialize LLM for triage: {e}") + if self.config.enable_pattern_learning: try: # Initialize storage with custom directory if provided @@ -267,17 +277,18 @@ def analyze(self, target: Union[str, Path]) -> AnalysisResult: # Create a wrapper function for parallel processing def process_file_wrapper( file_path: Path, - ) -> Tuple[Optional[FileMetrics], Optional[FileAnalysisError]]: + ) -> Tuple[ + Optional[FileMetrics], Optional[FileAnalysisError], Optional[AnalysisSkipWarning] + ]: try: file_metrics, skip_warn = self._analyze_file(file_path) - if skip_warn is not None: - result.semantic_skip_warnings.append(skip_warn) + # Warnings are collected by ParallelProcessor.process_files return # Update incremental tracker if self.incremental_tracker.enabled: self.incremental_tracker.update_file_state(file_path) - return file_metrics, None + return file_metrics, None, skip_warn except AnalysisError as e: logger.debug(f"Failed to analyze {file_path}: {e}") error = FileAnalysisError( @@ -286,7 +297,7 @@ def process_file_wrapper( error_type=e.__class__.__name__, recovery_suggestion=e.recovery_suggestion, ) - return None, error + return None, error, None except Exception as e: logger.error(f"Unexpected error analyzing {file_path}: {e}", exc_info=True) error = FileAnalysisError( @@ -295,16 +306,17 @@ def process_file_wrapper( error_type=e.__class__.__name__, recovery_suggestion="Check the file for syntax errors or encoding issues", ) - return None, error + return None, error, None # Process files in parallel - file_metrics_list, error_list = self.parallel_processor.process_files( + file_metrics_list, error_list, skip_warnings = self.parallel_processor.process_files( files, process_file_wrapper, ) result.file_metrics.extend(file_metrics_list) result.failed_files.extend(error_list) + result.semantic_skip_warnings.extend(skip_warnings) result.total_issues = sum(fm.issue_count for fm in file_metrics_list) else: # Sequential processing @@ -331,13 +343,12 @@ def process_file_wrapper( ) except Exception as e: logger.error(f"Unexpected error analyzing {file_path}: {e}", exc_info=True) - recovery_msg = "Check the file for syntax errors or encoding issues" result.failed_files.append( FileAnalysisError( file_path=file_path, error_message=str(e), error_type=e.__class__.__name__, - recovery_suggestion=recovery_msg, + recovery_suggestion="Check the file for syntax errors or encoding issues", ) ) @@ -346,7 +357,7 @@ def process_file_wrapper( skipped_count = len(result.semantic_skip_warnings) if total_analyzed > 0 and skipped_count / total_analyzed > 0.10: result.semantic_skip_summary = ( - f"⚠ Semantic analysis (taint) was skipped for {skipped_count} of " + f"(warning) Semantic analysis (taint) was skipped for {skipped_count} of " f"{total_analyzed} files ({skipped_count / total_analyzed * 100:.0f}%). " "Check logs for details. Common causes: unsupported syntax or very large files." ) @@ -371,7 +382,7 @@ def process_file_wrapper( analyzers_used=analyzer_names, ) - # End memory profiling + # Final memory snapshot if self.memory_profiler.enabled: self.memory_profiler.snapshot("analysis_end") diff = self.memory_profiler.compare("analysis_start", "analysis_end") @@ -514,6 +525,49 @@ def _analyze_file( # Run semantic analysis (TaintAnalyzer) with full exception isolation _, skip_warning = self._run_semantic_analysis(file_path, source_code) + # Optional LLM Triage to suppress false positives + if self.llm_orchestrator and metrics.issues: + try: + start_triage = time.time() + logger.debug( + f"Running LLM triage on {len(metrics.issues)} issues in {file_path.name}" + ) + confidence_map = self.llm_orchestrator.evaluate_issues_batch( + metrics.issues, source_code + ) + + filtered_issues = [] + for i, issue in enumerate(metrics.issues): + base_id = getattr(issue, "rule_id", None) or "issue" + line_number = getattr(issue, "line_number", None) + id_parts = [str(base_id)] + if line_number is not None: + id_parts.append(str(line_number)) + id_parts.append(str(i)) + issue_id = ":".join(id_parts) + + matched_confidence = 1.0 + for k, v in confidence_map.items(): + if k == issue_id or k.startswith(f"{issue_id}_"): + matched_confidence = v + break + + if matched_confidence >= 0.5: + filtered_issues.append(issue) + else: + logger.debug( + f"LLM suppressed false positive: {issue.message} (score: {matched_confidence})" + ) + + metrics.issues = filtered_issues + logger.debug( + f"Completed LLM triage in {time.time() - start_triage:.2f}s. Kept {len(metrics.issues)} issues." + ) + except Exception as e: + logger.warning( + f"LLM Triage failed for {file_path}, falling back to static results: {e}" + ) + # Record file analysis metrics if self.metrics_collector and self.config.metrics_detailed: analysis_time_ms = (time.time() - start_time) * 1000 diff --git a/refactron/core/workspace.py b/refactron/core/workspace.py index 88c476e..42f4314 100644 --- a/refactron/core/workspace.py +++ b/refactron/core/workspace.py @@ -37,11 +37,19 @@ def to_dict(self) -> Dict[str, Any]: @classmethod def from_dict(cls, data: Dict[str, Any]) -> "WorkspaceMapping": """Create from dictionary.""" + # Robust parsing to handle older configuration formats + repo_full_name = data.get("repo_full_name", "unknown/unknown") + repo_name = data.get("repo_name") + + # Derive repo_name from full_name if missing + if not repo_name and "/" in repo_full_name: + repo_name = repo_full_name.split("/")[-1] + return cls( - repo_name=data["repo_name"], - repo_full_name=data["repo_full_name"], - local_path=data["local_path"], - connected_at=data["connected_at"], + repo_name=repo_name or repo_full_name or "unknown", + repo_full_name=repo_full_name, + local_path=data.get("local_path", ""), + connected_at=data.get("connected_at", "unknown"), repo_id=data.get("repo_id"), ) diff --git a/refactron/llm/backend_client.py b/refactron/llm/backend_client.py index 436d179..95cd2c2 100644 --- a/refactron/llm/backend_client.py +++ b/refactron/llm/backend_client.py @@ -17,7 +17,7 @@ def __init__( backend_url: Optional[str] = None, model: str = "llama-3.3-70b-versatile", temperature: float = 0.2, - max_tokens: int = 2000, + max_tokens: int = 4000, ): """Initialize backend client. diff --git a/refactron/llm/client.py b/refactron/llm/client.py index d7c9ba6..6bf2a23 100644 --- a/refactron/llm/client.py +++ b/refactron/llm/client.py @@ -21,7 +21,7 @@ def __init__( api_key: Optional[str] = None, model: str = "llama-3.3-70b-versatile", # Updated to current model temperature: float = 0.2, - max_tokens: int = 2000, + max_tokens: int = 4000, ): """Initialize Groq client. diff --git a/refactron/llm/orchestrator.py b/refactron/llm/orchestrator.py index 36ac47d..a380d0c 100644 --- a/refactron/llm/orchestrator.py +++ b/refactron/llm/orchestrator.py @@ -4,6 +4,7 @@ import logging import os import re +import ast from pathlib import Path from typing import Dict, List, Optional, Union @@ -12,6 +13,8 @@ from refactron.llm.client import GroqClient from refactron.llm.models import RefactoringSuggestion, SuggestionStatus from refactron.llm.prompts import ( + BATCH_SUGGESTION_PROMPT, + BATCH_SUGGESTION_SYSTEM_PROMPT, BATCH_TRIAGE_PROMPT, BATCH_TRIAGE_SYSTEM_PROMPT, DOCUMENTATION_PROMPT, @@ -30,10 +33,12 @@ class LLMOrchestrator: def __init__( self, retriever: Optional[ContextRetriever] = None, + workspace_path: Optional[Path] = None, llm_client: Optional[Union[GroqClient, BackendLLMClient]] = None, safety_gate: Optional[SafetyGate] = None, ): self.retriever = retriever + self.workspace_path = workspace_path if llm_client: self.client = llm_client @@ -50,6 +55,82 @@ def __init__( self.safety_gate = safety_gate or SafetyGate() + # Auto-initialize retriever if missing but workspace is provided + if not self.retriever and self.workspace_path: + self._ensure_retriever() + + def _ensure_retriever(self) -> None: + """Attempt to load or build the context retriever.""" + from refactron.rag.retriever import ContextRetriever + + try: + self.retriever = ContextRetriever(self.workspace_path) + except RuntimeError: + logger.info("RAG index missing. Auto-building index...") + self.build_vector_index(self.workspace_path, summarize=False) + + # Try loading again after build + try: + self.retriever = ContextRetriever(self.workspace_path) + except RuntimeError as final_e: + logger.warning( + f"Context retrieval will be limited. Failed to load RAG index: {final_e}" + ) + + def build_vector_index(self, workspace_path: Path, summarize: bool = False) -> None: + """Create or update the RAG vector index for the workspace. + + Args: + workspace_path: Path to the workspace directory + summarize: Whether to use AI to summarize code for better retrieval + """ + # Lazy import to prevent circular dependency + from refactron.rag.indexer import RAGIndexer + + try: + indexer = RAGIndexer(workspace_path=workspace_path, llm_integration=self) + indexer.index_repository(summarize=summarize) + + # Reload the retriever if it exists to pick up new chunks + if self.retriever: + # Assuming context retriever shares the same path concept + from refactron.rag.retriever import ContextRetriever + + self.retriever = ContextRetriever(workspace_path) + logger.info("Successfully updated LLMOrchestrator vector retrieval context.") + except Exception as e: + logger.error(f"Failed to build vector index: {e}") + + def generate_chunk_summary(self, chunk_content: str) -> Optional[str]: + """Generate a semantic summary of a code chunk for RAG indexing. + + Args: + chunk_content: The python code chunk + + Returns: + A one-sentence description of the chunk, or None on failure. + """ + prompt = ( + "Analyze the following Python code snippet and provide a one-sentence " + "summary of its purpose, focusing on what it DOES (e.g. 'Calculates user permissions' " + "or 'Handles secure database connections').\n\n" + f"Code:\n{chunk_content}" + ) + + try: + summary = self.client.generate( + prompt=prompt, + system=( + "You are a senior software architect. " + "Provide a concise, semantic summary of code purpose." + ), + max_tokens=100, + ) + return summary.strip() + except Exception as e: + logger.warning(f"AI summarization failed: {e}") + return None + def generate_suggestion(self, issue: CodeIssue, original_code: str) -> RefactoringSuggestion: """Generate a refactoring suggestion for a code issue. @@ -114,11 +195,32 @@ def generate_suggestion(self, issue: CodeIssue, original_code: str) -> Refactori except (ValueError, TypeError, AttributeError): confidence = 0.5 # Fallback + proposed_code = data.get("proposed_code", "") + + # Aggressively clean up hallucinated markdown inside the JSON string + if proposed_code.startswith("```"): + lines = proposed_code.split("\n") + if lines[0].startswith("```"): + lines.pop(0) + if lines and lines[-1].startswith("```"): + lines.pop(-1) + proposed_code = "\n".join(lines).strip() + + # If the LLM accidentally wrapped the code in JSON braces + if proposed_code.startswith("{") and proposed_code.endswith("}"): + potential_code = proposed_code[1:-1].strip() + try: + # Only accept the stripped version if it's valid Python syntax + ast.parse(potential_code) + proposed_code = potential_code + except SyntaxError: + pass + suggestion = RefactoringSuggestion( issue=issue, original_code=original_code, context_files=[r.file_path for r in results] if self.retriever else [], - proposed_code=data.get("proposed_code", ""), + proposed_code=proposed_code, explanation=data.get("explanation", "No explanation provided."), reasoning=data.get("reasoning", ""), model_name=self.client.model, @@ -162,6 +264,121 @@ def generate_suggestion(self, issue: CodeIssue, original_code: str) -> Refactori return suggestion + def generate_batch_suggestion( + self, issues: List[CodeIssue], original_code: str + ) -> RefactoringSuggestion: + """Generate a single refactoring suggestion that fixes a batch of issues. + + Args: + issues: List of code issues to fix + original_code: The original source code of the file + + Returns: + A combined refactoring suggestion + """ + if not issues: + raise ValueError("No issues provided for batch suggestion") + + # 1. Retrieve Context + context_snippets = [] + if self.retriever: + try: + # Use the first few issues for context retrieval + query = " ".join([i.message for i in issues[:3]]) + results = self.retriever.retrieve_similar(query, top_k=3) + context_snippets = [r.content for r in results] + except Exception as e: + logger.warning(f"Context retrieval failed: {e}") + + rag_context = "\n\n".join(context_snippets) if context_snippets else "No context available." + + # 2. Format issues for prompt + issues_details = "" + for idx, issue in enumerate(issues, 1): + issues_details += ( + f"{idx}. {issue.category.value} (Line {issue.line_number}): {issue.message}\n" + ) + + # 3. Construct Prompt + prompt = BATCH_SUGGESTION_PROMPT.format( + issues_details=issues_details, + original_code=original_code, + rag_context=rag_context, + ) + + # 4. Call LLM + response_text = "N/A" + try: + response_text = self.client.generate( + prompt=prompt, system=BATCH_SUGGESTION_SYSTEM_PROMPT, temperature=0.2 + ) + + # Reuse cleaning and parsing logic + clean_text = self._clean_json_response(response_text) + data = json.loads(clean_text, strict=False) + + proposed_code = data.get("proposed_code", "") + + # Clean up hallucinations + if proposed_code.startswith("```"): + lines = proposed_code.split("\n") + if lines[0].startswith("```"): + lines.pop(0) + if lines and lines[-1].startswith("```"): + lines.pop(-1) + proposed_code = "\n".join(lines).strip() + + if proposed_code.startswith("{") and proposed_code.endswith("}"): + potential_code = proposed_code[1:-1].strip() + try: + ast.parse(potential_code) + proposed_code = potential_code + except SyntaxError: + pass + + suggestion = RefactoringSuggestion( + issue=issues[0], # Use first issue as primary reference + original_code=original_code, + context_files=[r.file_path for r in results] if self.retriever else [], + proposed_code=proposed_code, + explanation=data.get("explanation", "Combined fix for multiple issues."), + reasoning=data.get("reasoning", ""), + model_name=self.client.model, + confidence_score=float(data.get("confidence_score", 0.7)), + llm_confidence=float(data.get("confidence_score", 0.7)), + ) + + except Exception as e: + logger.error(f"LLM batch generation failed: {e}") + return RefactoringSuggestion( + issue=issues[0], + original_code=original_code, + context_files=[], + proposed_code="", + explanation=f"Batch generation failed: {str(e)}", + reasoning="", + model_name=self.client.model, + confidence_score=0.0, + status=SuggestionStatus.FAILED, + ) + + # 5. Safety Validation + try: + safety_result = self.safety_gate.validate(suggestion) + suggestion.safety_result = safety_result + suggestion.confidence_score = safety_result.score + + if not safety_result.passed: + suggestion.status = SuggestionStatus.REJECTED + else: + suggestion.status = SuggestionStatus.PENDING + + except Exception as e: + logger.error(f"Safety validation failed: {e}") + suggestion.status = SuggestionStatus.FAILED + + return suggestion + def generate_documentation( self, code: str, file_path: str = "unknown" ) -> RefactoringSuggestion: diff --git a/refactron/llm/prompts.py b/refactron/llm/prompts.py index 20c5603..aa9f1ca 100644 --- a/refactron/llm/prompts.py +++ b/refactron/llm/prompts.py @@ -19,6 +19,80 @@ } """ +BATCH_TRIAGE_SYSTEM_PROMPT = """\ +You are an expert software architect and code refactoring specialist. +Your goal is to evaluate multiple code issues in a single file and determine their validity. + +RESPONSE FORMAT: +You must output ONLY valid JSON. +- Escape all double quotes inside strings with backslash (e.g. \\"). +- Do not use trailing commas. +- Do not output markdown code blocks, just the raw JSON object. + +Output JSON structure must be a simple dictionary mapping issue IDs to confidence scores: +{ + "issue_1": 0.85, + "issue_2": 0.1 +} +""" + +BATCH_TRIAGE_PROMPT = """ +You are a code triage expert. Evaluate the following list of code issues found in a +single file and determine the confidence that each is a true positive (requiring +fixing) rather than a false positive. + +File Source Code: +``` +{source_code} +``` + +Relevant Context (RAG): +{rag_context} + +Issues to evaluate: +{issues_json} + +Return ONLY a JSON map where the keys are the issue IDs and the values are the +confidence scores (float between 0.0 and 1.0). +Do NOT return anything except the JSON object. +""" + +BATCH_SUGGESTION_SYSTEM_PROMPT = """You are an expert software architect and code refactoring specialist. +Your goal is to analyze multiple code issues in a file and provide a single, comprehensive fix that resolves all of them. + +RESPONSE FORMAT: +You must output ONLY valid JSON. +- Escape all double quotes inside strings with backslash (e.g. \\"). +- Do not use trailing commas. +- Do not output markdown code blocks, just the raw JSON object. +- Ensure newlines in strings are escaped as \\n. + +Output JSON structure: +{ + "explanation": "Summary of all fixes applied", + "proposed_code": "The complete fixed code block for the entire file/context", + "reasoning": "Briefly explain how you addressed the issues", + "confidence_score": "Float between 0.0 and 1.0 representing your confidence in this combined fix" +} +""" + +BATCH_SUGGESTION_PROMPT = """ +Fix the following code issues in the file: + +Issues: +{issues_details} + +Original Code: +```python +{original_code} +``` + +Relevant Context (RAG): +{rag_context} + +Provide a single, comprehensive fix that resolves ALL the listed issues while maintaining consistency with the codebase. +""" + SUGGESTION_PROMPT = """ Fix the following code issue: @@ -93,40 +167,3 @@ @@@END@@@ """ -BATCH_TRIAGE_SYSTEM_PROMPT = """\ -You are an expert software architect and code refactoring specialist. -Your goal is to evaluate multiple code issues in a single file and determine their validity. - -RESPONSE FORMAT: -You must output ONLY valid JSON. -- Escape all double quotes inside strings with backslash (e.g. \\"). -- Do not use trailing commas. -- Do not output markdown code blocks, just the raw JSON object. - -Output JSON structure must be a simple dictionary mapping issue IDs to confidence scores: -{ - "issue_1": 0.85, - "issue_2": 0.1 -} -""" - -BATCH_TRIAGE_PROMPT = """ -You are a code triage expert. Evaluate the following list of code issues found in a -single file and determine the confidence that each is a true positive (requiring -fixing) rather than a false positive. - -File Source Code: -``` -{source_code} -``` - -Relevant Context (RAG): -{rag_context} - -Issues to evaluate: -{issues_json} - -Return ONLY a JSON map where the keys are the issue IDs and the values are the -confidence scores (float between 0.0 and 1.0). -Do NOT return anything except the JSON object. -""" diff --git a/refactron/rag/indexer.py b/refactron/rag/indexer.py index 690cfe3..9fad604 100644 --- a/refactron/rag/indexer.py +++ b/refactron/rag/indexer.py @@ -5,19 +5,14 @@ import json from dataclasses import dataclass from pathlib import Path -from typing import Any, Dict, List, Optional, cast +from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast -try: - import chromadb - from chromadb.config import Settings - from sentence_transformers import SentenceTransformer +if TYPE_CHECKING: + from refactron.llm.orchestrator import LLMOrchestrator - CHROMA_AVAILABLE = True -except ImportError: - chromadb = None - Settings = None - SentenceTransformer = None - CHROMA_AVAILABLE = False +# RAG dependencies are loaded lazily in __init__ to prevent CLI crashes +# if libraries like PyTorch fail to initialize (common on some Windows environments). +CHROMA_AVAILABLE = None from refactron.rag.chunker import CodeChunk from refactron.rag.parser import CodeParser @@ -48,7 +43,7 @@ def __init__( workspace_path: Path, embedding_model: str = "all-MiniLM-L6-v2", collection_name: str = "code_chunks", - llm_client: Optional[GroqClient] = None, + llm_integration: Optional["LLMOrchestrator"] = None, ): """Initialize the RAG indexer. @@ -56,39 +51,71 @@ def __init__( workspace_path: Path to the workspace directory embedding_model: Name of the sentence-transformers model collection_name: Name of the ChromaDB collection - llm_client: Optional LLM client for code summarization + llm_integration: Optional LLM Orchestrator for code summarization """ - if not CHROMA_AVAILABLE: - raise RuntimeError( - "ChromaDB is not available. " - "Install with: pip install chromadb sentence-transformers" - ) - self.workspace_path = Path(workspace_path) self.index_path = self.workspace_path / ".rag" self.index_path.mkdir(exist_ok=True) + self.llm_integration = llm_integration + self.parser = CodeParser() - # Initialize LLM client for summarization - # GroqClient is used for type hints and potential init below - - self.llm_client = llm_client - - # Initialize embedding model - self.embedding_model_name = embedding_model - self.embedding_model = SentenceTransformer(embedding_model) - - # Initialize ChromaDB - self.client = chromadb.PersistentClient( - path=str(self.index_path / "chroma"), settings=Settings(anonymized_telemetry=False) - ) - - # Get or create collection - self.collection = self.client.get_or_create_collection( - name=collection_name, - metadata={"embedding_model": embedding_model, "hnsw:space": "cosine"}, - ) + # Lazy load dependencies to avoid crashing the whole CLI + # when PyTorch DLL initialization fails (WinError 1114) + global CHROMA_AVAILABLE + if CHROMA_AVAILABLE is None: + try: + import chromadb as _chromadb + from chromadb.config import Settings as _Settings + from sentence_transformers import SentenceTransformer as _SentenceTransformer + + globals()["chromadb"] = _chromadb + globals()["Settings"] = _Settings + globals()["SentenceTransformer"] = _SentenceTransformer + CHROMA_AVAILABLE = True + except (ImportError, OSError) as e: + import logging + + logging.getLogger(__name__).warning( + f"RAG dependencies failed to load, falling back to keyword mode: {e}" + ) + CHROMA_AVAILABLE = False + + if CHROMA_AVAILABLE: + self.mode = "vector" + self.embedding_model_name = embedding_model + try: + self.embedding_model = globals()["SentenceTransformer"](embedding_model) + except Exception as e: + import logging + + logging.getLogger(__name__).error(f"Failed to initialize embedding model: {e}") + self.mode = "keyword" + self.embedding_model_name = "keyword-fallback" + + if self.mode == "vector": + # Initialize ChromaDB + self.client = globals()["chromadb"].PersistentClient( + path=str(self.index_path / "chroma"), + settings=globals()["Settings"](anonymized_telemetry=False), + ) + + # Get or create collection + self.collection = self.client.get_or_create_collection( + name=collection_name, + metadata={"embedding_model": embedding_model, "hnsw:space": "cosine"}, + ) + else: + self.mode = "keyword" + self.embedding_model_name = "keyword-fallback" + import logging + + logging.getLogger(__name__).info( + "Initializing in Keyword Fallback mode (No PyTorch/Chroma needed)." + ) - self.parser = CodeParser() + # In keyword mode, we store chunks in a human-readable JSON file + if self.mode == "keyword": + self.chunk_storage = self.index_path / "keyword_chunks.json" def index_repository( self, repo_path: Optional[Path] = None, summarize: bool = False @@ -102,11 +129,11 @@ def index_repository( Returns: Statistics about the indexed content """ - if summarize and not self.llm_client: - from refactron.llm.client import GroqClient + if summarize and not self.llm_integration: + from refactron.llm.orchestrator import LLMOrchestrator try: - self.llm_client = GroqClient() + self.llm_integration = LLMOrchestrator() except Exception as e: print(f"Warning: Could not initialize AI for summarization: {e}") summarize = False @@ -177,7 +204,7 @@ def _index_file(self, file_path: Path, summarize: bool = False) -> List[CodeChun chunker = CodeChunker(self.parser) chunks = chunker.chunk_file(file_path) - if summarize and self.llm_client: + if summarize and self.llm_integration: for chunk in chunks: try: summary = self._summarize_chunk(chunk) @@ -195,31 +222,14 @@ def _index_file(self, file_path: Path, summarize: bool = False) -> List[CodeChun def _summarize_chunk(self, chunk: CodeChunk) -> Optional[str]: """Use AI to generate a brief semantic summary of a code chunk.""" - if not self.llm_client: + if not self.llm_integration: return None - prompt = ( - "Analyze the following Python code snippet and provide a one-sentence " - "summary of its purpose, focusing on what it DOES (e.g. 'Calculates user permissions' " - "or 'Handles secure database connections').\n\n" - f"Code:\n{chunk.content}" - ) - - try: - summary = self.llm_client.generate( - prompt=prompt, - system=( - "You are a senior software architect. " - "Provide a concise, semantic summary of code purpose." - ), - max_tokens=100, - ) - return summary.strip() - except Exception: - return None + # Delegate to the orchestration layer + return self.llm_integration.generate_chunk_summary(chunk.content) def add_chunks(self, chunks: List[CodeChunk]) -> None: - """Add code chunks to the vector index. + """Add code chunks to the vector index or keyword storage. Args: chunks: List of code chunks to add @@ -227,7 +237,40 @@ def add_chunks(self, chunks: List[CodeChunk]) -> None: if not chunks: return - # Prepare data for ChromaDB + if self.mode == "keyword": + # Keyword mode: append to JSON storage + current_data = [] + if self.chunk_storage.exists(): + try: + with open(self.chunk_storage, "r", encoding="utf-8") as f: + current_data = json.load(f) + except Exception: + current_data = [] + + for chunk in chunks: + chunk_dict = { + "content": chunk.content, + "metadata": { + "chunk_type": chunk.chunk_type, + "file_path": chunk.file_path, + "name": chunk.name, + "line_start": chunk.line_range[0], + "line_end": chunk.line_range[1], + }, + } + # Add extra metadata + for key, value in chunk.metadata.items(): + if value is not None: + chunk_dict["metadata"][key] = value + + current_data.append(chunk_dict) + + # Deduplicate or just overwrite for now (simpler) + with open(self.chunk_storage, "w", encoding="utf-8") as f: + json.dump(current_data, f, indent=2) + return + + # Prepare data for ChromaDB (Vector Mode) documents = [chunk.content for chunk in chunks] metadatas = [] for chunk in chunks: diff --git a/refactron/rag/retriever.py b/refactron/rag/retriever.py index 9fcd16f..cf81e13 100644 --- a/refactron/rag/retriever.py +++ b/refactron/rag/retriever.py @@ -6,17 +6,9 @@ from pathlib import Path from typing import List, Optional -try: - import chromadb - from chromadb.config import Settings - from sentence_transformers import SentenceTransformer - - CHROMA_AVAILABLE = True -except ImportError: - chromadb = None - Settings = None - SentenceTransformer = None - CHROMA_AVAILABLE = False +# RAG dependencies are loaded lazily in __init__ to prevent CLI crashes +# if libraries like PyTorch fail to initialize (common on some Windows environments). +CHROMA_AVAILABLE = None @dataclass @@ -48,30 +40,80 @@ def __init__( embedding_model: Name of the sentence-transformers model collection_name: Name of the ChromaDB collection """ - if not CHROMA_AVAILABLE: - raise RuntimeError( - "ChromaDB is not available. " - "Install with: pip install chromadb sentence-transformers" - ) - self.workspace_path = Path(workspace_path) self.index_path = self.workspace_path / ".rag" - # Initialize embedding model - self.embedding_model = SentenceTransformer(embedding_model) - - # Initialize ChromaDB client - self.client = chromadb.PersistentClient( - path=str(self.index_path / "chroma"), settings=Settings(anonymized_telemetry=False) - ) + # Lazy load dependencies to avoid crashing the whole CLI + # when PyTorch DLL initialization fails (WinError 1114) + global CHROMA_AVAILABLE + if CHROMA_AVAILABLE is None: + try: + import chromadb as _chromadb + from chromadb.config import Settings as _Settings + from sentence_transformers import SentenceTransformer as _SentenceTransformer + + globals()["chromadb"] = _chromadb + globals()["Settings"] = _Settings + globals()["SentenceTransformer"] = _SentenceTransformer + CHROMA_AVAILABLE = True + except (ImportError, OSError) as e: + import logging + + logging.getLogger(__name__).warning( + f"RAG dependencies failed to load, falling back to keyword mode: {e}" + ) + CHROMA_AVAILABLE = False + + if CHROMA_AVAILABLE: + self.mode = "vector" + # Initialize embedding model + try: + self.embedding_model = globals()["SentenceTransformer"](embedding_model) + except Exception as e: + import logging + + logging.getLogger(__name__).error(f"Failed to initialize embedding model: {e}") + self.mode = "keyword" + + if self.mode == "vector": + # Initialize ChromaDB client + self.client = globals()["chromadb"].PersistentClient( + path=str(self.index_path / "chroma"), + settings=globals()["Settings"](anonymized_telemetry=False), + ) - # Get collection - try: - self.collection = self.client.get_collection(name=collection_name) - except Exception: - raise RuntimeError( - f"Index not found at {self.index_path}. Run 'refactron rag index' first." - ) + # Get collection + try: + self.collection = self.client.get_collection(name=collection_name) + except Exception: + # If vector collection is missing but folder exists, + # we might have indexed in keyword mode before + if (self.index_path / "keyword_chunks.json").exists(): + self.mode = "keyword" + else: + raise RuntimeError( + f"Index not found at {self.index_path}. Run 'refactron rag index' first." + ) + else: + self.mode = "keyword" + + # In keyword mode, we load chunks from the JSON file + if self.mode == "keyword": + self.keyword_chunks = [] + chunk_file = self.index_path / "keyword_chunks.json" + if chunk_file.exists(): + import json + + try: + with open(chunk_file, "r", encoding="utf-8") as f: + self.keyword_chunks = json.load(f) + except Exception as e: + import logging + + logging.getLogger(__name__).error(f"Failed to load keyword index: {e}") + else: + # Only raise if we are actually trying to search + pass def retrieve_similar( self, query: str, top_k: int = 5, chunk_type: Optional[str] = None @@ -86,6 +128,55 @@ def retrieve_similar( Returns: List of retrieved contexts sorted by relevance """ + if self.mode == "keyword": + if not self.keyword_chunks: + import logging + + logging.getLogger(__name__).warning("No keyword chunks found for retrieval.") + return [] + + # Implementation of simple keyword search + query_terms = set(query.lower().split()) + results = [] + + for chunk in self.keyword_chunks: + # Filter by chunk type if requested + if chunk_type and chunk["metadata"].get("chunk_type") != chunk_type: + continue + + content = chunk["content"] + # Simple score: count query terms that appear in content + content_lower = content.lower() + score = 0 + for term in query_terms: + if term in content_lower: + score += 1 + + # Normalize score + if score > 0: + # Invert score to look like 'distance' (lower is better) + results.append((chunk, 1.0 - (score / len(query_terms)))) + + # Sort by distance (relevance) + results.sort(key=lambda x: x[1]) + + # Convert to RetrievedContext + retrieved = [] + for chunk_data, dist in results[:top_k]: + meta = chunk_data["metadata"] + retrieved.append( + RetrievedContext( + content=chunk_data["content"], + file_path=meta.get("file_path", "unknown"), + chunk_type=meta.get("chunk_type", "unknown"), + name=meta.get("name", "unknown"), + line_range=(meta.get("line_start", 0), meta.get("line_end", 0)), + distance=dist, + metadata=meta, + ) + ) + return retrieved + # Generate query embedding query_embedding = self.embedding_model.encode([query], show_progress_bar=False).tolist()[0] diff --git a/tests/test_config_management.py b/tests/test_config_management.py index c71b71a..50afda8 100644 --- a/tests/test_config_management.py +++ b/tests/test_config_management.py @@ -903,16 +903,16 @@ def test_parallel_processor_sequential_and_thread_modes(tmp_path: Path) -> None: def process_func(p: Path): if p.name == "bad.py": raise ValueError("boom") - return None, None + return None, None, None p_seq = ParallelProcessor(max_workers=1, use_processes=False, enabled=True) - _, errors = p_seq.process_files(files, process_func) + _, errors, skips = p_seq.process_files(files, process_func) assert p_seq.enabled is False assert len(errors) == 1 assert isinstance(errors[0], FileAnalysisError) p_thr = ParallelProcessor(max_workers=2, use_processes=False, enabled=True) - results, errors = p_thr.process_files(files, lambda p: (None, None)) + results, errors, skips = p_thr.process_files(files, lambda p: (None, None, None)) assert results == [] assert errors == [] assert p_thr.get_config()["max_workers"] == 2 @@ -1171,11 +1171,11 @@ def make_error(path): def success_func(p): - return make_metrics(p), None + return make_metrics(p), None, None def error_func(p): - return None, make_error(p) + return None, make_error(p), None def raises_func(p): @@ -1210,23 +1210,23 @@ def test_get_config(self): class TestSequentialProcessing: def test_empty_files(self): pp = ParallelProcessor(enabled=False) - results, errors = pp.process_files([], success_func) - assert results == [] and errors == [] + results, errors, skips = pp.process_files([], success_func) + assert results == [] and errors == [] and skips == [] def test_single_file_success(self): pp = ParallelProcessor(enabled=False) files = [Path("a.py")] - results, errors = pp.process_files(files, success_func) - assert len(results) == 1 and len(errors) == 0 + results, errors, skips = pp.process_files(files, success_func) + assert len(results) == 1 and len(errors) == 0 and len(skips) == 0 def test_single_file_error(self): pp = ParallelProcessor(enabled=False) - results, errors = pp.process_files([Path("a.py")], error_func) - assert len(results) == 0 and len(errors) == 1 + results, errors, skips = pp.process_files([Path("a.py")], error_func) + assert len(results) == 0 and len(errors) == 1 and len(skips) == 0 def test_single_file_exception(self): pp = ParallelProcessor(enabled=False) - results, errors = pp.process_files([Path("a.py")], raises_func) + results, errors, skips = pp.process_files([Path("a.py")], raises_func) assert len(errors) == 1 def test_progress_callback(self): @@ -1244,12 +1244,13 @@ class TestThreadedProcessing: def test_two_files_threads(self): pp = ParallelProcessor(max_workers=2, use_processes=False, enabled=True) files = [Path("a.py"), Path("b.py")] - results, errors = pp.process_files(files, success_func) + results, errors, skips = pp.process_files(files, success_func) assert len(results) == 2 + assert len(skips) == 0 def test_thread_error_handling(self): pp = ParallelProcessor(max_workers=2, use_processes=False, enabled=True) - results, errors = pp.process_files([Path("a.py"), Path("b.py")], raises_func) + results, errors, skips = pp.process_files([Path("a.py"), Path("b.py")], raises_func) assert len(errors) == 2 def test_thread_progress_callback(self): @@ -1264,7 +1265,7 @@ def test_thread_progress_callback(self): def test_single_file_goes_sequential(self): pp = ParallelProcessor(max_workers=4, use_processes=False, enabled=True) - results, errors = pp.process_files([Path("a.py")], success_func) + results, errors, skips = pp.process_files([Path("a.py")], success_func) assert len(results) == 1 @@ -1274,13 +1275,13 @@ def test_process_pool_falls_back_on_exception(self): with patch( "refactron.core.parallel.ProcessPoolExecutor", side_effect=Exception("spawn fail") ): - results, errors = pp.process_files([Path("a.py")], success_func) + results, errors, skips = pp.process_files([Path("a.py")], success_func) assert len(results) == 1 def test_process_pool_success(self): pp = ParallelProcessor(max_workers=2, use_processes=True, enabled=True) mock_future = MagicMock() - mock_future.result.return_value = (make_metrics(Path("a.py")), None) + mock_future.result.return_value = (make_metrics(Path("a.py")), None, None) mock_exec = MagicMock() mock_exec.__enter__ = lambda s: s mock_exec.__exit__ = MagicMock(return_value=False) diff --git a/tests/test_performance_optimization.py b/tests/test_performance_optimization.py index 0b66afc..ca60500 100644 --- a/tests/test_performance_optimization.py +++ b/tests/test_performance_optimization.py @@ -242,11 +242,12 @@ def test_sequential_processing(self): def process_func(file_path): # Simulate processing - return None, None + return None, None, None - results, errors = processor.process_files(files, process_func) + results, errors, skips = processor.process_files(files, process_func) assert len(results) == 0 # All return None assert len(errors) == 0 + assert len(skips) == 0 class TestMemoryProfiler: diff --git a/tests/test_symbol_table_incremental.py b/tests/test_symbol_table_incremental.py new file mode 100644 index 0000000..1291f27 --- /dev/null +++ b/tests/test_symbol_table_incremental.py @@ -0,0 +1,122 @@ +import json +import time +from pathlib import Path +from refactron.analysis.symbol_table import SymbolTableBuilder, SymbolType + + +def test_symbol_table_incremental_build(tmp_path): + project_root = tmp_path / "project" + project_root.mkdir() + + file1 = project_root / "module1.py" + file1.write_text("def func1(): pass\nclass Class1: pass") + + cache_dir = tmp_path / "cache" + builder = SymbolTableBuilder(cache_dir=cache_dir) + + # 1. First build + table = builder.build_for_project(project_root) + assert "func1" in table.exports + assert "Class1" in table.exports + + cache_file = cache_dir / "symbols.json" + assert cache_file.exists() + + with open(cache_file, "r") as f: + cache_data = json.load(f) + assert file1.resolve().as_posix() in cache_data["file_metadata"] + + # 2. Second build (no change) + # We'll monkeypatch _analyze_file to verify it's not called + original_analyze = builder._analyze_file + analyze_called = [] + + def mocked_analyze(path): + analyze_called.append(path) + return original_analyze(path) + + builder._analyze_file = mocked_analyze + table2 = builder.build_for_project(project_root) + + assert len(analyze_called) == 0 + assert "func1" in table2.exports + + # 3. Modify file (incremental update) + time.sleep(0.1) # Ensure mtime changes + file1.write_text("def func1_v2(): pass\nclass Class1: pass") + + analyze_called.clear() + table3 = builder.build_for_project(project_root) + + assert len(analyze_called) == 1 + assert "func1_v2" in table3.exports + assert "func1" not in table3.exports + assert "Class1" in table3.exports + + # 4. Add new file + file2 = project_root / "module2.py" + file2.write_text("var2 = 42") + + analyze_called.clear() + table4 = builder.build_for_project(project_root) + + assert len(analyze_called) == 1 + assert file2.resolve().as_posix() in [p.as_posix() for p in analyze_called] + assert "var2" in table4.exports + + # 5. Delete file + file1.unlink() + + analyze_called.clear() + table5 = builder.build_for_project(project_root) + + assert len(analyze_called) == 0 + assert "func1_v2" not in table5.exports + assert "Class1" not in table5.exports + assert "var2" in table5.exports + assert file1.resolve().as_posix() not in table5.file_metadata + + +def test_symbol_table_hash_validation(tmp_path): + """Verify that content change triggers re-analysis even if mtime stays the same.""" + project_root = tmp_path / "project" + project_root.mkdir() + + file1 = project_root / "module1.py" + file1.write_text("x = 1") + + cache_dir = tmp_path / "cache" + builder = SymbolTableBuilder(cache_dir=cache_dir) + + # Initial build + builder.build_for_project(project_root) + original_mtime = file1.stat().st_mtime + original_size = file1.stat().st_size + + # Modify content but keep same size and restore mtime (simulated) + # Actually, hard to keep same size AND restore mtime exactly in some FS, + # but we can try. + file1.write_text("y = 2") # same size "x = 1" vs "y = 2" + import os + + os.utime(file1, (original_mtime, original_mtime)) + + # Verify mtime/size match but content hash differs + assert file1.stat().st_mtime == original_mtime + assert file1.stat().st_size == original_size + + analyze_called = [] + original_analyze = builder._analyze_file + + def mocked_analyze(path): + analyze_called.append(path) + return original_analyze(path) + + builder._analyze_file = mocked_analyze + + table = builder.build_for_project(project_root) + + # Should detect change via hash + assert len(analyze_called) == 1 + assert "y" in table.exports + assert "x" not in table.exports From 9897c90ee3cfe68d1e241064816b5d8edefa0731 Mon Sep 17 00:00:00 2001 From: shrutu0929 Date: Thu, 16 Apr 2026 22:15:11 +0530 Subject: [PATCH 2/2] fix: sync tests with RAG keyword-fallback and LLMOrchestrator routing --- refactron/rag/indexer.py | 16 ++- tests/test_cli_patterns_rag.py | 10 ++ tests/test_groq_client.py | 2 +- tests/test_rag_indexer.py | 54 +++++--- tests/test_rag_retriever.py | 218 +++++++++++++-------------------- 5 files changed, 148 insertions(+), 152 deletions(-) diff --git a/refactron/rag/indexer.py b/refactron/rag/indexer.py index 9fad604..d5d81a2 100644 --- a/refactron/rag/indexer.py +++ b/refactron/rag/indexer.py @@ -10,6 +10,12 @@ if TYPE_CHECKING: from refactron.llm.orchestrator import LLMOrchestrator +# LLMOrchestrator is imported lazily at runtime to keep it patchable in tests +try: + from refactron.llm.orchestrator import LLMOrchestrator as _LLMOrchestrator # noqa: F401 +except Exception: + _LLMOrchestrator = None # type: ignore + # RAG dependencies are loaded lazily in __init__ to prevent CLI crashes # if libraries like PyTorch fail to initialize (common on some Windows environments). CHROMA_AVAILABLE = None @@ -225,8 +231,14 @@ def _summarize_chunk(self, chunk: CodeChunk) -> Optional[str]: if not self.llm_integration: return None - # Delegate to the orchestration layer - return self.llm_integration.generate_chunk_summary(chunk.content) + try: + # Delegate to the orchestration layer + return self.llm_integration.generate_chunk_summary(chunk.content) + except Exception as e: + import logging + + logging.getLogger(__name__).warning(f"Chunk summarization failed: {e}") + return None def add_chunks(self, chunks: List[CodeChunk]) -> None: """Add code chunks to the vector index or keyword storage. diff --git a/tests/test_cli_patterns_rag.py b/tests/test_cli_patterns_rag.py index 8c7a788..5c2c110 100644 --- a/tests/test_cli_patterns_rag.py +++ b/tests/test_cli_patterns_rag.py @@ -163,6 +163,16 @@ def retrieve_similar(self, query, top_k=5, chunk_type=None): # noqa: ARG002 monkeypatch.setattr("refactron.cli.rag.RAGIndexer", _FakeIndexer) monkeypatch.setattr("refactron.cli.rag.ContextRetriever", _FakeRetriever) + # The `rag index` command routes through LLMOrchestrator.build_vector_index + class _FakeOrchestrator: + def __init__(self, *args, **kwargs): # noqa: ARG002 + pass + + def build_vector_index(self, *args, **kwargs): # noqa: ARG002 + pass + + monkeypatch.setattr("refactron.cli.rag.LLMOrchestrator", _FakeOrchestrator) + with runner.isolated_filesystem(temp_dir=tmp_path): assert runner.invoke(rag, ["index"]).exit_code == 0 assert runner.invoke(rag, ["search", "find function"]).exit_code == 0 diff --git a/tests/test_groq_client.py b/tests/test_groq_client.py index 17612d6..d37d3d2 100644 --- a/tests/test_groq_client.py +++ b/tests/test_groq_client.py @@ -34,7 +34,7 @@ def test_client_initialization_with_api_key(self, mock_groq_api): assert client.api_key == "test_key_123" assert client.model == "llama-3.3-70b-versatile" assert client.temperature == 0.2 - assert client.max_tokens == 2000 + assert client.max_tokens == 4000 def test_client_initialization_from_env(self, mock_groq_api): """Test client initialization from environment variable.""" diff --git a/tests/test_rag_indexer.py b/tests/test_rag_indexer.py index 43f3ebb..11293dd 100644 --- a/tests/test_rag_indexer.py +++ b/tests/test_rag_indexer.py @@ -13,19 +13,25 @@ def make_indexer(tmp_path: Path): """Return a RAGIndexer with all heavy dependencies mocked.""" + import refactron.rag.indexer as _indexer_mod + mock_embedding_model = MagicMock() mock_embedding_model.encode.return_value = MagicMock(tolist=lambda: [[0.1, 0.2, 0.3]]) mock_collection = MagicMock() mock_chroma_client = MagicMock() mock_chroma_client.get_or_create_collection.return_value = mock_collection + mock_chromadb = MagicMock() + mock_chromadb.PersistentClient.return_value = mock_chroma_client + + # SentenceTransformer and chromadb are injected lazily into the module's globals(). + # We must force-set them before RAGIndexer.__init__ runs so it finds them. + _indexer_mod.CHROMA_AVAILABLE = True + _indexer_mod.__dict__["SentenceTransformer"] = MagicMock(return_value=mock_embedding_model) + _indexer_mod.__dict__["chromadb"] = mock_chromadb + _indexer_mod.__dict__["Settings"] = MagicMock() - with patch("refactron.rag.indexer.CHROMA_AVAILABLE", True), patch( - "refactron.rag.indexer.SentenceTransformer", return_value=mock_embedding_model - ), patch("refactron.rag.indexer.chromadb") as mock_chromadb, patch( - "refactron.rag.indexer.Settings" - ): - mock_chromadb.PersistentClient.return_value = mock_chroma_client + try: from refactron.rag.indexer import RAGIndexer indexer = RAGIndexer(workspace_path=tmp_path) @@ -33,6 +39,12 @@ def make_indexer(tmp_path: Path): indexer._mock_collection = mock_collection indexer._mock_embedding = mock_embedding_model return indexer + finally: + # Reset so other tests start fresh + _indexer_mod.CHROMA_AVAILABLE = None + _indexer_mod.__dict__.pop("SentenceTransformer", None) + _indexer_mod.__dict__.pop("chromadb", None) + _indexer_mod.__dict__.pop("Settings", None) # ── tests for IndexStats ────────────────────────────────────────────────────── @@ -59,11 +71,17 @@ def test_fields(self): class TestRAGIndexerInit: def test_raises_without_chromadb(self, tmp_path): - with patch("refactron.rag.indexer.CHROMA_AVAILABLE", False): + """When ChromaDB is unavailable the indexer falls back to keyword mode (no exception).""" + import refactron.rag.indexer as _mod + + _mod.CHROMA_AVAILABLE = False # force keyword path + try: from refactron.rag.indexer import RAGIndexer - with pytest.raises(RuntimeError, match="ChromaDB"): - RAGIndexer(workspace_path=tmp_path) + indexer = RAGIndexer(workspace_path=tmp_path) + assert indexer.mode == "keyword" + finally: + _mod.CHROMA_AVAILABLE = None # reset for other tests def test_creates_index_dir(self, tmp_path): indexer = make_indexer(tmp_path) @@ -151,7 +169,7 @@ def test_chunk_metadata_list_serialised_to_json(self, tmp_path): class TestSummarizeChunk: def test_returns_none_without_llm(self, tmp_path): indexer = make_indexer(tmp_path) - indexer.llm_client = None + indexer.llm_integration = None chunk = MagicMock() result = indexer._summarize_chunk(chunk) assert result is None @@ -159,8 +177,8 @@ def test_returns_none_without_llm(self, tmp_path): def test_returns_summary_with_llm(self, tmp_path): indexer = make_indexer(tmp_path) mock_llm = MagicMock() - mock_llm.generate.return_value = "Handles authentication." - indexer.llm_client = mock_llm + mock_llm.generate_chunk_summary.return_value = "Handles authentication." + indexer.llm_integration = mock_llm chunk = MagicMock() chunk.content = "def authenticate(): pass" result = indexer._summarize_chunk(chunk) @@ -169,8 +187,8 @@ def test_returns_summary_with_llm(self, tmp_path): def test_llm_exception_returns_none(self, tmp_path): indexer = make_indexer(tmp_path) mock_llm = MagicMock() - mock_llm.generate.side_effect = RuntimeError("LLM error") - indexer.llm_client = mock_llm + mock_llm.generate_chunk_summary.side_effect = RuntimeError("LLM error") + indexer.llm_integration = mock_llm chunk = MagicMock() chunk.content = "x = 1" result = indexer._summarize_chunk(chunk) @@ -227,16 +245,16 @@ def test_index_file_error_handled(self, tmp_path): def test_summarize_flag_initializes_llm(self, tmp_path): indexer = make_indexer(tmp_path) - indexer.llm_client = None + indexer.llm_integration = None # No files → summarize path never reached for real files - with patch("refactron.rag.indexer.GroqClient", MagicMock()): + with patch("refactron.llm.orchestrator.LLMOrchestrator", MagicMock()): stats = indexer.index_repository(tmp_path, summarize=True) assert stats.total_files == 0 def test_summarize_flag_llm_init_failure(self, tmp_path): indexer = make_indexer(tmp_path) - indexer.llm_client = None - with patch("refactron.rag.indexer.GroqClient", side_effect=Exception("key error")): + indexer.llm_integration = None + with patch("refactron.llm.orchestrator.LLMOrchestrator", side_effect=Exception("key error")): stats = indexer.index_repository(tmp_path, summarize=True) assert stats.total_files == 0 diff --git a/tests/test_rag_retriever.py b/tests/test_rag_retriever.py index 85f5dc4..1bcc926 100644 --- a/tests/test_rag_retriever.py +++ b/tests/test_rag_retriever.py @@ -1,94 +1,103 @@ """Tests for the RAG retriever module.""" -import tempfile from pathlib import Path -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import MagicMock, Mock import pytest -# Create a comprehensive mock for transformers -transformers_mock = MagicMock() -transformers_mock.__path__ = [] +from refactron.rag.retriever import ContextRetriever, RetrievedContext -with patch.dict( - "sys.modules", - { - "transformers": transformers_mock, - "transformers.configuration_utils": MagicMock(), - "transformers.utils": MagicMock(), - "transformers.models": MagicMock(), - "transformers.file_utils": MagicMock(), - "transformers.tokenization_utils_base": MagicMock(), - }, -): - from refactron.rag.retriever import ContextRetriever, RetrievedContext + +# ── helper: build a fully-mocked ContextRetriever ──────────────────────────── + + +def make_retriever(workspace_path: Path) -> ContextRetriever: + """Return a ContextRetriever with all heavy deps mocked via globals injection.""" + import refactron.rag.retriever as _mod + + mock_model = Mock() + mock_model.encode.return_value.tolist.return_value = [[0.1, 0.2]] + + mock_collection = Mock() + mock_chroma_client = Mock() + mock_chroma_client.get_collection.return_value = mock_collection + mock_chromadb = MagicMock() + mock_chromadb.PersistentClient.return_value = mock_chroma_client + + _mod.CHROMA_AVAILABLE = True + _mod.__dict__["SentenceTransformer"] = MagicMock(return_value=mock_model) + _mod.__dict__["chromadb"] = mock_chromadb + _mod.__dict__["Settings"] = MagicMock() + + try: + r = ContextRetriever(workspace_path) + r._mock_collection = mock_collection + r._mock_model = mock_model + return r + finally: + _mod.CHROMA_AVAILABLE = None + _mod.__dict__.pop("SentenceTransformer", None) + _mod.__dict__.pop("chromadb", None) + _mod.__dict__.pop("Settings", None) + + +# ── tests ───────────────────────────────────────────────────────────────────── class TestContextRetriever: """Test cases for ContextRetriever.""" @pytest.fixture - def temp_workspace(self): + def temp_workspace(self, tmp_path): """Create a temporary workspace directory.""" - with tempfile.TemporaryDirectory() as tmpdir: - workspace_path = Path(tmpdir) - (workspace_path / ".rag").mkdir() - yield workspace_path - - @patch("refactron.rag.retriever.CHROMA_AVAILABLE", True) - @patch("refactron.rag.retriever.chromadb") - @patch("refactron.rag.retriever.Settings") - @patch("refactron.rag.retriever.SentenceTransformer") - def test_retriever_initialization( - self, mock_transformer, mock_settings, mock_chroma, temp_workspace - ): - """Test retriever initialization.""" - mock_collection = Mock() - mock_client = Mock() - mock_client.get_collection.return_value = mock_collection - mock_chroma.PersistentClient.return_value = mock_client - - mock_model = Mock() - mock_transformer.return_value = mock_model - - retriever = ContextRetriever(temp_workspace) + (tmp_path / ".rag").mkdir() + return tmp_path + def test_retriever_initialization(self, temp_workspace): + """Test retriever initialization.""" + retriever = make_retriever(temp_workspace) assert retriever.workspace_path == temp_workspace assert retriever.index_path == temp_workspace / ".rag" - @patch("refactron.rag.retriever.CHROMA_AVAILABLE", False) def test_retriever_requires_chromadb(self, temp_workspace): - """Test that retriever requires ChromaDB.""" - with pytest.raises(RuntimeError, match="ChromaDB is not available"): - ContextRetriever(temp_workspace) - - @patch("refactron.rag.retriever.CHROMA_AVAILABLE", True) - @patch("refactron.rag.retriever.chromadb") - @patch("refactron.rag.retriever.Settings") - @patch("refactron.rag.retriever.SentenceTransformer") - def test_retriever_missing_index( - self, mock_transformer, mock_settings, mock_chroma, temp_workspace - ): - """Test that missing index raises error.""" - mock_client = Mock() - mock_client.get_collection.side_effect = Exception("Collection not found") - mock_chroma.PersistentClient.return_value = mock_client + """When ChromaDB is unavailable the retriever falls back to keyword mode.""" + import refactron.rag.retriever as _mod - mock_model = Mock() - mock_transformer.return_value = mock_model + _mod.CHROMA_AVAILABLE = False + try: + retriever = ContextRetriever(temp_workspace) + assert retriever.mode == "keyword" + finally: + _mod.CHROMA_AVAILABLE = None - with pytest.raises(RuntimeError, match="Index not found"): - ContextRetriever(temp_workspace) + def test_retriever_missing_index(self, temp_workspace): + """Test that a missing collection (no keyword fallback file) raises RuntimeError.""" + import refactron.rag.retriever as _mod - @patch("refactron.rag.retriever.CHROMA_AVAILABLE", True) - @patch("refactron.rag.retriever.chromadb") - @patch("refactron.rag.retriever.Settings") - @patch("refactron.rag.retriever.SentenceTransformer") - def test_retrieve_similar(self, mock_transformer, mock_settings, mock_chroma, temp_workspace): + mock_model = Mock() + mock_chroma_client = Mock() + mock_chroma_client.get_collection.side_effect = Exception("Collection not found") + mock_chromadb = MagicMock() + mock_chromadb.PersistentClient.return_value = mock_chroma_client + + _mod.CHROMA_AVAILABLE = True + _mod.__dict__["SentenceTransformer"] = MagicMock(return_value=mock_model) + _mod.__dict__["chromadb"] = mock_chromadb + _mod.__dict__["Settings"] = MagicMock() + try: + # No keyword_chunks.json present → should raise RuntimeError("Index not found") + with pytest.raises(RuntimeError, match="Index not found"): + ContextRetriever(temp_workspace) + finally: + _mod.CHROMA_AVAILABLE = None + _mod.__dict__.pop("SentenceTransformer", None) + _mod.__dict__.pop("chromadb", None) + _mod.__dict__.pop("Settings", None) + + def test_retrieve_similar(self, temp_workspace): """Test retrieving similar code chunks.""" - # Setup mock collection with results - mock_collection = Mock() - mock_collection.query.return_value = { + retriever = make_retriever(temp_workspace) + retriever._mock_collection.query.return_value = { "documents": [["def test(): pass"]], "metadatas": [ [ @@ -103,16 +112,6 @@ def test_retrieve_similar(self, mock_transformer, mock_settings, mock_chroma, te ], "distances": [[0.15]], } - - mock_client = Mock() - mock_client.get_collection.return_value = mock_collection - mock_chroma.PersistentClient.return_value = mock_client - - mock_model = Mock() - mock_model.encode.return_value.tolist.return_value = [[0.1, 0.2]] - mock_transformer.return_value = mock_model - - retriever = ContextRetriever(temp_workspace) results = retriever.retrieve_similar("test function", top_k=1) assert len(results) == 1 @@ -120,14 +119,10 @@ def test_retrieve_similar(self, mock_transformer, mock_settings, mock_chroma, te assert results[0].name == "test" assert results[0].chunk_type == "function" - @patch("refactron.rag.retriever.CHROMA_AVAILABLE", True) - @patch("refactron.rag.retriever.chromadb") - @patch("refactron.rag.retriever.Settings") - @patch("refactron.rag.retriever.SentenceTransformer") - def test_retrieve_by_file(self, mock_transformer, mock_settings, mock_chroma, temp_workspace): + def test_retrieve_by_file(self, temp_workspace): """Test retrieving chunks by file path.""" - mock_collection = Mock() - mock_collection.get.return_value = { + retriever = make_retriever(temp_workspace) + retriever._mock_collection.get.return_value = { "documents": ["def test(): pass"], "metadatas": [ { @@ -139,72 +134,33 @@ def test_retrieve_by_file(self, mock_transformer, mock_settings, mock_chroma, te } ], } - - mock_client = Mock() - mock_client.get_collection.return_value = mock_collection - mock_chroma.PersistentClient.return_value = mock_client - - mock_model = Mock() - mock_transformer.return_value = mock_model - - retriever = ContextRetriever(temp_workspace) results = retriever.retrieve_by_file("/test.py") assert len(results) == 1 assert results[0].file_path == "/test.py" - @patch("refactron.rag.retriever.CHROMA_AVAILABLE", True) - @patch("refactron.rag.retriever.chromadb") - @patch("refactron.rag.retriever.Settings") - @patch("refactron.rag.retriever.SentenceTransformer") - def test_retrieve_functions(self, mock_transformer, mock_settings, mock_chroma, temp_workspace): + def test_retrieve_functions(self, temp_workspace): """Test retrieving only function chunks.""" - mock_collection = Mock() - mock_collection.query.return_value = { + retriever = make_retriever(temp_workspace) + retriever._mock_collection.query.return_value = { "documents": [[]], "metadatas": [[]], "distances": [[]], } - - mock_client = Mock() - mock_client.get_collection.return_value = mock_collection - mock_chroma.PersistentClient.return_value = mock_client - - mock_model = Mock() - mock_model.encode.return_value.tolist.return_value = [[0.1, 0.2]] - mock_transformer.return_value = mock_model - - retriever = ContextRetriever(temp_workspace) retriever.retrieve_functions("test", top_k=5) # Verify that chunk_type filter was used - call_args = mock_collection.query.call_args + call_args = retriever._mock_collection.query.call_args assert call_args.kwargs.get("where") == {"chunk_type": "function"} - @patch("refactron.rag.retriever.CHROMA_AVAILABLE", True) - @patch("refactron.rag.retriever.chromadb") - @patch("refactron.rag.retriever.Settings") - @patch("refactron.rag.retriever.SentenceTransformer") - def test_retrieve_no_results( - self, mock_transformer, mock_settings, mock_chroma, temp_workspace - ): + def test_retrieve_no_results(self, temp_workspace): """Test retrieval with no results.""" - mock_collection = Mock() - mock_collection.query.return_value = { + retriever = make_retriever(temp_workspace) + retriever._mock_collection.query.return_value = { "documents": [[]], "metadatas": [[]], "distances": [[]], } - - mock_client = Mock() - mock_client.get_collection.return_value = mock_collection - mock_chroma.PersistentClient.return_value = mock_client - - mock_model = Mock() - mock_model.encode.return_value.tolist.return_value = [[0.1, 0.2]] - mock_transformer.return_value = mock_model - - retriever = ContextRetriever(temp_workspace) results = retriever.retrieve_similar("nonexistent", top_k=5) assert len(results) == 0