apple-foundation-models 0.2.2__cp312-cp312-macosx_26_0_universal2.whl

applefoundationmodels/error_codes.json

@@ -0,0 +1,163 @@
+[
+  {
+    "code": -1,
+    "name": "InitializationError",
+    "swift_case": null,
+    "parent": "FoundationModelsError",
+    "description": "Library initialization failed."
+  },
+  {
+    "code": -2,
+    "name": "NotAvailableError",
+    "swift_case": "errorNotAvailable",
+    "parent": "FoundationModelsError",
+    "description": "Apple Intelligence not available on this device."
+  },
+  {
+    "code": -3,
+    "name": "InvalidParametersError",
+    "swift_case": "errorInvalidParams",
+    "parent": "FoundationModelsError",
+    "description": "Invalid parameters provided to function."
+  },
+  {
+    "code": -4,
+    "name": "MemoryError",
+    "swift_case": "errorMemory",
+    "parent": "FoundationModelsError",
+    "description": "Memory allocation error."
+  },
+  {
+    "code": -5,
+    "name": "JSONParseError",
+    "swift_case": "errorJSONParse",
+    "parent": "FoundationModelsError",
+    "description": "JSON parsing or validation error."
+  },
+  {
+    "code": -6,
+    "name": "GenerationError",
+    "swift_case": "errorGeneration",
+    "parent": "FoundationModelsError",
+    "description": "Text generation error."
+  },
+  {
+    "code": -7,
+    "name": "TimeoutError",
+    "swift_case": "errorTimeout",
+    "parent": "FoundationModelsError",
+    "description": "Operation timeout."
+  },
+  {
+    "code": -8,
+    "name": "SessionNotFoundError",
+    "swift_case": null,
+    "parent": "FoundationModelsError",
+    "description": "Session ID not found."
+  },
+  {
+    "code": -9,
+    "name": "StreamNotFoundError",
+    "swift_case": null,
+    "parent": "FoundationModelsError",
+    "description": "Stream ID not found or already completed."
+  },
+  {
+    "code": -10,
+    "name": "GuardrailViolationError",
+    "swift_case": "errorGuardrailViolation",
+    "parent": "FoundationModelsError",
+    "description": "Content blocked by safety filters."
+  },
+  {
+    "code": -11,
+    "name": "ToolNotFoundError",
+    "swift_case": "errorToolNotFound",
+    "parent": "FoundationModelsError",
+    "description": "Tool callback not registered for session."
+  },
+  {
+    "code": -12,
+    "name": "ToolExecutionError",
+    "swift_case": "errorToolExecution",
+    "parent": "FoundationModelsError",
+    "description": "Tool execution failed or returned invalid result."
+  },
+  {
+    "code": -13,
+    "name": "BufferTooSmallError",
+    "swift_case": null,
+    "parent": "FoundationModelsError",
+    "description": "Internal buffer too small to serialize error message."
+  },
+  {
+    "code": -14,
+    "name": "ContextWindowExceededError",
+    "swift_case": "errorContextWindowExceeded",
+    "parent": "GenerationError",
+    "description": "Context window size limit exceeded (4096 tokens)."
+  },
+  {
+    "code": -15,
+    "name": "DecodingFailureError",
+    "swift_case": "errorDecodingFailure",
+    "parent": "GenerationError",
+    "description": "Failed to deserialize a valid generable type from model output."
+  },
+  {
+    "code": -16,
+    "name": "RateLimitedError",
+    "swift_case": "errorRateLimited",
+    "parent": "GenerationError",
+    "description": "Session has been rate limited."
+  },
+  {
+    "code": -17,
+    "name": "RefusalError",
+    "swift_case": "errorRefusal",
+    "parent": "GenerationError",
+    "description": "Session refused the request."
+  },
+  {
+    "code": -18,
+    "name": "ConcurrentRequestsError",
+    "swift_case": "errorConcurrentRequests",
+    "parent": "GenerationError",
+    "description": "Attempted to make a session respond to a second prompt while responding to the first."
+  },
+  {
+    "code": -19,
+    "name": "UnsupportedGuideError",
+    "swift_case": "errorUnsupportedGuide",
+    "parent": "GenerationError",
+    "description": "Generation guide with an unsupported pattern was used."
+  },
+  {
+    "code": -20,
+    "name": "UnsupportedLanguageError",
+    "swift_case": "errorUnsupportedLanguage",
+    "parent": "GenerationError",
+    "description": "Model was prompted to respond in a language that it does not support."
+  },
+  {
+    "code": -21,
+    "name": "AssetsUnavailableError",
+    "swift_case": "errorAssetsUnavailable",
+    "parent": "GenerationError",
+    "description": "Assets required for the session are unavailable."
+  },
+  {
+    "code": -98,
+    "name": "ToolCallError",
+    "swift_case": null,
+    "parent": "FoundationModelsError",
+    "description": "Tool call error (validation, schema, etc.)."
+  },
+  {
+    "code": -99,
+    "name": "UnknownError",
+    "swift_case": "errorUnknown",
+    "parent": "FoundationModelsError",
+    "description": "Unknown error occurred."
+  }
+]

applefoundationmodels/exceptions.py

@@ -0,0 +1,122 @@
+"""
+Exception hierarchy for applefoundationmodels Python bindings.
+
+Exceptions are generated from a single error_codes.json definition so both
+Python and Swift share the same source of truth for error metadata.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from importlib import resources
+from typing import Any, Dict, List, Optional, Tuple, Type
+
+__all__ = [
+    "FoundationModelsError",
+    "GenerationError",
+    "ERROR_CODE_TO_EXCEPTION",
+    "raise_for_error_code",
+    "get_error_definitions",
+]
+
+
+class FoundationModelsError(Exception):
+    """Base exception for all FoundationModels errors."""
+
+    def __init__(self, message: str, error_code: Optional[int] = None):
+        super().__init__(message)
+        self.message = message
+        self.error_code = error_code
+
+
+class GenerationError(FoundationModelsError):
+    """Text generation error."""
+
+    pass
+
+
+@dataclass(frozen=True)
+class ErrorDefinition:
+    """Structured representation of an error entry from the JSON definition."""
+
+    code: int
+    name: str
+    parent: str
+    description: str
+    swift_case: Optional[str]
+
+
+def _load_error_definitions() -> Tuple[ErrorDefinition, ...]:
+    data_path = resources.files(__package__).joinpath("error_codes.json")
+    with data_path.open("r", encoding="utf-8") as f:
+        raw: List[Dict[str, Any]] = json.load(f)
+
+    definitions: List[ErrorDefinition] = []
+    for entry in raw:
+        definitions.append(
+            ErrorDefinition(
+                code=int(entry["code"]),
+                name=str(entry["name"]),
+                parent=str(entry.get("parent", "FoundationModelsError")),
+                description=str(entry["description"]),
+                swift_case=entry.get("swift_case") or None,
+            )
+        )
+    return tuple(definitions)
+
+
+_ERROR_DEFINITIONS = _load_error_definitions()
+
+_PARENT_CLASS_MAP: Dict[str, Type[FoundationModelsError]] = {
+    "FoundationModelsError": FoundationModelsError,
+    "GenerationError": GenerationError,
+}
+
+
+def get_error_definitions() -> Tuple[ErrorDefinition, ...]:
+    """Return the error definitions used to build the exception hierarchy."""
+
+    return _ERROR_DEFINITIONS
+
+
+ERROR_CODE_TO_EXCEPTION: Dict[int, Type[FoundationModelsError]] = {}
+
+
+def _register_exception_classes() -> None:
+    for definition in _ERROR_DEFINITIONS:
+        parent_class = _PARENT_CLASS_MAP.get(definition.parent, FoundationModelsError)
+
+        if definition.name == "GenerationError":
+            ERROR_CODE_TO_EXCEPTION[definition.code] = GenerationError
+            continue
+
+        exception_class = type(
+            definition.name,
+            (parent_class,),
+            {"__doc__": definition.description},
+        )
+        globals()[definition.name] = exception_class
+        __all__.append(definition.name)
+        ERROR_CODE_TO_EXCEPTION[definition.code] = exception_class
+
+
+_register_exception_classes()
+
+
+def raise_for_error_code(error_code: int, message: str) -> None:
+    """
+    Raise the appropriate exception for a given error code.
+
+    Args:
+        error_code: The error code from the Swift API
+        message: Error message to include in the exception
+
+    Raises:
+        FoundationModelsError: The appropriate exception subclass for the error code
+    """
+
+    exception_class = ERROR_CODE_TO_EXCEPTION.get(
+        error_code, ERROR_CODE_TO_EXCEPTION.get(-99, FoundationModelsError)
+    )
+    raise exception_class(message, error_code)

applefoundationmodels/py.typed

@@ -0,0 +1,2 @@
+# Marker file for PEP 561
+# This package supports type checking

applefoundationmodels/pydantic_compat.py

@@ -0,0 +1,144 @@
+"""
+Pydantic compatibility utilities for applefoundationmodels.
+
+Provides optional Pydantic integration for structured output generation.
+"""
+
+from typing import Any, Dict, Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+# Try to import Pydantic, but don't fail if it's not installed
+try:
+    from pydantic import BaseModel
+
+    PYDANTIC_AVAILABLE = True
+except ImportError:
+    PYDANTIC_AVAILABLE = False
+    BaseModel = None  # type: ignore
+
+
+def require_pydantic() -> None:
+    """
+    Raise ImportError if Pydantic is not available.
+
+    Raises:
+        ImportError: If Pydantic is not installed
+    """
+    if not PYDANTIC_AVAILABLE:
+        raise ImportError(
+            "Pydantic is not installed. Install it with: pip install pydantic>=2.0"
+        )
+
+
+def model_to_schema(model: "BaseModel") -> Dict[str, Any]:
+    """
+    Convert a Pydantic model to JSON Schema.
+
+    Args:
+        model: Pydantic BaseModel class or instance
+
+    Returns:
+        JSON Schema dictionary
+
+    Raises:
+        ImportError: If Pydantic is not installed
+        ValueError: If model is not a valid Pydantic model
+
+    Example:
+        >>> from pydantic import BaseModel
+        >>> class Person(BaseModel):
+        ...     name: str
+        ...     age: int
+        >>> schema = model_to_schema(Person)
+        >>> print(schema)
+        {'type': 'object', 'properties': {...}, 'required': [...]}
+    """
+    require_pydantic()
+
+    if not hasattr(model, "model_json_schema"):
+        raise ValueError(
+            f"Expected Pydantic BaseModel, got {type(model).__name__}. "
+            "Make sure your model inherits from pydantic.BaseModel"
+        )
+
+    # Get JSON Schema from Pydantic model
+    schema = model.model_json_schema()
+
+    # Clean up schema (remove title, $defs if empty, etc.)
+    schema = _clean_schema(schema)
+
+    return schema
+
+
+def _clean_schema(schema: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Clean up Pydantic-generated schema for better compatibility.
+
+    Removes unnecessary fields like 'title' and simplifies nested definitions.
+    """
+    cleaned = schema.copy()
+
+    # Remove title if present
+    cleaned.pop("title", None)
+
+    # Remove $defs if empty or inline them if simple
+    if "$defs" in cleaned:
+        defs = cleaned["$defs"]
+        if not defs:
+            del cleaned["$defs"]
+
+    return cleaned
+
+
+def is_pydantic_model(obj: Any) -> bool:
+    """
+    Check if an object is a Pydantic model class or instance.
+
+    Args:
+        obj: Object to check
+
+    Returns:
+        True if obj is a Pydantic BaseModel class or instance
+    """
+    if not PYDANTIC_AVAILABLE:
+        return False
+
+    # Check if it has the Pydantic model_json_schema method
+    # This works for both classes and instances
+    if hasattr(obj, "model_json_schema"):
+        return True
+
+    return False
+
+
+def normalize_schema(schema: Union[Dict[str, Any], "BaseModel"]) -> Dict[str, Any]:
+    """
+    Normalize schema input to JSON Schema dict.
+
+    Accepts either a JSON Schema dictionary or a Pydantic model,
+    and returns a JSON Schema dictionary.
+
+    Args:
+        schema: JSON Schema dict or Pydantic BaseModel
+
+    Returns:
+        JSON Schema dictionary
+
+    Raises:
+        TypeError: If schema is neither dict nor Pydantic model
+        ImportError: If Pydantic is needed but not installed
+    """
+    # If it's already a dict, return it
+    if isinstance(schema, dict):
+        return schema
+
+    # If it's a Pydantic model, convert it
+    if is_pydantic_model(schema):
+        return model_to_schema(schema)
+
+    # Otherwise, raise an error
+    raise TypeError(
+        f"Expected JSON Schema dict or Pydantic BaseModel, got {type(schema).__name__}"
+    )

applefoundationmodels/session.py

@@ -0,0 +1,265 @@
+"""
+Session API for applefoundationmodels Python bindings.
+
+Provides session management, text generation, and async streaming support.
+"""
+
+import asyncio
+import json
+import logging
+from typing import (
+    Optional,
+    Dict,
+    Any,
+    Callable,
+    Union,
+    TYPE_CHECKING,
+    List,
+    cast,
+    Iterator,
+    overload,
+    Type,
+)
+from typing_extensions import Literal
+from queue import Queue, Empty
+import threading
+
+from .base_session import BaseSession, StreamQueueItem
+from .types import (
+    GenerationResponse,
+    StreamChunk,
+)
+from .pydantic_compat import normalize_schema
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class Session(BaseSession):
+    """
+    AI session for maintaining conversation state.
+
+    Sessions maintain conversation history and can be configured with tools
+    and instructions. Use as a context manager for automatic cleanup.
+
+    Usage:
+        with Session() as session:
+            response = session.generate("Hello!")
+            print(response.text)
+
+        # With configuration:
+        def get_weather(location: str) -> str:
+            '''Get current weather for a location.'''
+            return f"Weather in {location}: 22°C"
+
+        session = Session(
+            instructions="You are a helpful assistant.",
+            tools=[get_weather]
+        )
+        response = session.generate("What's the weather in Paris?")
+    """
+
+    def _call_ffi(self, func, *args, **kwargs):
+        """Execute FFI call synchronously."""
+        return func(*args, **kwargs)
+
+    def _create_stream_queue_adapter(self) -> BaseSession._StreamQueueAdapter:
+        """Return the queue adapter used by synchronous streaming."""
+        queue: Queue[StreamQueueItem] = Queue()
+
+        def push(item: StreamQueueItem) -> None:
+            queue.put(item)
+
+        def get_sync() -> StreamQueueItem:
+            while True:
+                try:
+                    return queue.get(timeout=0.1)
+                except Empty:
+                    continue
+
+        return BaseSession._StreamQueueAdapter(push=push, get_sync=get_sync)
+
+    def close(self) -> None:
+        """Close the session and cleanup resources."""
+        self._mark_closed()
+
+    # ========================================================================
+    # Type overloads for generate() method
+    #
+    # IMPORTANT: These overloads must be kept in sync with AsyncSession.generate()
+    # in async_session.py. The signatures are identical except for:
+    # - async keyword (Session: def generate() vs AsyncSession: async def generate())
+    # - Return type for streaming (Iterator vs AsyncIterator)
+    #
+    # When modifying these overloads, update both files to maintain consistency.
+    # ========================================================================
+
+    # Type overload for non-streaming text generation
+    @overload
+    def generate(
+        self,
+        prompt: str,
+        schema: None = None,
+        stream: Literal[False] = False,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+    ) -> GenerationResponse: ...
+
+    # Type overload for non-streaming structured generation
+    @overload
+    def generate(
+        self,
+        prompt: str,
+        schema: Union[Dict[str, Any], Type["BaseModel"]],
+        stream: Literal[False] = False,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+    ) -> GenerationResponse: ...
+
+    # Type overload for streaming generation (text only, no structured streaming)
+    @overload
+    def generate(
+        self,
+        prompt: str,
+        schema: None = None,
+        stream: Literal[True] = True,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+    ) -> Iterator[StreamChunk]: ...
+
+    def generate(
+        self,
+        prompt: str,
+        schema: Optional[Union[Dict[str, Any], Type["BaseModel"]]] = None,
+        stream: bool = False,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+    ) -> Union[GenerationResponse, Iterator[StreamChunk]]:
+        """
+        Generate text or structured output, with optional streaming.
+
+        This unified method supports three generation modes:
+        1. Text generation (schema=None, stream=False) -> GenerationResponse
+        2. Structured generation (schema=dict/model, stream=False) -> GenerationResponse
+        3. Streaming generation (schema=None, stream=True) -> Iterator[StreamChunk]
+
+        Args:
+            prompt: Input text prompt
+            schema: Optional JSON schema dict or Pydantic model for structured output
+            stream: If True, return an iterator of chunks instead of complete response
+            temperature: Sampling temperature (0.0-2.0, default: DEFAULT_TEMPERATURE)
+            max_tokens: Maximum tokens to generate (default: DEFAULT_MAX_TOKENS)
+
+        Returns:
+            - GenerationResponse with .text or .parsed property (if stream=False)
+            - Iterator[StreamChunk] yielding content deltas (if stream=True)
+
+        Raises:
+            RuntimeError: If session is closed
+            GenerationError: If generation fails
+            ValueError: If schema is provided with stream=True
+
+        Examples:
+            Text generation:
+                >>> response = session.generate("What is Python?")
+                >>> print(response.text)
+
+            Structured generation:
+                >>> schema = {"type": "object", "properties": {"name": {"type": "string"}}}
+                >>> response = session.generate("Extract name: John Doe", schema=schema)
+                >>> print(response.parsed)
+
+            Streaming generation:
+                >>> for chunk in session.generate("Tell me a story", stream=True):
+                ...     print(chunk.content, end='', flush=True)
+        """
+        plan = self._plan_generate_call(stream, schema, temperature, max_tokens)
+
+        if plan.mode == "stream":
+            # Streaming mode: return Iterator[StreamChunk]
+            return self._generate_stream_impl(prompt, plan.temperature, plan.max_tokens)
+        if plan.mode == "structured" and schema is not None:
+            # Structured generation mode
+            return self._generate_structured_impl(
+                prompt, schema, plan.temperature, plan.max_tokens
+            )
+
+        # Text generation mode
+        return self._generate_text_impl(prompt, plan.temperature, plan.max_tokens)
+
+    def _generate_text_impl(
+        self, prompt: str, temperature: float, max_tokens: int
+    ) -> GenerationResponse:
+        """Internal implementation for text generation."""
+        with self._generation_context() as start_length:
+            text = self._call_ffi(
+                self._ffi.generate,
+                prompt,
+                temperature,
+                max_tokens,
+            )
+            return self._build_generation_response(text, False, start_length)
+
+    def _generate_structured_impl(
+        self,
+        prompt: str,
+        schema: Union[Dict[str, Any], Type["BaseModel"]],
+        temperature: float,
+        max_tokens: int,
+    ) -> GenerationResponse:
+        """Internal implementation for structured generation."""
+        with self._generation_context() as start_length:
+            json_schema = normalize_schema(schema)
+            result = self._call_ffi(
+                self._ffi.generate_structured,
+                prompt,
+                json_schema,
+                temperature,
+                max_tokens,
+            )
+            return self._build_generation_response(result, True, start_length)
+
+    def _generate_stream_impl(
+        self, prompt: str, temperature: float, max_tokens: int
+    ) -> Iterator[StreamChunk]:
+        """Internal implementation for streaming generation."""
+        start_length = self._begin_generation()
+        adapter = self._create_stream_queue_adapter()
+        try:
+            # Use shared streaming implementation from base class
+            yield from self._stream_chunks_impl(
+                prompt, temperature, max_tokens, adapter
+            )
+        finally:
+            self._end_generation(start_length)
+
+    def get_history(self) -> List[Dict[str, Any]]:
+        """
+        Get conversation history.
+
+        Returns:
+            List of message dictionaries with 'role' and 'content' keys
+
+        Example:
+            >>> history = session.get_history()
+            >>> for msg in history:
+            ...     print(f"{msg['role']}: {msg['content']}")
+        """
+        self._check_closed()
+        result = self._call_ffi(self._ffi.get_history)
+        return cast(List[Dict[str, Any]], result)
+
+    def clear_history(self) -> None:
+        """
+        Clear conversation history.
+
+        Removes all messages from the session while keeping the session active.
+        """
+        self._check_closed()
+        self._call_ffi(self._ffi.clear_history)
+        # Reset to current transcript length (may include persistent instructions)
+        self._last_transcript_length = len(self.transcript)
+
+    # Properties inherited from BaseSession (transcript, last_generation_transcript)