cognify-code 0.2.0__py3-none-any.whl

ai_code_assistant/refactor/multi_file_editor.py

@@ -0,0 +1,346 @@
+"""Multi-file editor for coordinated refactoring operations."""
+
+import re
+import shutil
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+from ai_code_assistant.config import Config, get_language_by_extension
+from ai_code_assistant.llm import LLMManager
+from ai_code_assistant.editor.diff_handler import DiffHandler
+from ai_code_assistant.refactor.prompts import (
+    MULTI_FILE_EDIT_PROMPT,
+    RENAME_SYMBOL_PROMPT,
+    EXTRACT_TO_FILE_PROMPT,
+)
+from ai_code_assistant.refactor.change_plan import (
+    ChangePlan,
+    FileChange,
+    ChangeType,
+    RefactorResult,
+)
+from ai_code_assistant.refactor.analyzer import RefactorAnalyzer
+
+
+class MultiFileEditor:
+    """Editor for coordinated multi-file refactoring."""
+    
+    def __init__(self, config: Config, llm_manager: LLMManager):
+        """Initialize the multi-file editor.
+        
+        Args:
+            config: Application configuration
+            llm_manager: LLM manager for AI interactions
+        """
+        self.config = config
+        self.llm = llm_manager
+        self.diff_handler = DiffHandler()
+        self.analyzer = RefactorAnalyzer(config, llm_manager)
+    
+    def refactor(
+        self,
+        instruction: str,
+        files: List[Path],
+        dry_run: bool = False,
+        create_backup: bool = True,
+    ) -> RefactorResult:
+        """Perform multi-file refactoring.
+        
+        Args:
+            instruction: Refactoring instruction
+            files: List of files to potentially modify
+            dry_run: If True, don't apply changes
+            create_backup: If True, create backup before changes
+            
+        Returns:
+            RefactorResult with all changes
+        """
+        # Analyze to create plan
+        plan = self.analyzer.analyze(instruction, files)
+        
+        if not plan.changes:
+            return RefactorResult(
+                plan=plan,
+                error="No changes identified",
+            )
+        
+        # Read current file contents
+        file_contents = self._read_files([Path(c.file_path) for c in plan.changes])
+        
+        # Update plan with original content
+        for change in plan.changes:
+            if change.file_path in file_contents:
+                change.original_content = file_contents[change.file_path]
+        
+        # Generate changes
+        result = self._generate_changes(plan, file_contents)
+        
+        if result.error:
+            return result
+        
+        # Apply changes if not dry run
+        if not dry_run:
+            backup_dir = None
+            if create_backup:
+                backup_dir = self._create_backup(plan)
+                result.backup_dir = str(backup_dir) if backup_dir else None
+            
+            self._apply_changes(plan)
+            result.applied = True
+        
+        return result
+    
+    def rename_symbol(
+        self,
+        old_name: str,
+        new_name: str,
+        symbol_type: str,
+        files: List[Path],
+        dry_run: bool = False,
+    ) -> RefactorResult:
+        """Rename a symbol across multiple files.
+        
+        Args:
+            old_name: Current symbol name
+            new_name: New symbol name
+            symbol_type: Type of symbol (function, class, variable, etc.)
+            files: Files to search and modify
+            dry_run: If True, don't apply changes
+            
+        Returns:
+            RefactorResult with all changes
+        """
+        # Read file contents
+        file_contents = self._read_files(files)
+        
+        # Filter to files containing the symbol
+        relevant_files = {
+            path: content for path, content in file_contents.items()
+            if old_name in content
+        }
+        
+        if not relevant_files:
+            plan = ChangePlan(
+                instruction=f"Rename {symbol_type} '{old_name}' to '{new_name}'",
+                summary=f"Symbol '{old_name}' not found in any files",
+                changes=[],
+            )
+            return RefactorResult(plan=plan, error="Symbol not found")
+        
+        # Create plan
+        plan = ChangePlan(
+            instruction=f"Rename {symbol_type} '{old_name}' to '{new_name}'",
+            summary=f"Renaming {symbol_type} across {len(relevant_files)} files",
+            changes=[
+                FileChange(
+                    file_path=path,
+                    change_type=ChangeType.MODIFY,
+                    description=f"Rename '{old_name}' to '{new_name}'",
+                    original_content=content,
+                )
+                for path, content in relevant_files.items()
+            ],
+        )
+        
+        # Generate changes using LLM
+        formatted_contents = self._format_file_contents(relevant_files)
+        
+        try:
+            response = self.llm.invoke_with_template(
+                RENAME_SYMBOL_PROMPT,
+                old_name=old_name,
+                new_name=new_name,
+                symbol_type=symbol_type,
+                file_contents=formatted_contents,
+            )
+            
+            self._parse_multi_file_response(plan, response)
+            
+        except Exception as e:
+            return RefactorResult(plan=plan, error=f"LLM error: {str(e)}")
+        
+        # Apply if not dry run
+        if not dry_run:
+            self._apply_changes(plan)
+            return RefactorResult(plan=plan, applied=True)
+        
+        return RefactorResult(plan=plan)
+    
+    def _read_files(self, files: List[Path]) -> Dict[str, str]:
+        """Read contents of files."""
+        contents = {}
+        for file_path in files:
+            if file_path.exists():
+                try:
+                    contents[str(file_path)] = file_path.read_text(encoding="utf-8")
+                except Exception:
+                    pass
+        return contents
+    
+    def _format_file_contents(self, contents: Dict[str, str]) -> str:
+        """Format file contents for prompt."""
+        parts = []
+        for file_path, content in contents.items():
+            lang = get_language_by_extension(self.config, Path(file_path)) or "text"
+            parts.append(f"### {file_path}\n```{lang}\n{content}\n```\n")
+        return "\n".join(parts)
+
+    def _generate_changes(
+        self,
+        plan: ChangePlan,
+        file_contents: Dict[str, str],
+    ) -> RefactorResult:
+        """Generate actual file changes using LLM.
+
+        Args:
+            plan: The change plan
+            file_contents: Current file contents
+
+        Returns:
+            RefactorResult with generated changes
+        """
+        # Format change plan for prompt
+        change_plan_text = "\n".join([
+            f"- {c.file_path}: {c.description} ({c.change_type.value})"
+            for c in plan.changes
+        ])
+
+        formatted_contents = self._format_file_contents(file_contents)
+
+        try:
+            response = self.llm.invoke_with_template(
+                MULTI_FILE_EDIT_PROMPT,
+                instruction=plan.instruction,
+                change_plan=change_plan_text,
+                file_contents=formatted_contents,
+            )
+
+            self._parse_multi_file_response(plan, response)
+
+            # Generate diffs for each change
+            for change in plan.changes:
+                if change.new_content and change.original_content:
+                    change.diff = self.diff_handler.generate_diff(
+                        original=change.original_content,
+                        modified=change.new_content,
+                        filename=change.file_path,
+                    )
+
+            return RefactorResult(plan=plan)
+
+        except Exception as e:
+            return RefactorResult(plan=plan, error=f"LLM error: {str(e)}")
+
+    def _parse_multi_file_response(self, plan: ChangePlan, response: str) -> None:
+        """Parse LLM response containing multiple file contents.
+
+        Args:
+            plan: Change plan to update
+            response: LLM response with file contents
+        """
+        # Pattern to match file sections
+        pattern = r"###\s*FILE:\s*([^\n]+)\s*\n```\w*\s*\n(.*?)```"
+        matches = re.findall(pattern, response, re.DOTALL)
+
+        # Map new content to changes
+        for file_path, content in matches:
+            file_path = file_path.strip()
+            content = content.strip()
+
+            # Find matching change in plan
+            for change in plan.changes:
+                if change.file_path == file_path or file_path.endswith(change.file_path):
+                    change.new_content = content
+                    break
+
+    def _create_backup(self, plan: ChangePlan) -> Optional[Path]:
+        """Create backup of all files to be modified.
+
+        Args:
+            plan: Change plan with files to backup
+
+        Returns:
+            Path to backup directory or None
+        """
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        backup_dir = Path(f".refactor_backup_{timestamp}")
+
+        try:
+            backup_dir.mkdir(exist_ok=True)
+
+            for change in plan.changes:
+                if change.change_type in (ChangeType.MODIFY, ChangeType.DELETE):
+                    src = Path(change.file_path)
+                    if src.exists():
+                        # Preserve directory structure in backup
+                        dest = backup_dir / change.file_path
+                        dest.parent.mkdir(parents=True, exist_ok=True)
+                        shutil.copy2(src, dest)
+
+            return backup_dir
+
+        except Exception:
+            return None
+
+    def _apply_changes(self, plan: ChangePlan) -> None:
+        """Apply all changes in the plan.
+
+        Args:
+            plan: Change plan with changes to apply
+        """
+        for change in plan.get_ordered_changes():
+            try:
+                path = Path(change.file_path)
+
+                if change.change_type == ChangeType.DELETE:
+                    if path.exists():
+                        path.unlink()
+                    change.applied = True
+
+                elif change.change_type == ChangeType.CREATE:
+                    path.parent.mkdir(parents=True, exist_ok=True)
+                    path.write_text(change.new_content, encoding="utf-8")
+                    change.applied = True
+
+                elif change.change_type == ChangeType.RENAME:
+                    if path.exists() and change.new_path:
+                        new_path = Path(change.new_path)
+                        new_path.parent.mkdir(parents=True, exist_ok=True)
+                        shutil.move(str(path), str(new_path))
+                    change.applied = True
+
+                elif change.change_type == ChangeType.MODIFY:
+                    if change.new_content:
+                        path.write_text(change.new_content, encoding="utf-8")
+                        change.applied = True
+                    else:
+                        change.error = "No new content generated"
+
+            except Exception as e:
+                change.error = str(e)
+
+    def restore_backup(self, backup_dir: Path) -> bool:
+        """Restore files from backup.
+
+        Args:
+            backup_dir: Path to backup directory
+
+        Returns:
+            True if successful
+        """
+        if not backup_dir.exists():
+            return False
+
+        try:
+            for backup_file in backup_dir.rglob("*"):
+                if backup_file.is_file():
+                    # Get relative path from backup dir
+                    rel_path = backup_file.relative_to(backup_dir)
+                    dest = Path(rel_path)
+                    dest.parent.mkdir(parents=True, exist_ok=True)
+                    shutil.copy2(backup_file, dest)
+            return True
+        except Exception:
+            return False
+

ai_code_assistant/refactor/prompts.py

@@ -0,0 +1,175 @@
+"""Prompt templates for multi-file refactoring."""
+
+from langchain_core.prompts import ChatPromptTemplate
+
+# System prompt for analyzing refactoring scope
+ANALYZE_REFACTOR_SYSTEM = """You are an expert code architect analyzing a codebase for refactoring.
+
+Your task is to analyze the provided code files and determine what changes are needed to accomplish the refactoring goal.
+
+For each file that needs changes, provide:
+1. The file path
+2. The type of change (modify, create, delete, rename)
+3. A description of what changes are needed
+4. The priority (high, medium, low)
+5. Dependencies on other file changes
+
+Return your analysis as JSON with this structure:
+{
+    "summary": "Brief description of the refactoring plan",
+    "affected_files": [
+        {
+            "file_path": "path/to/file.py",
+            "change_type": "modify|create|delete|rename",
+            "description": "What changes are needed",
+            "priority": "high|medium|low",
+            "depends_on": ["other/file.py"]
+        }
+    ],
+    "risks": ["List of potential risks or breaking changes"],
+    "estimated_complexity": "low|medium|high"
+}"""
+
+ANALYZE_REFACTOR_PROMPT = ChatPromptTemplate.from_messages([
+    ("system", ANALYZE_REFACTOR_SYSTEM),
+    ("human", """Analyze the following codebase for this refactoring task:
+
+**Refactoring Goal:** {instruction}
+
+**Files in scope:**
+{file_contents}
+
+Analyze what changes are needed and return a JSON plan.""")
+])
+
+# System prompt for generating multi-file changes
+MULTI_FILE_EDIT_SYSTEM = """You are an expert code refactoring assistant performing coordinated changes across multiple files.
+
+When making changes:
+1. Ensure consistency across all files
+2. Update all imports and references
+3. Maintain backward compatibility where possible
+4. Follow existing code style and patterns
+5. Add appropriate type hints and docstrings
+
+For each file, return the COMPLETE modified content wrapped in a code block with the filename as a comment."""
+
+MULTI_FILE_EDIT_PROMPT = ChatPromptTemplate.from_messages([
+    ("system", MULTI_FILE_EDIT_SYSTEM),
+    ("human", """Apply the following refactoring across these files:
+
+**Refactoring Goal:** {instruction}
+
+**Change Plan:**
+{change_plan}
+
+**Current File Contents:**
+{file_contents}
+
+For each file that needs changes, return the complete modified content in this format:
+
+### FILE: path/to/file.py
+```python
+# complete file content here
+```
+
+### FILE: path/to/another.py
+```javascript
+// complete file content here
+```
+
+Return ALL files that need changes with their complete content.""")
+])
+
+# Prompt for renaming symbols across files
+RENAME_SYMBOL_PROMPT = ChatPromptTemplate.from_messages([
+    ("system", """You are an expert at renaming symbols across a codebase.
+
+When renaming:
+1. Update all occurrences of the symbol
+2. Update imports and exports
+3. Update docstrings and comments that reference the symbol
+4. Preserve all other code exactly as-is
+
+Return each modified file with its complete content."""),
+    ("human", """Rename the symbol across these files:
+
+**Old Name:** {old_name}
+**New Name:** {new_name}
+**Symbol Type:** {symbol_type}
+
+**Files:**
+{file_contents}
+
+Return each modified file with complete content in this format:
+
+### FILE: path/to/file.py
+```python
+# complete file content
+```""")
+])
+
+# Prompt for extracting code to new file
+EXTRACT_TO_FILE_PROMPT = ChatPromptTemplate.from_messages([
+    ("system", """You are an expert at extracting code into separate modules.
+
+When extracting:
+1. Move the specified code to a new file
+2. Add appropriate imports to the new file
+3. Update the original file to import from the new location
+4. Ensure all references are updated
+
+Return both the new file and modified original file."""),
+    ("human", """Extract code to a new file:
+
+**What to Extract:** {instruction}
+**New File Path:** {new_file_path}
+
+**Original File ({original_file}):**
+```{language}
+{original_content}
+```
+
+Return both files with complete content:
+
+### FILE: {new_file_path}
+```{language}
+# new file content
+```
+
+### FILE: {original_file}
+```{language}
+# modified original file
+```""")
+])
+
+# Prompt for adding feature across multiple files
+ADD_FEATURE_MULTI_PROMPT = ChatPromptTemplate.from_messages([
+    ("system", """You are an expert at adding features that span multiple files.
+
+When adding features:
+1. Follow existing patterns in the codebase
+2. Add necessary imports
+3. Update configuration if needed
+4. Add appropriate tests structure
+5. Maintain consistency with existing code style
+
+Return all files that need to be created or modified."""),
+    ("human", """Add the following feature across the codebase:
+
+**Feature Description:** {instruction}
+
+**Existing Files:**
+{file_contents}
+
+**Files to Create/Modify:**
+{target_files}
+
+Return all files with complete content in this format:
+
+### FILE: path/to/file.py
+```python
+# complete file content
+```""")
+])
+

ai_code_assistant/retrieval/__init__.py

@@ -0,0 +1,19 @@
+"""
+Codebase retrieval module for semantic search.
+
+This module provides functionality to:
+- Index code files into a vector database
+- Search for relevant code using natural language queries
+- Watch for file changes and update the index
+"""
+
+from .indexer import CodebaseIndexer
+from .search import CodebaseSearch
+from .chunker import CodeChunker
+
+__all__ = [
+    "CodebaseIndexer",
+    "CodebaseSearch", 
+    "CodeChunker",
+]
+