algomath-extract 1.0.0

package/src/generation/llm_generator.py

@@ -0,0 +1,223 @@
+"""LLM-based code generator.
+
+This module provides the LLMCodeGenerator class for generating
+Python code using LLM assistance for complex expressions.
+"""
+
+import ast
+import re
+from typing import Any, Dict, List, Optional
+
+from src.extraction.schema import Algorithm
+from src.generation.code_generator import GeneratedCode
+from src.generation.errors import LLMGenerationError
+from src.generation.prompts import (
+    format_code_generation_prompt,
+    format_complex_expression_prompt,
+)
+from src.generation.types import TypeInferrer, ValidationResult
+
+
+class LLMCodeGenerator:
+    """
+    Generate code using LLM for complex constructs.
+    
+    Falls back to template generation when LLM is unavailable
+    or for simple constructs.
+    """
+    
+    def __init__(self,
+                 model: str = "claude-3-5-haiku",
+                 temperature: float = 0.2,
+                 timeout: int = 30):
+        """
+        Initialize the LLM code generator.
+        
+        Args:
+            model: LLM model to use
+            temperature: Sampling temperature
+            timeout: Request timeout in seconds
+        """
+        self.model = model
+        self.temperature = temperature
+        self.timeout = timeout
+        self.type_inferrer = TypeInferrer()
+    
+    def generate(self, algorithm: Algorithm) -> GeneratedCode:
+        """
+        Generate code using LLM for complex constructs.
+        
+        Uses hybrid approach:
+        - LLM for complex expressions templates can't handle
+        - Type inference for type hints
+        - AST validation for syntax checking
+        
+        Args:
+            algorithm: Algorithm to generate code for
+            
+        Returns:
+            GeneratedCode with source and metadata
+            
+        Raises:
+            LLMGenerationError: If LLM generation fails
+        """
+        try:
+            # Format prompt for LLM
+            messages = format_code_generation_prompt(algorithm)
+            
+            # Call LLM (simulated - would use actual LLM API)
+            code = self._call_llm(messages)
+            
+            if not code:
+                raise LLMGenerationError("LLM returned empty response")
+            
+            # Validate syntax
+            validation = self._validate_syntax(code)
+            
+            return GeneratedCode(
+                source=code,
+                algorithm_name=algorithm.name,
+                imports=[],
+                validation_result=validation
+            )
+            
+        except Exception as e:
+            raise LLMGenerationError(f"LLM generation failed: {e}")
+    
+    def generate_docstring(self, algorithm: Algorithm) -> str:
+        """
+        Generate comprehensive docstring via LLM.
+        
+        Args:
+            algorithm: Algorithm to generate docstring for
+            
+        Returns:
+            Formatted docstring
+        """
+        try:
+            from src.generation.prompts import format_docstring_generation_prompt
+            messages = format_docstring_generation_prompt(algorithm)
+            docstring = self._call_llm(messages)
+            return docstring or self._fallback_docstring(algorithm)
+        except Exception:
+            return self._fallback_docstring(algorithm)
+    
+    def generate_complex_expression(self,
+                                    expression: str,
+                                    context: Dict[str, Any] = None) -> str:
+        """
+        Generate Python code for complex mathematical expression.
+        
+        Args:
+            expression: Mathematical expression text
+            context: Variable context (available vars, types)
+            
+        Returns:
+            Python code string
+        """
+        messages = format_complex_expression_prompt(expression, context)
+        code = self._call_llm(messages)
+        return code or expression
+    
+    def _call_llm(self, messages: List[Dict[str, str]]) -> Optional[str]:
+        """
+        Call LLM API with messages.
+        
+        This is a stub - actual implementation would use
+        Anthropic, OpenAI, or other LLM API.
+        
+        Args:
+            messages: List of message dicts
+            
+        Returns:
+            LLM response text or None
+        """
+        # TODO: Implement actual LLM API call
+        # For now, return None to trigger fallback
+        return None
+    
+    def _validate_syntax(self, code: str) -> ValidationResult:
+        """Validate Python syntax."""
+        try:
+            ast.parse(code)
+            return ValidationResult(is_valid=True, errors=[], warnings=[])
+        except SyntaxError as e:
+            return ValidationResult(
+                is_valid=False,
+                errors=[{
+                    'message': f"Syntax error: {e.msg}",
+                    'line': e.lineno,
+                    'type': 'syntax',
+                    'text': e.text
+                }],
+                warnings=[]
+            )
+    
+    def _fallback_docstring(self, algorithm: Algorithm) -> str:
+        """Generate basic docstring when LLM fails."""
+        lines = []
+        lines.append(algorithm.description or f"{algorithm.name} implementation.")
+        lines.append("")
+        
+        if algorithm.inputs:
+            lines.append("Args:")
+            for inp in algorithm.inputs:
+                name = inp.get('name', '')
+                desc = inp.get('description', '')
+                ptype = self.type_inferrer.infer_variable_type(name, inp)
+                lines.append(f"    {name} ({ptype}): {desc}")
+            lines.append("")
+        
+        if algorithm.outputs:
+            lines.append("Returns:")
+            for out in algorithm.outputs:
+                name = out.get('name', '')
+                desc = out.get('description', '')
+                lines.append(f"    {name}: {desc}")
+        
+        return '\n'.join(lines)
+    
+    def _generate_stub(self, algorithm: Algorithm) -> GeneratedCode:
+        """
+        Generate stub code when LLM fails.
+        
+        Args:
+            algorithm: Algorithm to generate stub for
+            
+        Returns:
+            GeneratedCode with stub implementation
+        """
+        sig = self.type_inferrer.infer_function_signature(algorithm)
+        
+        stub = f'''def {sig.name}({sig.format_params()}) -> {sig.return_type}:
+    """Stub implementation - LLM generation failed.
+    
+    Steps: {[s.description for s in algorithm.steps]}
+    """
+    raise NotImplementedError("Algorithm not yet implemented")
+'''
+        
+        return GeneratedCode(
+            source=stub,
+            algorithm_name=algorithm.name,
+            imports=[],
+            validation_result=ValidationResult(is_valid=True, errors=[], warnings=[])
+        )
+    
+    def generate_with_fallback(self, algorithm: Algorithm) -> GeneratedCode:
+        """
+        Generate code with graceful fallback.
+        
+        Args:
+            algorithm: Algorithm to generate code for
+            
+        Returns:
+            GeneratedCode (may be stub on failure)
+        """
+        try:
+            return self.generate(algorithm)
+        except LLMGenerationError:
+            return self._generate_stub(algorithm)
+
+
+__all__ = ["LLMCodeGenerator"]

package/src/generation/persistence.py

@@ -0,0 +1,221 @@
+"""Code persistence utilities.
+
+This module handles saving generated code to the filesystem
+with metadata.
+"""
+
+import json
+import re
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from src.extraction.schema import Algorithm
+from src.generation.review import ReviewState
+
+
+class CodePersistence:
+    """
+    Persist generated code to filesystem.
+    
+    Saves code and metadata to .algomath/algorithms/{name}/
+    """
+    
+    def __init__(self, base_path: str = ".algomath"):
+        """
+        Initialize persistence manager.
+        
+        Args:
+            base_path: Base directory for persistence
+        """
+        self.base_path = Path(base_path)
+        self.algorithms_path = self.base_path / "algorithms"
+    
+    def save_generated_code(self,
+                           algorithm: Algorithm,
+                           code: str,
+                           review_state: Optional[ReviewState] = None) -> str:
+        """
+        Save generated code to .algomath/algorithms/{name}/generated.py.
+        
+        Args:
+            algorithm: Algorithm metadata
+            code: Python code to save
+            review_state: Optional review state
+            
+        Returns:
+            Path to saved file
+        """
+        # Create algorithm directory
+        algo_dir = self.algorithms_path / self._sanitize_name(algorithm.name)
+        algo_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Save generated.py
+        generated_file = algo_dir / "generated.py"
+        
+        # Add header comment with metadata
+        header = self._generate_header(algorithm, review_state)
+        full_code = header + "\n\n" + code
+        
+        generated_file.write_text(full_code, encoding='utf-8')
+        
+        # Save metadata
+        self._save_metadata(algo_dir, algorithm, review_state)
+        
+        return str(generated_file)
+    
+    def _generate_header(self,
+                        algorithm: Algorithm,
+                        review_state: Optional[ReviewState]) -> str:
+        """
+        Generate file header with metadata.
+        
+        Args:
+            algorithm: Algorithm metadata
+            review_state: Optional review state
+            
+        Returns:
+            Header comment string
+        """
+        lines = ['"""']
+        lines.append(f"Generated code for: {algorithm.name}")
+        lines.append(f"Description: {algorithm.description}")
+        lines.append(f"Generated at: {datetime.now().isoformat()}")
+        
+        if review_state:
+            lines.append(f"Status: {'Approved' if review_state.is_approved else 'Pending review'}")
+            lines.append(f"Edited: {'Yes' if review_state.is_edited else 'No'}")
+        
+        lines.append('"""')
+        return '\n'.join(lines)
+    
+    def _save_metadata(self,
+                      algo_dir: Path,
+                      algorithm: Algorithm,
+                      review_state: Optional[ReviewState]):
+        """
+        Save metadata JSON alongside code.
+        
+        Args:
+            algo_dir: Algorithm directory
+            algorithm: Algorithm metadata
+            review_state: Optional review state
+        """
+        metadata: Dict[str, Any] = {
+            'name': algorithm.name,
+            'description': algorithm.description,
+            'inputs': algorithm.inputs,
+            'outputs': algorithm.outputs,
+            'step_count': len(algorithm.steps),
+            'saved_at': datetime.now().isoformat(),
+        }
+        
+        if review_state:
+            metadata['review'] = {
+                'is_edited': review_state.is_edited,
+                'is_approved': review_state.is_approved,
+                'approved_at': review_state.approved_at.isoformat() if review_state.approved_at else None
+            }
+        
+        metadata_file = algo_dir / "metadata.json"
+        metadata_file.write_text(json.dumps(metadata, indent=2))
+    
+    def load_generated_code(self, algorithm_name: str) -> Optional[str]:
+        """
+        Load previously generated code.
+        
+        Args:
+            algorithm_name: Algorithm name
+            
+        Returns:
+            Code string or None if not found
+        """
+        algo_dir = self.algorithms_path / self._sanitize_name(algorithm_name)
+        generated_file = algo_dir / "generated.py"
+        
+        if generated_file.exists():
+            content = generated_file.read_text()
+            # Strip header
+            if content.startswith('"""'):
+                end = content.find('"""', 3)
+                if end != -1:
+                    return content[end + 3:].strip()
+            return content
+        return None
+    
+    def load_metadata(self, algorithm_name: str) -> Optional[Dict[str, Any]]:
+        """
+        Load metadata for an algorithm.
+        
+        Args:
+            algorithm_name: Algorithm name
+            
+        Returns:
+            Metadata dict or None if not found
+        """
+        algo_dir = self.algorithms_path / self._sanitize_name(algorithm_name)
+        metadata_file = algo_dir / "metadata.json"
+        
+        if metadata_file.exists():
+            return json.loads(metadata_file.read_text())
+        return None
+    
+    def _sanitize_name(self, name: str) -> str:
+        """
+        Convert algorithm name to filesystem-safe directory name.
+        
+        Args:
+            name: Algorithm name
+            
+        Returns:
+            Sanitized name
+        """
+        # Remove special characters
+        safe = re.sub(r'[^\w\s-]', '', name)
+        # Replace spaces and dashes with underscores
+        safe = re.sub(r'[-\s]+', '_', safe)
+        return safe.lower().strip('_')
+
+
+def save_to_context(context: Any,
+                   algorithm: Algorithm,
+                   review: 'CodeReviewInterface') -> str:
+    """
+    Save approved code to context.
+    
+    Args:
+        context: ContextManager instance
+        algorithm: Algorithm metadata
+        review: CodeReviewInterface with approval
+        
+    Returns:
+        Path to saved file
+        
+    Raises:
+        ValueError: If code not approved
+    """
+    persistence = CodePersistence()
+    
+    if not review.is_approved():
+        raise ValueError("Code must be approved before saving")
+    
+    code = review.get_modified_code()
+    path = persistence.save_generated_code(
+        algorithm,
+        code,
+        review.state
+    )
+    
+    # Update context
+    context.save_code(code)
+    # Assuming context has save_code_path method
+    if hasattr(context, 'save_code_path'):
+        context.save_code_path(path)
+    
+    return path
+
+
+__all__ = [
+    'CodePersistence',
+    'save_to_context',
+]

package/src/generation/prompts.py

@@ -0,0 +1,202 @@
+"""LLM prompts for code generation.
+
+This module provides prompts for generating Python code from
+structured algorithms using LLM assistance.
+"""
+
+from typing import Any, Dict, List
+
+from src.extraction.schema import Algorithm
+
+
+# System prompts
+CODE_GENERATION_SYSTEM_PROMPT = """You are a Python code generator for mathematical algorithms.
+
+Your task: Generate executable Python code from structured algorithm steps.
+
+Rules:
+1. Use type hints for all parameters and return values
+2. Include Google-style docstrings with Args, Returns sections
+3. Use numpy (np) for matrix/vector operations
+4. Use math module for mathematical functions (sqrt, sin, cos, etc.)
+5. Follow PEP 8 naming conventions (snake_case)
+6. Generate clean, well-formatted code
+7. Handle edge cases gracefully
+
+Output format: Return ONLY valid Python code, no markdown fences, no explanations.
+"""
+
+
+CODE_GENERATION_USER_TEMPLATE = """Generate Python code for this algorithm:
+
+Algorithm Name: {name}
+Description: {description}
+
+Inputs:
+{inputs}
+
+Outputs:
+{outputs}
+
+Steps:
+{steps}
+
+Requirements:
+- Include type hints
+- Include comprehensive docstring
+- Use numpy for arrays/matrices
+- Generate helper functions if needed
+- Handle all edge cases
+
+Generate only the Python function code."""
+
+
+COMPLEX_EXPRESSION_SYSTEM_PROMPT = """You are a Python expression generator.
+
+Convert mathematical expressions into valid Python code.
+
+Examples:
+- "sum of A[i] for i from 1 to n" -> "sum(A[i] for i in range(n))"
+- "minimum distance from u to v" -> "min(dist[u], dist[v])"
+- "argmin of distances" -> "np.argmin(distances)"
+- "set of visited nodes" -> "visited = set()"
+
+Use standard Python operators and functions."""
+
+
+DOCSTRING_GENERATION_PROMPT = """Generate a Google-style docstring for this algorithm.
+
+Include:
+- One-line summary
+- Detailed description (from algorithm description)
+- Args section with types and descriptions
+- Returns section with type and description
+- Raises section (if applicable)
+- Complexity notation (if inferable: O(n), O(n²), etc.)
+- Step references (links to step IDs)
+
+Example format:
+Triple-quote
+{one_line_summary}
+
+{detailed_description}
+
+Args:
+{param}: {type} - {description}
+
+Returns:
+{type}: {description}
+
+Raises:
+ValueError: When {condition}
+
+Time Complexity: {complexity}
+Space Complexity: {complexity}
+Triple-quote
+"""
+
+
+def format_inputs(inputs: List[Dict[str, Any]]) -> str:
+    """Format inputs for prompt."""
+    if not inputs:
+        return "  (none)"
+    lines = []
+    for inp in inputs:
+        name = inp.get('name', 'unnamed')
+        desc = inp.get('description', '')
+        lines.append(f"  - {name}: {desc}")
+    return '\n'.join(lines)
+
+
+def format_outputs(outputs: List[Dict[str, Any]]) -> str:
+    """Format outputs for prompt."""
+    if not outputs:
+        return "  (none)"
+    lines = []
+    for out in outputs:
+        name = out.get('name', 'unnamed')
+        desc = out.get('description', '')
+        lines.append(f"  - {name}: {desc}")
+    return '\n'.join(lines)
+
+
+def format_steps(steps: List[Any]) -> str:
+    """Format steps for prompt."""
+    if not steps:
+        return "  (none)"
+    lines = []
+    for i, step in enumerate(steps, 1):
+        lines.append(f"  {i}. [{step.type.value if hasattr(step, 'type') else 'step'}] {step.description}")
+    return '\n'.join(lines)
+
+
+def format_code_generation_prompt(algorithm: Algorithm) -> List[Dict[str, str]]:
+    """
+    Format messages for LLM code generation.
+    
+    Args:
+        algorithm: Algorithm to generate code for
+        
+    Returns:
+        List of message dicts for LLM API
+    """
+    return [
+        {"role": "system", "content": CODE_GENERATION_SYSTEM_PROMPT},
+        {"role": "user", "content": CODE_GENERATION_USER_TEMPLATE.format(
+            name=algorithm.name,
+            description=algorithm.description,
+            inputs=format_inputs(algorithm.inputs),
+            outputs=format_outputs(algorithm.outputs),
+            steps=format_steps(algorithm.steps)
+        )}
+    ]
+
+
+def format_complex_expression_prompt(expression: str, context: Dict[str, Any] = None) -> List[Dict[str, str]]:
+    """
+    Format messages for complex expression generation.
+    
+    Args:
+        expression: Mathematical expression to convert
+        context: Variable context
+        
+    Returns:
+        List of message dicts for LLM API
+    """
+    context_str = ""
+    if context:
+        vars_str = '\n'.join(f"  - {k}: {v}" for k, v in context.items())
+        context_str = f"\nAvailable variables:\n{vars_str}"
+    
+    return [
+        {"role": "system", "content": COMPLEX_EXPRESSION_SYSTEM_PROMPT},
+        {"role": "user", "content": f"Convert to Python: {expression}{context_str}"}
+    ]
+
+
+def format_docstring_generation_prompt(algorithm: Algorithm) -> List[Dict[str, str]]:
+    """
+    Format messages for docstring generation.
+    
+    Args:
+        algorithm: Algorithm to generate docstring for
+        
+    Returns:
+        List of message dicts for LLM API
+    """
+    return [
+        {"role": "system", "content": DOCSTRING_GENERATION_PROMPT},
+        {"role": "user", "content": f"Generate docstring for algorithm: {algorithm.name}\n\n{algorithm.description}"}
+    ]
+
+
+# Exports
+__all__ = [
+    "CODE_GENERATION_SYSTEM_PROMPT",
+    "CODE_GENERATION_USER_TEMPLATE",
+    "COMPLEX_EXPRESSION_SYSTEM_PROMPT",
+    "DOCSTRING_GENERATION_PROMPT",
+    "format_code_generation_prompt",
+    "format_complex_expression_prompt",
+    "format_docstring_generation_prompt",
+]