flashlite 0.1.0__py3-none-any.whl

flashlite/structured/outputs.py

@@ -0,0 +1,189 @@
+"""Structured output parsing and validation."""
+
+import json
+import logging
+import re
+from typing import TypeVar
+
+from pydantic import BaseModel, ValidationError
+
+from ..types import CompletionResponse
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class StructuredOutputError(Exception):
+    """Error parsing or validating structured output."""
+
+    def __init__(
+        self,
+        message: str,
+        raw_content: str | None = None,
+        validation_errors: list[dict] | None = None,
+    ):
+        super().__init__(message)
+        self.raw_content = raw_content
+        self.validation_errors = validation_errors or []
+
+
+def parse_json_response(content: str) -> dict:
+    """
+    Parse JSON from an LLM response.
+
+    Handles common cases where the model wraps JSON in markdown code blocks
+    or includes extra text.
+
+    Args:
+        content: The raw response content from the LLM
+
+    Returns:
+        Parsed JSON as a dict
+
+    Raises:
+        StructuredOutputError: If JSON cannot be parsed
+    """
+    # Strip whitespace
+    content = content.strip()
+
+    # Try direct parse first
+    try:
+        return json.loads(content)
+    except json.JSONDecodeError:
+        pass
+
+    # Try to extract JSON from markdown code blocks
+    # Match ```json ... ``` or ``` ... ```
+    json_block_pattern = r"```(?:json)?\s*([\s\S]*?)```"
+    matches = re.findall(json_block_pattern, content)
+
+    for match in matches:
+        try:
+            return json.loads(match.strip())
+        except json.JSONDecodeError:
+            continue
+
+    # Try to find JSON object boundaries
+    # Look for first { and last }
+    first_brace = content.find("{")
+    last_brace = content.rfind("}")
+
+    if first_brace != -1 and last_brace != -1 and last_brace > first_brace:
+        try:
+            return json.loads(content[first_brace : last_brace + 1])
+        except json.JSONDecodeError:
+            pass
+
+    # Try to find JSON array boundaries
+    first_bracket = content.find("[")
+    last_bracket = content.rfind("]")
+
+    if first_bracket != -1 and last_bracket != -1 and last_bracket > first_bracket:
+        try:
+            return json.loads(content[first_bracket : last_bracket + 1])
+        except json.JSONDecodeError:
+            pass
+
+    raise StructuredOutputError(
+        f"Could not parse JSON from response: {content[:200]}...",
+        raw_content=content,
+    )
+
+
+def validate_response(
+    response: CompletionResponse,
+    model: type[T],
+) -> T:
+    """
+    Parse and validate a completion response against a Pydantic model.
+
+    Args:
+        response: The completion response to validate
+        model: The Pydantic model class to validate against
+
+    Returns:
+        A validated instance of the model
+
+    Raises:
+        StructuredOutputError: If parsing or validation fails
+    """
+    content = response.content
+
+    # Parse JSON
+    try:
+        data = parse_json_response(content)
+    except StructuredOutputError:
+        raise
+
+    # Validate against model
+    try:
+        return model.model_validate(data)
+    except ValidationError as e:
+        errors = e.errors()
+        error_messages = []
+        for err in errors:
+            loc = ".".join(str(x) for x in err["loc"])
+            msg = err["msg"]
+            error_messages.append(f"  - {loc}: {msg}")
+
+        raise StructuredOutputError(
+            "Validation failed:\n" + "\n".join(error_messages),
+            raw_content=content,
+            validation_errors=[dict(e) for e in errors],
+        )
+
+
+def format_validation_error_for_retry(error: StructuredOutputError) -> str:
+    """
+    Format a validation error as feedback for the model to retry.
+
+    This creates a message that explains what went wrong so the model
+    can correct its response.
+
+    Args:
+        error: The structured output error
+
+    Returns:
+        A formatted error message for the retry prompt
+    """
+    lines = ["Your previous response had the following errors:", ""]
+
+    if error.validation_errors:
+        for err in error.validation_errors:
+            loc = ".".join(str(x) for x in err.get("loc", []))
+            msg = err.get("msg", "Unknown error")
+            input_val = err.get("input")
+
+            if loc:
+                lines.append(f"- Field '{loc}': {msg}")
+                if input_val is not None:
+                    lines.append(f"  Got: {input_val!r}")
+            else:
+                lines.append(f"- {msg}")
+    else:
+        lines.append(f"- {str(error)}")
+
+    lines.append("")
+    lines.append("Please correct these errors and respond with valid JSON.")
+
+    return "\n".join(lines)
+
+
+def extract_json_from_content(content: str) -> str | None:
+    """
+    Extract just the JSON portion from content that may contain other text.
+
+    Returns the JSON string if found, None otherwise.
+
+    Args:
+        content: The raw content
+
+    Returns:
+        The extracted JSON string or None
+    """
+    try:
+        data = parse_json_response(content)
+        return json.dumps(data)
+    except StructuredOutputError:
+        return None

flashlite/structured/schema.py

@@ -0,0 +1,165 @@
+"""JSON schema generation from Pydantic models."""
+
+import json
+from typing import Any
+
+from pydantic import BaseModel
+
+
+def generate_json_schema(model: type[BaseModel]) -> dict[str, Any]:
+    """
+    Generate a JSON schema from a Pydantic model.
+
+    This uses Pydantic's built-in schema generation and formats it
+    for use with LLM structured output features.
+
+    Args:
+        model: A Pydantic BaseModel class
+
+    Returns:
+        JSON schema dict suitable for LLM prompts or response_format
+    """
+    # Use Pydantic's built-in schema generation
+    schema = model.model_json_schema()
+
+    # Clean up schema for LLM consumption
+    schema = _simplify_schema(schema)
+
+    return schema
+
+
+def _simplify_schema(schema: dict[str, Any]) -> dict[str, Any]:
+    """
+    Simplify a JSON schema for LLM consumption.
+
+    Removes Pydantic-specific fields and flattens $defs references
+    where possible for cleaner prompts.
+    """
+    # Remove Pydantic-specific metadata
+    keys_to_remove = ["title", "$defs", "definitions"]
+
+    # If there are $defs, inline them
+    defs = schema.get("$defs", schema.get("definitions", {}))
+
+    def resolve_refs(obj: Any) -> Any:
+        """Recursively resolve $ref references."""
+        if isinstance(obj, dict):
+            if "$ref" in obj:
+                ref_path = obj["$ref"]
+                # Extract the definition name from "#/$defs/Name"
+                if ref_path.startswith("#/$defs/"):
+                    def_name = ref_path.split("/")[-1]
+                    if def_name in defs:
+                        # Return a copy of the resolved definition
+                        resolved = resolve_refs(defs[def_name].copy())
+                        # Remove title from inlined definitions
+                        resolved.pop("title", None)
+                        return resolved
+                elif ref_path.startswith("#/definitions/"):
+                    def_name = ref_path.split("/")[-1]
+                    if def_name in defs:
+                        resolved = resolve_refs(defs[def_name].copy())
+                        resolved.pop("title", None)
+                        return resolved
+            return {k: resolve_refs(v) for k, v in obj.items() if k not in keys_to_remove}
+        elif isinstance(obj, list):
+            return [resolve_refs(item) for item in obj]
+        return obj
+
+    result = resolve_refs(schema)
+
+    # Remove top-level title if present
+    result.pop("title", None)
+
+    return result
+
+
+def schema_to_prompt(model: type[BaseModel]) -> str:
+    """
+    Convert a Pydantic model to a prompt-friendly schema description.
+
+    This generates a human-readable description of the expected JSON
+    structure that can be included in system prompts.
+
+    Args:
+        model: A Pydantic BaseModel class
+
+    Returns:
+        A string describing the expected JSON format
+    """
+    schema = generate_json_schema(model)
+
+    lines = ["You must respond with valid JSON matching this schema:", ""]
+    lines.append("```json")
+    lines.append(json.dumps(schema, indent=2))
+    lines.append("```")
+    lines.append("")
+    lines.append("Important:")
+    lines.append("- Respond ONLY with the JSON object, no other text")
+    lines.append("- Ensure all required fields are present")
+    lines.append("- Use the exact field names and types specified")
+
+    return "\n".join(lines)
+
+
+def get_field_descriptions(model: type[BaseModel]) -> dict[str, str]:
+    """
+    Extract field descriptions from a Pydantic model.
+
+    Args:
+        model: A Pydantic BaseModel class
+
+    Returns:
+        Dict mapping field names to their descriptions
+    """
+    descriptions = {}
+    for field_name, field_info in model.model_fields.items():
+        if field_info.description:
+            descriptions[field_name] = field_info.description
+    return descriptions
+
+
+def format_schema_for_openai(model: type[BaseModel]) -> dict[str, Any]:
+    """
+    Format a schema for OpenAI's structured outputs feature.
+
+    OpenAI's structured outputs (response_format with json_schema)
+    requires a specific format with name and strict fields.
+
+    Args:
+        model: A Pydantic BaseModel class
+
+    Returns:
+        Dict formatted for OpenAI's response_format parameter
+    """
+    schema = model.model_json_schema()
+
+    return {
+        "type": "json_schema",
+        "json_schema": {
+            "name": model.__name__,
+            "strict": True,
+            "schema": schema,
+        },
+    }
+
+
+def is_supported_type(model: type[BaseModel]) -> bool:
+    """
+    Check if a Pydantic model uses only supported types for structured outputs.
+
+    Most types are supported, but this can help identify potential issues.
+
+    Args:
+        model: A Pydantic BaseModel class
+
+    Returns:
+        True if all field types are supported
+    """
+    # For now, we support all types that Pydantic can serialize to JSON schema
+    # This could be expanded to check for specific unsupported types
+    try:
+        model.model_json_schema()
+        return True
+    except Exception:
+        return False

flashlite/templating/__init__.py

@@ -0,0 +1,11 @@
+"""Jinja templating for prompts."""
+
+from .engine import TemplateEngine
+from .filters import register_default_filters
+from .registry import TemplateRegistry
+
+__all__ = [
+    "TemplateEngine",
+    "TemplateRegistry",
+    "register_default_filters",
+]

flashlite/templating/engine.py

@@ -0,0 +1,217 @@
+"""Jinja template engine for prompt templating."""
+
+from pathlib import Path
+from typing import Any
+
+from jinja2 import (
+    BaseLoader,
+    Environment,
+    FileSystemLoader,
+    StrictUndefined,
+    Template,
+    TemplateNotFound,
+    Undefined,
+    UndefinedError,
+)
+
+from ..types import TemplateError
+from .filters import register_default_filters
+from .registry import TemplateRegistry
+
+
+class RegistryLoader(BaseLoader):
+    """Jinja loader that loads templates from a TemplateRegistry."""
+
+    def __init__(self, registry: TemplateRegistry):
+        self.registry = registry
+
+    def get_source(
+        self, environment: Environment, template: str
+    ) -> tuple[str, str | None, Any]:
+        if not self.registry.has(template):
+            raise TemplateNotFound(template)
+        source = self.registry.get_source(template)
+        return source, template, lambda: True
+
+
+class TemplateEngine:
+    """
+    Main template engine for rendering prompts.
+
+    Features:
+    - File-based templates from a directory
+    - In-memory template registry
+    - Custom filters for common operations
+    - Strict undefined checking (errors on missing variables)
+    - Variable validation
+    """
+
+    def __init__(
+        self,
+        template_dir: Path | str | None = None,
+        strict: bool = True,
+    ):
+        """
+        Initialize the template engine.
+
+        Args:
+            template_dir: Optional directory to load templates from
+            strict: If True, raise error on undefined variables
+        """
+        self.registry = TemplateRegistry()
+        self._template_dir = Path(template_dir) if template_dir else None
+
+        # Create loaders
+        loaders: list[BaseLoader] = [RegistryLoader(self.registry)]
+        if self._template_dir and self._template_dir.exists():
+            loaders.append(FileSystemLoader(str(self._template_dir)))
+
+        # Create Jinja environment
+        # Using ChoiceLoader to try registry first, then filesystem
+        from jinja2 import ChoiceLoader
+
+        self.env = Environment(
+            loader=ChoiceLoader(loaders),
+            undefined=StrictUndefined if strict else Undefined,
+            autoescape=False,  # Don't HTML-escape (we're making prompts, not HTML)
+            trim_blocks=True,
+            lstrip_blocks=True,
+        )
+
+        # Register default filters
+        register_default_filters(self.env)
+
+        # Load templates from directory if provided
+        if self._template_dir and self._template_dir.exists():
+            self.registry.load_from_directory(self._template_dir)
+
+    def render(
+        self,
+        template: str,
+        variables: dict[str, Any] | None = None,
+        validate: bool = True,
+    ) -> str:
+        """
+        Render a template with variables.
+
+        Args:
+            template: Template name or inline template string
+            variables: Variables to pass to the template
+            validate: If True, validate that all variables are provided
+
+        Returns:
+            Rendered template string
+
+        Raises:
+            TemplateError: If template not found or rendering fails
+        """
+        variables = variables or {}
+
+        try:
+            # Try to get from registry/filesystem first
+            try:
+                tpl = self.env.get_template(template)
+            except TemplateNotFound:
+                # Treat as inline template string
+                tpl = self.env.from_string(template)
+
+            if validate:
+                self._validate_variables(tpl, variables)
+
+            return tpl.render(**variables)
+
+        except UndefinedError as e:
+            raise TemplateError(f"Missing template variable: {e}") from e
+        except Exception as e:
+            raise TemplateError(f"Template rendering failed: {e}") from e
+
+    def render_string(
+        self,
+        template_string: str,
+        variables: dict[str, Any] | None = None,
+    ) -> str:
+        """
+        Render an inline template string.
+
+        Args:
+            template_string: The template source string
+            variables: Variables to pass to the template
+
+        Returns:
+            Rendered string
+        """
+        variables = variables or {}
+        try:
+            tpl = self.env.from_string(template_string)
+            return tpl.render(**variables)
+        except UndefinedError as e:
+            raise TemplateError(f"Missing template variable: {e}") from e
+        except Exception as e:
+            raise TemplateError(f"Template rendering failed: {e}") from e
+
+    def register(self, name: str, template: str) -> None:
+        """
+        Register an in-memory template.
+
+        Args:
+            name: Template name
+            template: Template source string
+        """
+        self.registry.register(name, template)
+
+    def get_variables(self, template: str) -> set[str]:
+        """
+        Get the variables used in a template.
+
+        Note: This uses AST analysis and may not catch all dynamic variables.
+
+        Args:
+            template: Template name or inline template string
+
+        Returns:
+            Set of variable names
+        """
+        from jinja2 import meta
+
+        try:
+            try:
+                self.env.get_template(template)  # Validate template exists
+                source = self.registry.get_source(template)
+            except TemplateNotFound:
+                source = template
+
+            ast = self.env.parse(source)
+            return meta.find_undeclared_variables(ast)
+        except Exception:
+            return set()
+
+    def _validate_variables(self, template: Template, variables: dict[str, Any]) -> None:
+        """Validate that all required variables are provided."""
+        from jinja2 import meta
+
+        if not hasattr(template, "source"):
+            return
+
+        try:
+            ast = self.env.parse(template.source)
+            required = meta.find_undeclared_variables(ast)
+            provided = set(variables.keys())
+            missing = required - provided
+
+            if missing:
+                raise TemplateError(
+                    f"Missing required template variables: {', '.join(sorted(missing))}"
+                )
+        except TemplateError:
+            raise
+        except Exception:
+            # If we can't parse, let rendering catch the error
+            pass
+
+    def add_filter(self, name: str, func: Any) -> None:
+        """Add a custom filter to the environment."""
+        self.env.filters[name] = func
+
+    def add_global(self, name: str, value: Any) -> None:
+        """Add a global variable available in all templates."""
+        self.env.globals[name] = value