ai-lib-python 0.5.0__py3-none-any.whl

ai_lib_python/structured/schema.py

@@ -0,0 +1,316 @@
+"""
+JSON Schema generation utilities.
+
+Provides schema generation from Python types and Pydantic models.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, get_args, get_origin
+
+
+def json_schema_from_type(python_type: type) -> dict[str, Any]:
+    """Generate JSON schema from a Python type.
+
+    Args:
+        python_type: Python type to convert
+
+    Returns:
+        JSON schema dictionary
+
+    Example:
+        >>> schema = json_schema_from_type(str)
+        >>> print(schema)
+        {"type": "string"}
+
+        >>> from typing import List
+        >>> schema = json_schema_from_type(List[int])
+        >>> print(schema)
+        {"type": "array", "items": {"type": "integer"}}
+    """
+    # Handle None
+    if python_type is type(None):
+        return {"type": "null"}
+
+    # Handle basic types
+    type_mapping = {
+        str: {"type": "string"},
+        int: {"type": "integer"},
+        float: {"type": "number"},
+        bool: {"type": "boolean"},
+        bytes: {"type": "string", "format": "byte"},
+    }
+
+    if python_type in type_mapping:
+        return type_mapping[python_type]
+
+    # Handle generic types
+    origin = get_origin(python_type)
+    args = get_args(python_type)
+
+    # Handle list/List
+    if origin is list:
+        if args:
+            return {
+                "type": "array",
+                "items": json_schema_from_type(args[0]),
+            }
+        return {"type": "array"}
+
+    # Handle dict/Dict
+    if origin is dict:
+        schema: dict[str, Any] = {"type": "object"}
+        if len(args) >= 2:
+            schema["additionalProperties"] = json_schema_from_type(args[1])
+        return schema
+
+    # Handle tuple/Tuple
+    if origin is tuple:
+        if args:
+            return {
+                "type": "array",
+                "items": [json_schema_from_type(arg) for arg in args],
+                "minItems": len(args),
+                "maxItems": len(args),
+            }
+        return {"type": "array"}
+
+    # Handle Union types (including | syntax)
+    if origin is type(int | str):  # Python 3.10+ union
+        schemas = [json_schema_from_type(arg) for arg in args]
+        # Check if it's Optional (Union with None)
+        none_schemas = [s for s in schemas if s.get("type") == "null"]
+        other_schemas = [s for s in schemas if s.get("type") != "null"]
+
+        if none_schemas and len(other_schemas) == 1:
+            # Optional type
+            return {**other_schemas[0], "nullable": True}
+        return {"anyOf": schemas}
+
+    # Handle Any
+    if python_type is Any:
+        return {}
+
+    # Try to handle as Pydantic model
+    if hasattr(python_type, "model_json_schema"):
+        return json_schema_from_pydantic(python_type)
+
+    # Default to object
+    return {"type": "object"}
+
+
+def json_schema_from_pydantic(model: type) -> dict[str, Any]:
+    """Generate JSON schema from a Pydantic model.
+
+    Args:
+        model: Pydantic model class
+
+    Returns:
+        JSON schema dictionary
+
+    Raises:
+        ValueError: If model is not a Pydantic model
+
+    Example:
+        >>> from pydantic import BaseModel
+        >>> class User(BaseModel):
+        ...     name: str
+        ...     age: int
+        >>> schema = json_schema_from_pydantic(User)
+    """
+    if not hasattr(model, "model_json_schema"):
+        raise ValueError(f"{model} is not a Pydantic model")
+
+    return model.model_json_schema()
+
+
+class SchemaGenerator:
+    """Generator for JSON schemas with customization options.
+
+    Example:
+        >>> generator = SchemaGenerator()
+        >>> generator.add_property("name", str, description="User's name")
+        >>> generator.add_property("age", int, minimum=0)
+        >>> schema = generator.build()
+    """
+
+    def __init__(
+        self,
+        title: str | None = None,
+        description: str | None = None,
+    ) -> None:
+        """Initialize schema generator.
+
+        Args:
+            title: Schema title
+            description: Schema description
+        """
+        self._title = title
+        self._description = description
+        self._properties: dict[str, dict[str, Any]] = {}
+        self._required: list[str] = []
+        self._additional_properties: bool | dict[str, Any] = False
+
+    def add_property(
+        self,
+        name: str,
+        python_type: type,
+        *,
+        description: str | None = None,
+        required: bool = True,
+        default: Any = None,
+        enum: list[Any] | None = None,
+        minimum: float | None = None,
+        maximum: float | None = None,
+        min_length: int | None = None,
+        max_length: int | None = None,
+        pattern: str | None = None,
+    ) -> SchemaGenerator:
+        """Add a property to the schema.
+
+        Args:
+            name: Property name
+            python_type: Property type
+            description: Property description
+            required: Whether property is required
+            default: Default value
+            enum: Allowed values
+            minimum: Minimum value (for numbers)
+            maximum: Maximum value (for numbers)
+            min_length: Minimum length (for strings)
+            max_length: Maximum length (for strings)
+            pattern: Regex pattern (for strings)
+
+        Returns:
+            Self for chaining
+        """
+        prop_schema = json_schema_from_type(python_type)
+
+        if description:
+            prop_schema["description"] = description
+        if default is not None:
+            prop_schema["default"] = default
+        if enum:
+            prop_schema["enum"] = enum
+        if minimum is not None:
+            prop_schema["minimum"] = minimum
+        if maximum is not None:
+            prop_schema["maximum"] = maximum
+        if min_length is not None:
+            prop_schema["minLength"] = min_length
+        if max_length is not None:
+            prop_schema["maxLength"] = max_length
+        if pattern:
+            prop_schema["pattern"] = pattern
+
+        self._properties[name] = prop_schema
+
+        if required:
+            self._required.append(name)
+
+        return self
+
+    def add_object_property(
+        self,
+        name: str,
+        nested_schema: dict[str, Any],
+        *,
+        description: str | None = None,
+        required: bool = True,
+    ) -> SchemaGenerator:
+        """Add a nested object property.
+
+        Args:
+            name: Property name
+            nested_schema: Nested JSON schema
+            description: Property description
+            required: Whether property is required
+
+        Returns:
+            Self for chaining
+        """
+        prop_schema = nested_schema.copy()
+        if description:
+            prop_schema["description"] = description
+
+        self._properties[name] = prop_schema
+
+        if required:
+            self._required.append(name)
+
+        return self
+
+    def allow_additional_properties(
+        self, allowed: bool | type = True
+    ) -> SchemaGenerator:
+        """Configure additional properties.
+
+        Args:
+            allowed: True to allow any, False to disallow, or type to restrict
+
+        Returns:
+            Self for chaining
+        """
+        if isinstance(allowed, bool):
+            self._additional_properties = allowed
+        else:
+            self._additional_properties = json_schema_from_type(allowed)
+        return self
+
+    def build(self) -> dict[str, Any]:
+        """Build the JSON schema.
+
+        Returns:
+            JSON schema dictionary
+        """
+        schema: dict[str, Any] = {
+            "type": "object",
+            "properties": self._properties,
+        }
+
+        if self._title:
+            schema["title"] = self._title
+        if self._description:
+            schema["description"] = self._description
+        if self._required:
+            schema["required"] = self._required
+        if self._additional_properties is not True:
+            schema["additionalProperties"] = self._additional_properties
+
+        return schema
+
+    def to_json(self, indent: int = 2) -> str:
+        """Convert schema to JSON string.
+
+        Args:
+            indent: JSON indentation
+
+        Returns:
+            JSON string
+        """
+        return json.dumps(self.build(), indent=indent)
+
+    @classmethod
+    def from_pydantic(cls, model: type) -> SchemaGenerator:
+        """Create generator from Pydantic model.
+
+        Args:
+            model: Pydantic model class
+
+        Returns:
+            SchemaGenerator instance
+        """
+        generator = cls(
+            title=getattr(model, "__name__", None),
+            description=getattr(model, "__doc__", None),
+        )
+
+        # Get schema from Pydantic
+        schema = json_schema_from_pydantic(model)
+
+        # Copy properties
+        generator._properties = schema.get("properties", {})
+        generator._required = schema.get("required", [])
+
+        return generator

ai_lib_python/structured/validator.py

@@ -0,0 +1,334 @@
+"""
+Output validation for structured responses.
+
+Validates JSON output against schemas and Pydantic models.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any, TypeVar
+
+T = TypeVar("T")
+
+
+class ValidationError(Exception):
+    """Error raised when validation fails.
+
+    Attributes:
+        message: Error message
+        path: JSON path to the error location
+        value: The invalid value
+    """
+
+    def __init__(
+        self,
+        message: str,
+        path: str | None = None,
+        value: Any = None,
+    ) -> None:
+        """Initialize validation error.
+
+        Args:
+            message: Error message
+            path: JSON path to error
+            value: Invalid value
+        """
+        self.message = message
+        self.path = path
+        self.value = value
+        super().__init__(self._format_message())
+
+    def _format_message(self) -> str:
+        """Format the error message."""
+        msg = self.message
+        if self.path:
+            msg = f"{self.path}: {msg}"
+        return msg
+
+
+@dataclass
+class ValidationResult:
+    """Result of validation.
+
+    Attributes:
+        valid: Whether validation passed
+        errors: List of validation errors
+        data: Validated/parsed data
+    """
+
+    valid: bool = True
+    errors: list[str] = field(default_factory=list)
+    data: Any = None
+
+    def __bool__(self) -> bool:
+        """Return True if validation passed."""
+        return self.valid
+
+    def raise_if_invalid(self) -> None:
+        """Raise ValidationError if validation failed."""
+        if not self.valid:
+            raise ValidationError("; ".join(self.errors))
+
+
+class OutputValidator:
+    """Validator for structured output.
+
+    Validates JSON strings and dictionaries against schemas
+    or Pydantic models.
+
+    Example:
+        >>> from pydantic import BaseModel
+        >>> class User(BaseModel):
+        ...     name: str
+        ...     age: int
+        >>>
+        >>> validator = OutputValidator(User)
+        >>> result = validator.validate('{"name": "Alice", "age": 30}')
+        >>> print(result.valid)  # True
+        >>> print(result.data)   # User(name='Alice', age=30)
+    """
+
+    def __init__(
+        self,
+        schema: dict[str, Any] | type | None = None,
+        strict: bool = True,
+    ) -> None:
+        """Initialize validator.
+
+        Args:
+            schema: JSON schema dict or Pydantic model class
+            strict: Whether to use strict validation
+        """
+        self._pydantic_model: type | None = None
+        self._json_schema: dict[str, Any] | None = None
+        self._strict = strict
+
+        if schema is not None:
+            if isinstance(schema, dict):
+                self._json_schema = schema
+            elif hasattr(schema, "model_validate"):
+                # Pydantic model
+                self._pydantic_model = schema
+                self._json_schema = schema.model_json_schema()
+            else:
+                raise ValueError(
+                    "Schema must be a JSON schema dict or Pydantic model class"
+                )
+
+    def validate(self, data: str | dict[str, Any]) -> ValidationResult:
+        """Validate data against the schema.
+
+        Args:
+            data: JSON string or dictionary to validate
+
+        Returns:
+            ValidationResult with validation status and parsed data
+        """
+        # Parse JSON if string
+        if isinstance(data, str):
+            try:
+                parsed = json.loads(data)
+            except json.JSONDecodeError as e:
+                return ValidationResult(
+                    valid=False,
+                    errors=[f"Invalid JSON: {e}"],
+                )
+        else:
+            parsed = data
+
+        # Validate against Pydantic model if available
+        if self._pydantic_model is not None:
+            return self._validate_pydantic(parsed)
+
+        # Validate against JSON schema if available
+        if self._json_schema is not None:
+            return self._validate_json_schema(parsed)
+
+        # No schema, just return parsed data
+        return ValidationResult(valid=True, data=parsed)
+
+    def validate_or_raise(self, data: str | dict[str, Any]) -> Any:
+        """Validate data and raise if invalid.
+
+        Args:
+            data: Data to validate
+
+        Returns:
+            Validated/parsed data
+
+        Raises:
+            ValidationError: If validation fails
+        """
+        result = self.validate(data)
+        result.raise_if_invalid()
+        return result.data
+
+    def parse(self, data: str | dict[str, Any], model: type[T]) -> T:
+        """Parse and validate data into a Pydantic model.
+
+        Args:
+            data: Data to parse
+            model: Pydantic model class
+
+        Returns:
+            Model instance
+
+        Raises:
+            ValidationError: If validation fails
+        """
+        if isinstance(data, str):
+            try:
+                parsed = json.loads(data)
+            except json.JSONDecodeError as e:
+                raise ValidationError(f"Invalid JSON: {e}") from e
+        else:
+            parsed = data
+
+        try:
+            return model.model_validate(parsed)
+        except Exception as e:
+            raise ValidationError(str(e)) from e
+
+    def _validate_pydantic(self, data: dict[str, Any]) -> ValidationResult:
+        """Validate against Pydantic model.
+
+        Args:
+            data: Data to validate
+
+        Returns:
+            ValidationResult
+        """
+        try:
+            validated = self._pydantic_model.model_validate(data)
+            return ValidationResult(valid=True, data=validated)
+        except Exception as e:
+            return ValidationResult(
+                valid=False,
+                errors=[str(e)],
+            )
+
+    def _validate_json_schema(self, data: dict[str, Any]) -> ValidationResult:
+        """Validate against JSON schema.
+
+        Args:
+            data: Data to validate
+
+        Returns:
+            ValidationResult
+        """
+        errors: list[str] = []
+
+        # Basic type checking
+        schema_type = self._json_schema.get("type")
+        if schema_type == "object" and not isinstance(data, dict):
+            return ValidationResult(
+                valid=False,
+                errors=[f"Expected object, got {type(data).__name__}"],
+            )
+
+        # Validate required properties
+        required = self._json_schema.get("required", [])
+        for prop in required:
+            if prop not in data:
+                errors.append(f"Missing required property: {prop}")
+
+        # Validate property types
+        properties = self._json_schema.get("properties", {})
+        for prop_name, prop_schema in properties.items():
+            if prop_name in data:
+                prop_errors = self._validate_property(
+                    data[prop_name], prop_schema, prop_name
+                )
+                errors.extend(prop_errors)
+
+        # Check additional properties
+        additional_props = self._json_schema.get("additionalProperties", True)
+        if additional_props is False:
+            extra_props = set(data.keys()) - set(properties.keys())
+            for prop in extra_props:
+                errors.append(f"Additional property not allowed: {prop}")
+
+        return ValidationResult(
+            valid=len(errors) == 0,
+            errors=errors,
+            data=data if len(errors) == 0 else None,
+        )
+
+    def _validate_property(
+        self,
+        value: Any,
+        schema: dict[str, Any],
+        path: str,
+    ) -> list[str]:
+        """Validate a single property.
+
+        Args:
+            value: Property value
+            schema: Property schema
+            path: Property path
+
+        Returns:
+            List of error messages
+        """
+        errors: list[str] = []
+        prop_type = schema.get("type")
+
+        # Type validation
+        type_checks = {
+            "string": lambda v: isinstance(v, str),
+            "integer": lambda v: isinstance(v, int) and not isinstance(v, bool),
+            "number": lambda v: isinstance(v, (int, float)) and not isinstance(v, bool),
+            "boolean": lambda v: isinstance(v, bool),
+            "array": lambda v: isinstance(v, list),
+            "object": lambda v: isinstance(v, dict),
+            "null": lambda v: v is None,
+        }
+
+        if prop_type and prop_type in type_checks:
+            # Handle nullable
+            if schema.get("nullable") and value is None:
+                return []
+
+            if not type_checks[prop_type](value):
+                errors.append(f"{path}: Expected {prop_type}, got {type(value).__name__}")
+                return errors
+
+        # String constraints
+        if prop_type == "string" and isinstance(value, str):
+            if "minLength" in schema and len(value) < schema["minLength"]:
+                errors.append(f"{path}: String too short (min {schema['minLength']})")
+            if "maxLength" in schema and len(value) > schema["maxLength"]:
+                errors.append(f"{path}: String too long (max {schema['maxLength']})")
+            if "pattern" in schema:
+                import re
+
+                if not re.match(schema["pattern"], value):
+                    errors.append(f"{path}: String does not match pattern")
+
+        # Number constraints
+        if prop_type in ("integer", "number") and isinstance(value, (int, float)):
+            if "minimum" in schema and value < schema["minimum"]:
+                errors.append(f"{path}: Value below minimum ({schema['minimum']})")
+            if "maximum" in schema and value > schema["maximum"]:
+                errors.append(f"{path}: Value above maximum ({schema['maximum']})")
+
+        # Enum constraint
+        if "enum" in schema and value not in schema["enum"]:
+            errors.append(f"{path}: Value not in allowed enum values")
+
+        # Array validation
+        if prop_type == "array" and isinstance(value, list):
+            if "minItems" in schema and len(value) < schema["minItems"]:
+                errors.append(f"{path}: Array too short (min {schema['minItems']})")
+            if "maxItems" in schema and len(value) > schema["maxItems"]:
+                errors.append(f"{path}: Array too long (max {schema['maxItems']})")
+            if "items" in schema:
+                for i, item in enumerate(value):
+                    item_errors = self._validate_property(
+                        item, schema["items"], f"{path}[{i}]"
+                    )
+                    errors.extend(item_errors)
+
+        return errors