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

applefoundationmodels/tools.py

@@ -0,0 +1,285 @@
+"""
+Tool calling utilities for applefoundationmodels.
+
+Provides utilities for extracting JSON schemas from Python functions
+and managing tool registrations.
+"""
+
+import inspect
+from typing import (
+    Callable,
+    Dict,
+    Any,
+    Optional,
+    get_type_hints,
+    get_origin,
+    get_args,
+    Union,
+)
+from .exceptions import ToolCallError
+
+
+# Type mapping for efficient schema generation
+# Maps Python types and their string names to JSON schema types
+_TYPE_MAP = {
+    str: "string",
+    int: "integer",
+    float: "number",
+    bool: "boolean",
+    "str": "string",
+    "int": "integer",
+    "float": "number",
+    "bool": "boolean",
+}
+
+
+def is_optional_type(python_type: Any) -> bool:
+    """
+    Check if a type is Optional (Union[X, None]).
+
+    Args:
+        python_type: Python type hint
+
+    Returns:
+        True if the type is Optional, False otherwise
+    """
+    origin = get_origin(python_type)
+    if origin is Union:
+        args = get_args(python_type)
+        return type(None) in args
+    return False
+
+
+def unwrap_optional(python_type: Any) -> Any:
+    """
+    Extract the non-None type from Optional[X] or Union[X, None].
+
+    Args:
+        python_type: Python type hint that may be Optional
+
+    Returns:
+        The unwrapped type (X from Optional[X])
+    """
+    origin = get_origin(python_type)
+    if origin is Union:
+        args = get_args(python_type)
+        # Filter out None and return the first non-None type
+        non_none_types = [arg for arg in args if arg is not type(None)]
+        if non_none_types:
+            return non_none_types[0]
+    return python_type
+
+
+def _handle_list_type(python_type: Any) -> Dict[str, Any]:
+    """
+    Extract list/array type with optional item schema.
+
+    Args:
+        python_type: List or generic list type hint
+
+    Returns:
+        JSON Schema for array type
+    """
+    args = get_args(python_type)
+    if args:
+        items_schema = python_type_to_json_schema(args[0])
+        return {"type": "array", "items": items_schema}
+    return {"type": "array"}
+
+
+def _handle_dict_type(python_type: Any) -> Dict[str, Any]:
+    """
+    Extract dict/object type with optional value schema.
+
+    Args:
+        python_type: Dict or generic dict type hint
+
+    Returns:
+        JSON Schema for object type
+    """
+    args = get_args(python_type)
+    if len(args) == 2 and args[0] is str:
+        value_schema = python_type_to_json_schema(args[1])
+        return {"type": "object", "additionalProperties": value_schema}
+    return {"type": "object"}
+
+
+def python_type_to_json_schema(python_type: Any) -> Dict[str, Any]:
+    """
+    Convert a Python type hint to a JSON Schema type definition.
+
+    Args:
+        python_type: Python type hint
+
+    Returns:
+        JSON Schema type definition
+
+    Raises:
+        ToolCallError: If type cannot be converted
+    """
+    # Handle None type
+    if python_type is type(None):
+        return {"type": "null"}
+
+    # Handle Optional[X] / Union[X, None] - unwrap and process the inner type
+    if is_optional_type(python_type):
+        inner_type = unwrap_optional(python_type)
+        return python_type_to_json_schema(inner_type)
+
+    # Check basic types and string annotations via unified lookup table
+    if python_type in _TYPE_MAP:
+        return {"type": _TYPE_MAP[python_type]}
+
+    # Get origin for generic types
+    origin = get_origin(python_type)
+
+    # Handle container types
+    if python_type is list or origin is list:
+        return _handle_list_type(python_type)
+
+    if python_type is dict or origin is dict:
+        return _handle_dict_type(python_type)
+
+    # Default fallback for unknown types
+    return {"type": "string"}
+
+
+def extract_function_schema(func: Callable) -> Dict[str, Any]:
+    """
+    Extract JSON Schema from a Python function's signature and docstring.
+
+    Args:
+        func: Python function to extract schema from
+
+    Returns:
+        Dictionary containing:
+        - name: Function name
+        - description: Function description from docstring
+        - parameters: JSON Schema for function parameters
+
+    Raises:
+        ToolCallError: If schema cannot be extracted
+    """
+    try:
+        # Get function signature
+        sig = inspect.signature(func)
+
+        # Get type hints
+        try:
+            type_hints = get_type_hints(func)
+        except Exception:
+            # If type hints fail (e.g., forward references), inspect parameters directly
+            type_hints = {}
+
+        # Extract parameter schemas
+        properties = {}
+        required = []
+
+        for param_name, param in sig.parameters.items():
+            # Skip self and cls parameters
+            if param_name in ("self", "cls"):
+                continue
+
+            # Get type annotation
+            param_type = type_hints.get(param_name, param.annotation)
+
+            # Handle parameters without type hints
+            if param_type is inspect.Parameter.empty:
+                # Default to string type if no annotation
+                param_schema = {"type": "string"}
+            else:
+                param_schema = python_type_to_json_schema(param_type)
+
+            properties[param_name] = param_schema
+
+            # Mark as required if no default value
+            if param.default is inspect.Parameter.empty:
+                required.append(param_name)
+
+        # Build parameters schema
+        parameters_schema = {
+            "type": "object",
+            "properties": properties,
+        }
+
+        if required:
+            parameters_schema["required"] = required
+
+        # Extract description from docstring
+        description = ""
+        if func.__doc__:
+            # Get the first line or paragraph of the docstring
+            lines = func.__doc__.strip().split("\n")
+            description = lines[0].strip()
+
+        # Get function name
+        name = func.__name__
+
+        return {
+            "name": name,
+            "description": description,
+            "parameters": parameters_schema,
+        }
+
+    except Exception as e:
+        raise ToolCallError(
+            f"Failed to extract schema from function '{func.__name__}': {e}", -98
+        ) from e
+
+
+def register_tool_for_function(func: Callable) -> Dict[str, Any]:
+    """
+    Extract schema from a function and prepare it for registration.
+
+    This is used when registering tools with a session. It extracts the
+    schema from the function signature, attaches metadata to the function
+    object, and returns the schema ready for FFI registration.
+
+    Args:
+        func: Function to register as a tool
+
+    Returns:
+        Complete tool schema ready for registration
+
+    Raises:
+        ToolCallError: If schema cannot be extracted
+    """
+    schema = extract_function_schema(func)
+
+    # Use shared helper to attach metadata
+    return attach_tool_metadata(func, schema)
+
+
+def attach_tool_metadata(
+    func: Callable,
+    schema: Dict[str, Any],
+    description: Optional[str] = None,
+    name: Optional[str] = None,
+) -> Dict[str, Any]:
+    """
+    Attach tool metadata to a function and return final schema.
+
+    This is a shared helper used by both the standalone @tool decorator
+    and the Session.tool() method to avoid code duplication.
+
+    Args:
+        func: Function to attach metadata to
+        schema: Base schema from extract_function_schema
+        description: Optional override for description
+        name: Optional override for name
+
+    Returns:
+        Final schema with overrides applied
+    """
+    # Override with provided values
+    if description is not None:
+        schema["description"] = description
+    if name is not None:
+        schema["name"] = name
+
+    # Attach metadata to function
+    func._tool_name = schema["name"]  # type: ignore[attr-defined]
+    func._tool_description = schema["description"]  # type: ignore[attr-defined]
+    func._tool_parameters = schema["parameters"]  # type: ignore[attr-defined]
+
+    return schema

applefoundationmodels/types.py

@@ -0,0 +1,284 @@
+"""
+Type definitions for libai Python bindings.
+
+This module provides TypedDicts, enums, and type aliases for type-safe
+interaction with the library.
+"""
+
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import (
+    TypedDict,
+    Optional,
+    Callable,
+    Any,
+    Union,
+    Dict,
+    List,
+    Type,
+    TYPE_CHECKING,
+    cast,
+)
+from typing_extensions import NotRequired
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+
+class Result(IntEnum):
+    """
+    Result codes for AI operations.
+
+    These codes indicate the success or failure state of library operations.
+    """
+
+    SUCCESS = 0
+    INIT_FAILED = -1
+    NOT_AVAILABLE = -2
+    INVALID_PARAMS = -3
+    MEMORY = -4
+    JSON_PARSE = -5
+    GENERATION = -6
+    TIMEOUT = -7
+    SESSION_NOT_FOUND = -8
+    STREAM_NOT_FOUND = -9
+    GUARDRAIL_VIOLATION = -10
+    TOOL_NOT_FOUND = -11
+    TOOL_EXECUTION = -12
+    BUFFER_TOO_SMALL = -13
+    UNKNOWN = -99
+
+
+class Availability(IntEnum):
+    """
+    Apple Intelligence availability status.
+
+    Indicates whether Apple Intelligence is available and ready for use
+    on the current device and system configuration.
+    """
+
+    AVAILABLE = 1
+    DEVICE_NOT_ELIGIBLE = -1
+    NOT_ENABLED = -2
+    MODEL_NOT_READY = -3
+    AVAILABILITY_UNKNOWN = -99
+
+
+class SessionConfig(TypedDict, total=False):
+    """
+    Session configuration options.
+
+    Configuration for creating an AI session. Sessions maintain conversation
+    state and can be configured with tools and instructions.
+
+    Attributes:
+        instructions: Optional system instructions to guide AI behavior
+    """
+
+    instructions: NotRequired[Optional[str]]
+
+
+class GenerationParams(TypedDict, total=False):
+    """
+    Text generation parameters.
+
+    Controls various aspects of AI text generation including randomness
+    and length limits.
+
+    Attributes:
+        temperature: Generation randomness (0.0 = deterministic, 2.0 = very random)
+        max_tokens: Maximum response tokens (0 = use system default)
+    """
+
+    temperature: NotRequired[float]
+    max_tokens: NotRequired[int]
+
+
+@dataclass
+class Function:
+    """
+    Function call information from a tool call.
+
+    Represents the function that was called, including its name and arguments.
+    Follows OpenAI's pattern for tool call representation.
+
+    Attributes:
+        name: The name of the function that was called
+        arguments: JSON string containing the function arguments
+
+    Example:
+        >>> func = Function(name="get_weather", arguments='{"location": "Paris"}')
+        >>> print(func.name)
+        get_weather
+    """
+
+    name: str
+    arguments: str  # JSON string of arguments
+
+
+@dataclass
+class ToolCall:
+    """
+    A tool call made during generation.
+
+    Represents a single tool/function call that occurred during text generation.
+    Follows OpenAI's pattern where tool calls are exposed directly on the response.
+
+    Attributes:
+        id: Unique identifier for this tool call
+        type: Type of tool call (currently only "function" is supported)
+        function: The function call details (name and arguments)
+
+    Example:
+        >>> tool_call = ToolCall(
+        ...     id="call_123",
+        ...     type="function",
+        ...     function=Function(name="get_weather", arguments='{"location": "Paris"}')
+        ... )
+        >>> print(tool_call.function.name)
+        get_weather
+    """
+
+    id: str
+    type: str  # "function" - matches OpenAI's pattern
+    function: Function
+
+
+@dataclass
+class GenerationResponse:
+    """
+    Response from non-streaming generation.
+
+    Provides a unified interface for both text and structured generation results.
+    Use the .text property for text responses and .parsed for structured outputs.
+
+    Attributes:
+        content: The generated content (str for text, dict for structured)
+        is_structured: True if response is structured JSON, False for text
+        tool_calls: List of tool calls made during generation (None if no tools called)
+        finish_reason: Reason generation stopped ("stop", "tool_calls", "length", etc.)
+        metadata: Optional metadata about the generation
+
+    Example (text):
+        >>> response = session.generate("Hello")
+        >>> print(response.text)
+
+    Example (structured):
+        >>> response = session.generate("Extract name and age", schema={...})
+        >>> data = response.parsed
+        >>> person = Person(**data)  # Parse into Pydantic model
+
+    Example (tool calls):
+        >>> response = session.generate("What's the weather in Paris?")
+        >>> if response.tool_calls:
+        ...     for tool_call in response.tool_calls:
+        ...         print(f"Called {tool_call.function.name}")
+    """
+
+    content: Union[str, Dict[str, Any]]
+    is_structured: bool
+    tool_calls: Optional[List[ToolCall]] = None
+    finish_reason: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+    @property
+    def text(self) -> str:
+        """
+        Get response as text.
+
+        Returns:
+            The generated text
+
+        Raises:
+            ValueError: If response is structured (use .parsed instead)
+        """
+        if self.is_structured:
+            raise ValueError(
+                "Response is structured output. Use .parsed property instead of .text"
+            )
+        return cast(str, self.content)
+
+    @property
+    def parsed(self) -> Dict[str, Any]:
+        """
+        Get response as structured data.
+
+        Returns:
+            The parsed JSON dictionary
+
+        Raises:
+            ValueError: If response is text (use .text instead)
+        """
+        if not self.is_structured:
+            raise ValueError(
+                "Response is text output. Use .text property instead of .parsed"
+            )
+        return cast(Dict[str, Any], self.content)
+
+    def parse_as(self, model: "Type[BaseModel]") -> "BaseModel":
+        """
+        Parse structured response into a Pydantic model.
+
+        Args:
+            model: Pydantic BaseModel class to parse into
+
+        Returns:
+            Instantiated Pydantic model
+
+        Raises:
+            ValueError: If response is not structured
+            ImportError: If pydantic is not installed
+
+        Example:
+            >>> from pydantic import BaseModel
+            >>> class Person(BaseModel):
+            ...     name: str
+            ...     age: int
+            >>> response = session.generate("Extract: Alice is 28", schema=Person)
+            >>> person = response.parse_as(Person)
+            >>> print(person.name, person.age)
+        """
+        return model(**self.parsed)
+
+
+@dataclass
+class StreamChunk:
+    """
+    A chunk from streaming generation.
+
+    Represents a single delta in the streaming response. Multiple chunks
+    combine to form the complete response.
+
+    Attributes:
+        content: The text content delta for this chunk
+        finish_reason: Reason streaming ended (None for intermediate chunks)
+        index: Chunk sequence index (usually 0 for single-stream responses)
+
+    Example:
+        >>> for chunk in session.generate("Tell a story", stream=True):
+        ...     print(chunk.content, end='', flush=True)
+        ...     if chunk.finish_reason:
+        ...         print(f"\\n[Finished: {chunk.finish_reason}]")
+    """
+
+    content: str
+    finish_reason: Optional[str] = None
+    index: int = 0
+
+
+# Callback type aliases
+StreamCallback = Callable[[Optional[str]], None]
+"""
+Callback function for streaming text generation.
+
+Called incrementally during streaming generation for each token or chunk.
+None indicates completion or error.
+"""
+
+ToolCallback = Callable[[dict], Any]
+"""
+Callback function for tool execution.
+
+Receives tool parameters as a dict and should return the tool result.
+The result will be automatically JSON-serialized.
+"""