ai-coding-assistant 0.5.0__py3-none-any.whl

coding_assistant/operations/generator.py

@@ -0,0 +1,347 @@
+"""Code generation with extraction and validation."""
+
+import re
+from typing import List, Dict, Optional, Tuple
+from dataclasses import dataclass
+from coding_assistant.operations.validator import SyntaxValidator, ValidationResult
+
+
+@dataclass
+class CodeBlock:
+    """Represents a code block extracted from text."""
+    language: str
+    code: str
+    start_line: int
+    end_line: int
+    file_path: Optional[str] = None
+    is_valid: bool = False
+    validation_error: Optional[str] = None
+
+    def __repr__(self):
+        status = "✓ valid" if self.is_valid else f"✗ invalid: {self.validation_error}"
+        return f"CodeBlock({self.language}, lines {self.start_line}-{self.end_line}, {status})"
+
+
+class CodeGenerator:
+    """Generate and validate code from LLM responses."""
+
+    def __init__(self):
+        """Initialize the code generator."""
+        self.validator = SyntaxValidator()
+
+    def extract_code_blocks(self, text: str, validate: bool = True) -> List[CodeBlock]:
+        """
+        Extract code blocks from markdown-formatted text.
+
+        Supports fenced code blocks with language specifiers:
+        ```python
+        code here
+        ```
+
+        Args:
+            text: The text containing code blocks (typically LLM response)
+            validate: Whether to validate extracted code blocks
+
+        Returns:
+            List of CodeBlock objects
+        """
+        code_blocks = []
+
+        # Pattern for fenced code blocks: ```language\ncode\n```
+        pattern = r'```(\w+)?\n(.*?)```'
+
+        lines = text.split('\n')
+        current_line = 0
+
+        # Find all code blocks
+        for match in re.finditer(pattern, text, re.DOTALL):
+            language = match.group(1) or 'text'
+            code = match.group(2)
+
+            # Calculate line numbers
+            start_pos = match.start()
+            start_line = text[:start_pos].count('\n') + 1
+            end_line = start_line + code.count('\n')
+
+            # Skip non-code languages
+            if language.lower() in ('text', 'plaintext', 'output', 'bash', 'shell', 'sh'):
+                continue
+
+            # Normalize language names
+            language = self._normalize_language(language)
+
+            # Create code block
+            block = CodeBlock(
+                language=language,
+                code=code.strip(),
+                start_line=start_line,
+                end_line=end_line
+            )
+
+            # Validate if requested
+            if validate and language in self.validator.SUPPORTED_LANGUAGES:
+                validation_result = self.validator.validate(block.code, language)
+                block.is_valid = validation_result.is_valid
+                if not validation_result.is_valid:
+                    block.validation_error = validation_result.error_message
+
+            code_blocks.append(block)
+
+        return code_blocks
+
+    def extract_with_context(self, text: str) -> List[Dict]:
+        """
+        Extract code blocks with surrounding context.
+
+        Returns code blocks along with explanatory text before/after them.
+
+        Args:
+            text: The text containing code and explanations
+
+        Returns:
+            List of dicts with 'type' (code/text), 'content', and metadata
+        """
+        sections = []
+
+        # Split by code blocks but keep them
+        pattern = r'(```\w*\n.*?```)'
+        parts = re.split(pattern, text, flags=re.DOTALL)
+
+        for part in parts:
+            if part.strip():
+                if part.startswith('```'):
+                    # It's a code block
+                    blocks = self.extract_code_blocks(part)
+                    for block in blocks:
+                        sections.append({
+                            'type': 'code',
+                            'language': block.language,
+                            'content': block.code,
+                            'is_valid': block.is_valid,
+                            'validation_error': block.validation_error
+                        })
+                else:
+                    # It's text/explanation
+                    sections.append({
+                        'type': 'text',
+                        'content': part.strip()
+                    })
+
+        return sections
+
+    def generate_from_specification(self, specification: str,
+                                    context: Optional[List[Dict]] = None,
+                                    llm_client = None) -> List[CodeBlock]:
+        """
+        Generate code from a specification using an LLM.
+
+        Args:
+            specification: Description of what code to generate
+            context: Optional context (similar code, documentation, etc.)
+            llm_client: LLM client to use for generation
+
+        Returns:
+            List of generated and validated CodeBlock objects
+
+        Raises:
+            ValueError: If no LLM client provided
+        """
+        if not llm_client:
+            raise ValueError("LLM client required for code generation")
+
+        # Build prompt
+        prompt = self._build_generation_prompt(specification, context)
+
+        # Generate code
+        response = llm_client.generate(prompt, stream=False)
+
+        # Extract and validate code blocks
+        code_blocks = self.extract_code_blocks(response, validate=True)
+
+        return code_blocks
+
+    def validate_code_block(self, block: CodeBlock) -> ValidationResult:
+        """
+        Validate a code block.
+
+        Args:
+            block: The CodeBlock to validate
+
+        Returns:
+            ValidationResult
+        """
+        return self.validator.validate(block.code, block.language, block.file_path)
+
+    def attempt_fix(self, block: CodeBlock, llm_client=None) -> Optional[CodeBlock]:
+        """
+        Attempt to fix invalid code using LLM.
+
+        Args:
+            block: The invalid CodeBlock to fix
+            llm_client: LLM client for generating fixes
+
+        Returns:
+            Fixed CodeBlock if successful, None otherwise
+        """
+        if block.is_valid:
+            return block
+
+        if not llm_client:
+            return None
+
+        # Build fix prompt
+        prompt = f"""The following {block.language} code has a syntax error:
+
+```{block.language}
+{block.code}
+```
+
+Error: {block.validation_error}
+
+Please fix the syntax error and provide only the corrected code in a code block.
+"""
+
+        try:
+            response = llm_client.generate(prompt, stream=False)
+            fixed_blocks = self.extract_code_blocks(response, validate=True)
+
+            if fixed_blocks and fixed_blocks[0].is_valid:
+                return fixed_blocks[0]
+        except Exception:
+            pass
+
+        return None
+
+    def _normalize_language(self, language: str) -> str:
+        """Normalize language identifiers."""
+        language = language.lower()
+
+        # Map common aliases
+        alias_map = {
+            'py': 'python',
+            'js': 'javascript',
+            'ts': 'typescript',
+            'jsx': 'jsx',
+            'tsx': 'tsx',
+            'javascript': 'javascript',
+            'typescript': 'typescript',
+            'python': 'python',
+        }
+
+        return alias_map.get(language, language)
+
+    def _build_generation_prompt(self, specification: str,
+                                 context: Optional[List[Dict]] = None) -> str:
+        """Build a prompt for code generation."""
+        prompt = "Generate code based on the following specification:\n\n"
+        prompt += f"{specification}\n\n"
+
+        if context:
+            prompt += "Context (related code from the project):\n\n"
+            for ctx in context:
+                if 'file_path' in ctx:
+                    prompt += f"## {ctx['file_path']}\n"
+                prompt += f"```{ctx.get('language', 'python')}\n"
+                prompt += f"{ctx.get('content', '')}\n"
+                prompt += "```\n\n"
+
+        prompt += "\nProvide only the code in properly formatted code blocks with language specifiers."
+
+        return prompt
+
+    def extract_file_operations(self, text: str) -> List[Dict]:
+        """
+        Extract file operations from LLM response.
+
+        Looks for patterns like:
+        - "Create file: path/to/file.py"
+        - "Modify file: path/to/file.py"
+        - "Delete file: path/to/file.py"
+
+        Returns:
+            List of file operations with type, path, and code
+        """
+        operations = []
+
+        # Pattern for file operations
+        file_op_pattern = r'(Create|Modify|Update|Delete)\s+(?:file:?\s+)?`?([^\s`]+\.(?:py|js|ts|jsx|tsx))`?'
+
+        # Find all file operations
+        matches = re.finditer(file_op_pattern, text, re.IGNORECASE)
+
+        # Also extract code blocks
+        code_blocks = self.extract_code_blocks(text, validate=True)
+        code_block_idx = 0
+
+        for match in matches:
+            operation_type = match.group(1).lower()
+            file_path = match.group(2)
+
+            # Map operation types
+            if operation_type in ('create', 'modify', 'update'):
+                # Try to find associated code block
+                code = None
+                if code_block_idx < len(code_blocks):
+                    code = code_blocks[code_block_idx].code
+                    code_block_idx += 1
+
+                operations.append({
+                    'type': 'create' if operation_type == 'create' else 'modify',
+                    'file_path': file_path,
+                    'code': code,
+                    'validated': code_blocks[code_block_idx - 1].is_valid if code else False
+                })
+            elif operation_type == 'delete':
+                operations.append({
+                    'type': 'delete',
+                    'file_path': file_path
+                })
+
+        return operations
+
+    def format_code(self, code: str, language: str) -> str:
+        """
+        Format code according to language conventions.
+
+        Args:
+            code: The code to format
+            language: The programming language
+
+        Returns:
+            Formatted code
+        """
+        # For Python, could use black
+        # For JS/TS, could use prettier
+        # For now, just return as-is
+        # This is a placeholder for future formatting integration
+        return code
+
+    def count_lines(self, code: str) -> int:
+        """Count non-empty lines in code."""
+        return len([line for line in code.split('\n') if line.strip()])
+
+    def estimate_complexity(self, code: str, language: str) -> Dict:
+        """
+        Estimate code complexity.
+
+        Returns:
+            Dict with metrics like lines, functions, classes, etc.
+        """
+        metrics = {
+            'lines': self.count_lines(code),
+            'total_lines': len(code.split('\n')),
+            'language': language
+        }
+
+        if language == 'python':
+            # Count Python-specific constructs
+            metrics['functions'] = code.count('def ')
+            metrics['classes'] = code.count('class ')
+            metrics['imports'] = code.count('import ') + code.count('from ')
+        elif language in ('javascript', 'typescript', 'jsx', 'tsx'):
+            # Count JS/TS-specific constructs
+            metrics['functions'] = code.count('function ') + code.count('=> ')
+            metrics['classes'] = code.count('class ')
+            metrics['imports'] = code.count('import ') + code.count('require(')
+
+        return metrics