scry-run 0.1.0__py3-none-any.whl

scry_run/generator.py

@@ -0,0 +1,698 @@
+"""Code generator using pluggable LLM backends with structured JSON output."""
+
+from __future__ import annotations
+
+import ast
+import os
+import sys
+import time
+import random
+import json
+from dataclasses import dataclass
+from typing import Optional
+
+from scry_run.console import warning
+
+
+@dataclass
+class GeneratedCode:
+    """Result of code generation."""
+
+    code: str
+    code_type: str  # method, property, classmethod, staticmethod
+    docstring: str
+    dependencies: list[str]
+    packages: list[str] = None  # PyPI packages to install
+
+    def __post_init__(self):
+        if self.packages is None:
+            self.packages = []
+
+
+# JSON schema for structured output
+CODE_GENERATION_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "code": {
+            "type": "string",
+            "description": "The Python code for the requested method/property. Must be valid Python that can be exec'd."
+        },
+        "code_type": {
+            "type": "string",
+            "enum": ["method", "property", "classmethod", "staticmethod"],
+            "description": "The type of code being generated"
+        },
+        "docstring": {
+            "type": "string",
+            "description": "A brief description of what this code does"
+        },
+        "dependencies": {
+            "type": "array",
+            "items": {"type": "string"},
+            "description": "List of import statements needed (e.g., 'import json', 'from pathlib import Path')"
+        },
+        "packages": {
+            "type": "array",
+            "items": {"type": "string"},
+            "description": "List of PyPI packages to install (e.g., 'pygame', 'requests>=2.0'). Only include packages not in the standard library."
+        }
+    },
+    "required": ["code", "code_type", "docstring", "dependencies", "packages"]
+}
+
+
+# =============================================================================
+# User-Friendly Exception Classes
+# =============================================================================
+
+class ScryRunError(Exception):
+    """Base exception for scry-run errors with user-friendly messages."""
+    
+    def __init__(self, message: str, hint: Optional[str] = None):
+        self.message = message
+        self.hint = hint
+        super().__init__(self._format_message())
+    
+    def _format_message(self) -> str:
+        msg = f"[scry-run] {self.message}"
+        if self.hint:
+            msg += f"\n         💡 {self.hint}"
+        return msg
+
+
+class APIKeyError(ScryRunError):
+    """API key is missing or invalid."""
+    pass
+
+
+class RateLimitError(ScryRunError):
+    """Rate limit exceeded."""
+    pass
+
+
+class QuotaExceededError(ScryRunError):
+    """API quota exhausted."""
+    pass
+
+
+class ModelNotFoundError(ScryRunError):
+    """Requested model doesn't exist or isn't accessible."""
+    pass
+
+
+class ContentBlockedError(ScryRunError):
+    """Content was blocked by safety filters."""
+    pass
+
+
+class NetworkError(ScryRunError):
+    """Network connectivity issue."""
+    pass
+
+
+class CodeGenerationError(ScryRunError):
+    """Code generation failed."""
+    pass
+
+
+class CodeValidationError(ScryRunError):
+    """Generated code failed to parse."""
+    pass
+
+
+def _handle_api_error(error: Exception, model: str) -> ScryRunError:
+    """Convert API exceptions to user-friendly errors.
+
+    Args:
+        error: The original exception from the backend
+        model: The model name being used (for error messages)
+
+    Returns:
+        A user-friendly ScryRunError subclass
+    """
+    error_str = str(error).lower()
+    error_type = type(error).__name__
+
+    # Authentication errors
+    if "api_key" in error_str or "authentication" in error_str or "401" in error_str:
+        return APIKeyError(
+            "Invalid API key",
+            "Check your API key configuration"
+        )
+    
+    # Rate limiting
+    if "rate" in error_str and "limit" in error_str or "429" in error_str or "resource_exhausted" in error_str:
+        return RateLimitError(
+            "Rate limit exceeded",
+            "Wait a moment and try again. Consider using a slower request rate."
+        )
+    
+    # Quota exceeded
+    if "quota" in error_str or "billing" in error_str:
+        return QuotaExceededError(
+            "API quota exhausted",
+            "Check your API quota and billing settings"
+        )
+    
+    # Model not found
+    if "model" in error_str and ("not found" in error_str or "not exist" in error_str or "404" in error_str):
+        return ModelNotFoundError(
+            f"Model '{model}' not found",
+            f"Try a different model. Set SCRY_MODEL or backend-specific env var."
+        )
+    
+    # Content blocked by safety filters
+    if "safety" in error_str or "blocked" in error_str or "harm" in error_str:
+        return ContentBlockedError(
+            "Content blocked by safety filters",
+            "Try rephrasing your class/method names or docstrings"
+        )
+    
+    # Network errors
+    if any(x in error_str for x in ["connection", "timeout", "network", "dns", "ssl"]):
+        return NetworkError(
+            "Network error connecting to backend",
+            "Check your internet connection and firewall settings"
+        )
+
+    # Permission/access errors
+    if "permission" in error_str or "403" in error_str or "access" in error_str:
+        return APIKeyError(
+            "API access denied",
+            "Check your API key permissions"
+        )
+
+    # Server errors
+    if "500" in error_str or "503" in error_str or "server" in error_str:
+        return NetworkError(
+            "Backend server error",
+            "The API is temporarily unavailable. Try again in a few seconds."
+        )
+    
+    # Fallback: wrap unknown errors
+    return CodeGenerationError(
+        f"API error: {error_type}: {error}",
+        "Check the error message above. If persistent, file an issue."
+    )
+
+
+class GenerationLimitError(ScryRunError):
+    """Global generation limit exceeded."""
+    pass
+
+
+class CodeGenerator:
+    """Generates code using pluggable backends.
+
+    The generator maintains a persistent session with the backend.
+    Full codebase context is sent once at startup, then each generation
+    request sends only the minimal "frame" (class/method info).
+
+    Configure via SCRY_BACKEND env var or backend parameter.
+    """
+
+    MAX_RETRIES = 5
+    
+    # Global counter for all generations in this process
+    _total_generations_count = 0
+    
+    # Track context per-generator instance
+    _context_hash: Optional[str] = None
+    
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: Optional[str] = None,
+        backend: Optional[str] = None,
+    ):
+        """Initialize the code generator.
+
+        Args:
+            api_key: Deprecated, not used.
+            model: Model to use. Defaults to backend-specific default.
+            backend: Backend to use: "claude", "frozen", or "auto" (default).
+                     Can also be set via SCRY_BACKEND env var.
+        """
+        from scry_run.backends import get_backend
+
+        # Pass model through - backends handle their own defaults
+        self.model = model or os.environ.get("SCRY_MODEL")
+        self._backend = get_backend(name=backend, model=self.model)
+        self._context_provided = False
+
+    
+    def _format_installed_packages(self, packages: list[str] | None) -> str:
+        """Format installed packages for the prompt."""
+        if not packages:
+            return "No packages are currently installed in this app's environment."
+
+        pkg_list = "\n".join(f"- {pkg}" for pkg in packages)
+        return f"""The following packages are already installed:
+{pkg_list}
+
+Prefer using these before requesting new packages. Only add to "packages" if you need something not listed above."""
+
+    def _build_context_prompt(
+        self,
+        context: str,
+        installed_packages: list[str] | None = None,
+    ) -> str:
+        """Build the full context prompt (sent once at session start).
+        
+        Args:
+            context: Full codebase context with hole marker
+            installed_packages: Packages already installed
+            
+        Returns:
+            The formatted context string for the system prompt
+        """
+        return f"""## Codebase Context
+
+The following is the full codebase you are working with:
+
+```python
+{context}
+```
+
+## Installed Packages
+
+{self._format_installed_packages(installed_packages)}
+
+## Critical Requirements for Code Generation
+
+### Code Quality
+1. The code MUST be valid Python that can be executed with `exec()`
+2. Complete method/property definition starting with `def` or `@property`
+3. Use appropriate decorators (@classmethod, @staticmethod, @property) if needed
+4. STRICTLY USE TYPE HINTS for all parameters and return values
+
+### Application-Level Thinking
+5. READ THE CLASS DOCSTRING - it describes what the application does
+6. EXTRAPOLATE sensibly based on the application's purpose
+7. Generate code that a HUMAN would write for this application
+
+### Dependencies
+8. List ALL imports needed in the dependencies field
+9. Prefer standard library when possible
+"""
+    
+    def _build_frame_prompt(
+        self,
+        class_name: str,
+        attr_name: str,
+        is_classmethod: bool = False,
+        runtime_context: str = "",
+    ) -> str:
+        """Build the minimal frame prompt (sent for each generation).
+        
+        This assumes the full codebase has already been provided.
+        Runtime context (call stack, variables) is included since it's dynamic.
+        
+        Args:
+            class_name: Name of the class needing the attribute
+            attr_name: Name of the attribute to generate
+            is_classmethod: Whether this is a class-level attribute
+            runtime_context: Dynamic context (call stack, local variables)
+            
+        Returns:
+            The minimal prompt for this specific generation
+        """
+        attr_type = "class method/property" if is_classmethod else "instance method/property"
+        
+        # Include runtime context if provided
+        context_section = ""
+        if runtime_context:
+            context_section = f"\n## Runtime Context\n\n{runtime_context}\n"
+        
+        return f"""Generate the `{attr_name}` {attr_type} for the `{class_name}` class.
+{context_section}
+Respond with a JSON object containing:
+- code: The Python code (complete method/property definition with type hints)
+- code_type: One of "method", "property", "classmethod", "staticmethod"
+- docstring: Brief description of what this code does
+- dependencies: List of import statements needed
+- packages: List of PyPI packages to install (if any external packages needed)
+
+IMPORTANT: Output ONLY the JSON object, no markdown, no explanation, no code fences."""
+
+    def _build_prompt(
+        self,
+        context: str,
+        class_name: str,
+        attr_name: str,
+        is_classmethod: bool = False,
+        installed_packages: list[str] | None = None,
+    ) -> str:
+        """Build the full prompt for code generation (legacy/fallback).
+        
+        This is used when the backend doesn't support set_context().
+        
+        Args:
+            context: Full codebase context with hole marker
+            class_name: Name of the class needing the attribute
+            attr_name: Name of the attribute to generate
+            is_classmethod: Whether this is a class-level attribute
+            
+        Returns:
+            The formatted prompt string
+        """
+        attr_type = "class method/property" if is_classmethod else "instance method/property"
+        
+        return f"""You are an expert Python code generator for the scry-run dynamic code generation system.
+
+## Your Role
+
+You are generating code for an APPLICATION, not just isolated functions. Read the class docstring carefully - it describes the PURPOSE of the entire application. Use this to EXTRAPOLATE what this method should do.
+
+## Codebase Context
+
+The following is the full codebase. There is a marker showing where code is missing:
+
+```python
+{context}
+```
+
+## Task
+
+Generate the `{attr_name}` {attr_type} for the `{class_name}` class.
+
+## Critical Requirements
+
+### Code Quality
+1. The code MUST be valid Python that can be executed with `exec()`
+2. Complete method/property definition starting with `def` or `@property`
+3. Use appropriate decorators (@classmethod, @staticmethod, @property) if needed
+4. STRICTLY USE TYPE HINTS for all parameters and return values (e.g. `def foo(x: int) -> str:`)
+   - Import `typing` types (List, Dict, Optional, Any) as needed and add to dependencies
+
+### Rich Contextual Comments
+4. Include a DETAILED docstring explaining:
+   - What this method does
+   - How it fits into the overall application described in the class docstring
+   - Any assumptions made about the application's behavior
+5. Add inline comments for non-obvious logic
+6. Document parameter types and return types
+
+### Application-Level Thinking
+7. READ THE CLASS DOCSTRING - it describes what the application does
+8. EXTRAPOLATE sensibly: if the app is a "todo list", a `save()` method should persist todos
+9. Infer reasonable behavior from:
+   - The application description (class docstring)
+   - The method/property name
+   - Other existing methods
+   - How similar applications typically work
+10. Generate code that a HUMAN would write for this application, not minimal stubs
+11. Think about what OTHER PARAMETERS this method might need beyond the obvious ones:
+    - Configuration options (e.g., format, encoding, verbosity)
+    - Optional filters or modifiers
+    - Callbacks or hooks for extensibility
+    - Make these optional with sensible defaults
+
+### Dependencies
+11. List ALL imports needed in the dependencies field
+12. Prefer standard library when possible
+13. The code should work standalone once dependencies are imported
+14. Prefer standard library when possible
+15. The code should work standalone once dependencies are imported
+
+### Packages
+
+The app has its own virtual environment. If you need external packages:
+1. Add import statements to "dependencies" (e.g., "import pygame")
+2. Add PyPI package names to "packages" (e.g., "pygame")
+
+{self._format_installed_packages(installed_packages)}
+
+## Output Format
+
+Respond with a JSON object containing:
+- code: The Python code (complete method/property definition with rich comments)
+- code_type: One of "method", "property", "classmethod", "staticmethod"
+- docstring: Brief description of what this code does
+- dependencies: List of import statements needed
+- packages: List of PyPI packages to install (if any external packages needed)
+
+IMPORTANT: Output ONLY the JSON object, no markdown, no explanation, no code fences."""
+    
+    def _fix_indentation(self, code: str) -> str:
+        """Fix common indentation issues in generated code.
+        
+        The LLM sometimes returns code with leading indentation or
+        inconsistent whitespace. This method normalizes it.
+        
+        Args:
+            code: Raw generated code
+            
+        Returns:
+            Code with fixed indentation
+        """
+        import textwrap
+        from scry_run.logging import get_logger
+        logger = get_logger()
+        
+        original_code = code
+        
+        # Strip leading/trailing empty lines
+        lines = code.split('\n')
+        while lines and not lines[0].strip():
+            lines.pop(0)
+        while lines and not lines[-1].strip():
+            lines.pop()
+        
+        if not lines:
+            return code
+        
+        # Check if first non-empty line is indented
+        first_line = lines[0]
+        if first_line and first_line[0] in ' \t':
+            # Use textwrap.dedent to remove common leading whitespace
+            code = textwrap.dedent('\n'.join(lines))
+            logger.debug("Fixed indentation in generated code", f"ORIGINAL:\n{original_code}\n\nFIXED:\n{code}")
+        else:
+            code = '\n'.join(lines)
+        
+        return code
+    
+    def _validate_code(self, code: str) -> None:
+        """Validate that the code parses correctly.
+        
+        Args:
+            code: Python code to validate
+            
+        Raises:
+            CodeValidationError: If code fails to parse
+        """
+        try:
+            ast.parse(code)
+        except SyntaxError as e:
+            raise CodeValidationError(f"Generated code has syntax error: {e}")
+
+    
+    def _parse_response(self, response_text: str) -> GeneratedCode:
+        """Parse the JSON response from the LLM.
+        
+        Args:
+            response_text: Raw response text from LLM
+            
+        Returns:
+            GeneratedCode object
+            
+        Raises:
+            CodeGenerationError: If response cannot be parsed
+        """
+        try:
+            data = json.loads(response_text)
+        except json.JSONDecodeError as e:
+            raise CodeGenerationError(f"Failed to parse JSON response: {e}")
+        
+        required_fields = ["code", "code_type", "docstring", "dependencies", "packages"]
+        for field in required_fields:
+            if field not in data:
+                raise CodeGenerationError(f"Missing required field: {field}")
+
+        return GeneratedCode(
+            code=data["code"],
+            code_type=data["code_type"],
+            docstring=data["docstring"],
+            dependencies=data["dependencies"],
+            packages=data["packages"],
+        )
+    def generate(
+        self,
+        context: str,
+        class_name: str,
+        attr_name: str,
+        is_classmethod: bool = False,
+        installed_packages: list[str] | None = None,
+    ) -> GeneratedCode:
+        """Generate code for a missing attribute.
+
+        Args:
+            context: Full codebase context with hole marker
+            class_name: Name of the class needing the attribute
+            attr_name: Name of the attribute to generate
+            is_classmethod: Whether this is a class-level attribute
+            installed_packages: List of packages already installed in the app's venv
+
+        Returns:
+            GeneratedCode with the generated code and metadata
+
+        Raises:
+            CodeGenerationError: If generation fails after retries
+            CodeValidationError: If generated code doesn't parse
+            APIKeyError: If API key is invalid
+            RateLimitError: If rate limit is exceeded
+            QuotaExceededError: If quota is exhausted
+            NetworkError: If network issues occur
+            GenerationLimitError: If global generation limit is exceeded
+        """
+        # Check global generation limit
+        max_limit_str = os.environ.get("SCRY_MAX_GENERATIONS", "100")
+        if max_limit_str:
+            try:
+                limit = int(max_limit_str)
+                if CodeGenerator._total_generations_count >= limit:
+                    raise GenerationLimitError(
+                        f"Global generation limit of {limit} reached.",
+                        f"Increase SCRY_MAX_GENERATIONS env var if this is intended."
+                    )
+            except ValueError:
+                # Ignore invalid integer values in env var
+                pass
+
+        # Increment counter
+        CodeGenerator._total_generations_count += 1
+
+        # Log generation start
+        from scry_run.logging import get_logger
+        logger = get_logger()
+        
+        # Set generation context for frozen backend error messages
+        if hasattr(self._backend, 'set_generation_context'):
+            self._backend.set_generation_context(class_name, attr_name)
+
+        # Check if backend supports persistent context
+        supports_context = hasattr(self._backend, 'set_context')
+        
+        # Extract runtime context from the full context (it's appended at the end)
+        # The runtime context is dynamic and must be included in every request
+        runtime_context = ""
+        if "# === RUNTIME CONTEXT ===" in context:
+            idx = context.find("# === RUNTIME CONTEXT ===")
+            runtime_context = context[idx:]
+        
+        if supports_context and not self._context_provided:
+            # First generation: send full context to backend (without runtime context)
+            # Runtime context is dynamic and will be sent with each frame
+            static_context = context[:context.find("# === RUNTIME CONTEXT ===")] if runtime_context else context
+            context_prompt = self._build_context_prompt(static_context, installed_packages)
+            self._backend.set_context(context_prompt)
+            self._context_provided = True
+            logger.info("Generator: Provided full context to backend session")
+            
+            # For this generation, send frame with runtime context
+            prompt = self._build_frame_prompt(class_name, attr_name, is_classmethod, runtime_context)
+        elif supports_context:
+            # Subsequent generations: send frame with runtime context
+            prompt = self._build_frame_prompt(class_name, attr_name, is_classmethod, runtime_context)
+        else:
+            # Backend doesn't support persistent context, use full prompt
+            prompt = self._build_prompt(context, class_name, attr_name, is_classmethod, installed_packages)
+        
+        logger.log_generation_start(class_name, attr_name, prompt)
+
+        # Delegate to backend with retries
+        last_error: Optional[Exception] = None
+
+        for attempt in range(self.MAX_RETRIES):
+            try:
+                result = self._backend.generate_code(prompt)
+                
+                # Log raw response
+                logger.log_generation_response(result.code if hasattr(result, 'code') else str(result))
+                
+                # Convert backend result to our GeneratedCode type
+                from scry_run.backends.base import GeneratedCode as BackendResult
+                if isinstance(result, BackendResult):
+                    result = GeneratedCode(
+                        code=result.code,
+                        code_type=result.code_type,
+                        docstring=result.docstring,
+                        dependencies=result.dependencies,
+                        packages=result.packages,
+                    )
+
+                # Fix indentation issues before validation
+                result = GeneratedCode(
+                    code=self._fix_indentation(result.code),
+                    code_type=result.code_type,
+                    docstring=result.docstring,
+                    dependencies=result.dependencies,
+                    packages=result.packages,
+                )
+                
+                # Validate the generated code
+                self._validate_code(result.code)
+
+                
+                # Log success
+                logger.log_generation_success(class_name, attr_name, result.code)
+                
+                return result
+                
+            except (CodeValidationError, CodeGenerationError) as e:
+                last_error = e
+                # Log the validation error with full details
+                code_str = result.code if 'result' in dir() and hasattr(result, 'code') else None
+                logger.log_validation_error(e, code_str)
+                
+                if attempt >= self.MAX_RETRIES - 1:
+                    raise
+                warning(f"Validation error, retrying ({attempt + 1}/{self.MAX_RETRIES})...")
+            except Exception as e:
+                # Log the error
+                logger.error(f"Backend error: {e}", str(e))
+
+                # Convert to friendly error
+                friendly_error = _handle_api_error(e, self.model)
+                last_error = friendly_error
+
+                if isinstance(friendly_error, (RateLimitError, NetworkError)):
+                    wait_time = 2 ** attempt
+                    warning(f"{friendly_error.message}. Retrying in {wait_time:.1f}s...")
+                    time.sleep(wait_time)
+                else:
+                    raise friendly_error
+        
+        raise CodeGenerationError(
+            f"Failed to generate valid code after {self.MAX_RETRIES} attempts",
+            f"Last error: {last_error}"
+        )
+
+
+    def generate_freeform(
+        self,
+        prompt: str,
+        system_instruction: Optional[str] = "You are a helpful coding assistant.",
+    ) -> str:
+        """Generate freeform text from a prompt.
+        
+        Args:
+            prompt: User prompt
+            system_instruction: System instruction to guide the model behavior
+            
+        Returns:
+            Generated text content
+            
+        Raises:
+            ScryRunError: On API errors
+        """
+        # Prepend system instruction if provided
+        full_prompt = prompt
+        if system_instruction:
+            full_prompt = f"{system_instruction}\n\n{prompt}"
+        
+        try:
+            return self._backend.generate_text(full_prompt)
+        except Exception as e:
+            raise _handle_api_error(e, self.model)