mirascope 1.25.7__py3-none-any.whl → 2.0.0a0__py3-none-any.whl

mirascope/llm/tools/tool_schema.py

@@ -0,0 +1,246 @@
+"""The `ToolSchema` class for defining tools that LLMs can request be called."""
+
+from __future__ import annotations
+
+import inspect
+from collections import namedtuple
+from dataclasses import dataclass
+from typing import (
+    Annotated,
+    Any,
+    Generic,
+    TypeVar,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+from docstring_parser import parse
+from pydantic import BaseModel, Field, create_model
+from pydantic.fields import FieldInfo
+
+from ..content import ToolCall
+from .protocols import AsyncContextToolFn, AsyncToolFn, ContextToolFn, ToolFn
+
+ToolFnT = TypeVar(
+    "ToolFnT",
+    bound=ToolFn | AsyncToolFn | ContextToolFn | AsyncContextToolFn,
+    covariant=True,
+)
+ToolSchemaT = TypeVar("ToolSchemaT", bound="ToolSchema")
+
+DocstringArg = namedtuple("DocstringArg", ["name", "description"])
+
+FORMAT_TOOL_NAME = "__mirascope_formatted_output_tool__"
+"""Reserved name of the formatted output tool. 
+
+Any call to a tool with this name is NOT considered a regular tool call, but will instead
+be converted into textual output containing the arguments to the tool call.
+"""
+# TODO: Investigate whether using this hardcoded tool name has any adverse impact
+# on model performance, compared to a tool name that references the name of the formatted
+# class.
+
+
+@dataclass
+class ParsedDocstring:
+    args: list[DocstringArg]
+
+
+def _parse_docstring_params(docstring: str | None) -> ParsedDocstring:
+    """Parse parameter descriptions from a docstring.
+
+    Uses docstring-parser library which supports ReST, Google, Numpydoc-style and
+    Epydoc docstrings
+
+    Args:
+        docstring: The function's docstring
+
+    Returns:
+        ParsedDocstring containing parameter descriptions
+    """
+    if not docstring:
+        return ParsedDocstring(args=[])
+
+    parsed = parse(docstring)
+    args = []
+
+    for param in parsed.params:
+        if param.description:
+            args.append(
+                DocstringArg(name=param.arg_name, description=param.description)
+            )
+
+    return ParsedDocstring(args=args)
+
+
+class ToolParameterSchema(BaseModel):
+    """JSON Schema for tool parameters (always an object with properties).
+
+    This contains real JSON Schema as generated by Pydantic, with full support
+    for complex schemas like anyOf, nested objects, validation constraints, etc.
+    Including $defs for complex type references.
+    """
+
+    properties: dict[str, dict[str, Any]] = Field(default_factory=dict)
+    """Dictionary mapping parameter names to their JSON Schema definitions."""
+
+    required: list[str] = Field(default_factory=list)
+    """List of required parameter names."""
+
+    additionalProperties: bool = False
+    """Whether additional properties beyond those defined are allowed."""
+
+    defs: dict[str, dict[str, Any]] | None = Field(default=None, alias="$defs")
+    """JSON Schema definitions for complex types referenced via $ref."""
+
+
+@dataclass
+class ToolSchema(Generic[ToolFnT]):
+    """Underlying schema defining a tool that can be used by LLMs.
+
+    A ToolSchema represents a function that can be called by an LLM during a call.
+    It includes metadata like name, description, and parameter schema.
+
+    This class is not instantiated directly but created by the `@tool()` decorator.
+    """
+
+    fn: ToolFnT
+    """The function that implements the tool's functionality."""
+
+    name: str
+    """The name of the tool, used by the LLM to identify which tool to call."""
+
+    description: str
+    """Description of what the tool does, extracted from the function's docstring."""
+
+    parameters: ToolParameterSchema
+    """JSON Schema describing the parameters accepted by the tool.
+    
+    The serialized parameters table is cached for efficient hash lookups (e.g. when
+    caching provider-encoded tool representations in a LRU cache). Therefore,
+    it should **not be modified** after the ToolSchema is created.
+    """
+
+    strict: bool
+    """Whether the tool should use strict mode when supported by the model."""
+
+    def __hash__(self) -> int:
+        if not hasattr(self, "_hash"):
+            self._hash = hash(
+                (
+                    self.name,
+                    self.description,
+                    self.strict,
+                    self.parameters.model_dump_json(),
+                )
+            )
+        return self._hash
+
+    def __init__(
+        self,
+        fn: ToolFnT,
+        *,
+        strict: bool = False,
+        is_context_tool: bool = False,
+    ) -> None:
+        """Create a `ToolSchema` by inspecting a function and its docstring.
+
+        Uses Pydantic's create_model to dynamically build a model from the function
+        signature, then extracts the JSON schema from it. This leverages Pydantic's
+        robust type conversion and union handling.
+
+        Args:
+            fn: The function to extract schema from
+            strict: Whether the tool should use strict mode when supported
+            is_context_tool: Whether this is a context tool (skips the context parameter)
+
+        Returns:
+            a `ToolSchema` representing the function
+
+        Raises:
+            ValueError: If the tool has a reserved name.
+        """
+        name = fn.__name__
+        if name.startswith(FORMAT_TOOL_NAME):
+            raise ValueError(
+                f"Cannot use reserved name {FORMAT_TOOL_NAME} as tool name."
+            )
+        description = inspect.cleandoc(fn.__doc__) if fn.__doc__ else name
+
+        param_descriptions = _parse_docstring_params(fn.__doc__)
+
+        field_definitions = {}
+        hints = get_type_hints(fn, include_extras=True)
+
+        context_param_skipped = False
+        for param in inspect.signature(fn).parameters.values():
+            # Skip args that are provided by Python or Mirascope, not LLM-generated.
+            # TODO: Handling of FromCallArgs
+            if param.name in ("self", "cls"):
+                continue
+
+            # For context tools, skip the first non-self/cls parameter (the context parameter)
+            if is_context_tool and not context_param_skipped:
+                context_param_skipped = True
+                continue
+
+            param_type = hints.get(param.name, Any)
+            default = ... if param.default is inspect.Parameter.empty else param.default
+            field_info = None
+
+            if get_origin(param_type) is Annotated:
+                args = get_args(param_type)
+                param_type = args[0]
+                for annotation in args[1:]:
+                    if isinstance(annotation, FieldInfo):
+                        field_info = annotation
+                        break
+
+            if field_info is not None:
+                field_value = Field(
+                    default=default,
+                    description=field_info.description,
+                )
+            else:
+                docstring_description = None
+                for arg in param_descriptions.args:
+                    if arg.name == param.name:
+                        docstring_description = arg.description
+                        break
+                if docstring_description:
+                    field_value = Field(
+                        default=default, description=docstring_description
+                    )
+                else:
+                    field_value = default
+
+            field_definitions[param.name] = (param_type, field_value)
+
+        TempModel = create_model("TempModel", **field_definitions)
+
+        schema = TempModel.model_json_schema()
+
+        parameters = ToolParameterSchema(
+            properties=schema.get("properties", {}),
+            required=schema.get("required", []),
+            additionalProperties=False,
+        )
+        if "$defs" in schema:
+            parameters.defs = schema["$defs"]
+
+        self.fn = fn
+        self.name = name
+        self.description = description
+        self.parameters = parameters
+        self.strict = strict
+
+    def can_execute(self, tool_call: ToolCall) -> bool:
+        """Check if a `ToolCall` can be executed by tools with this `ToolSchema`.
+
+        This method is a convenient way to determine if a `ToolCall` is likely intended
+        to be executed by a tool with this `ToolSchema`. It does so by checking
+        whether the name on the call matches the name on the schema. No other validation
+        is performed.
+        """
+        return tool_call.name == self.name

mirascope/llm/tools/toolkit.py

@@ -0,0 +1,152 @@
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar
+
+from ..content import ToolCall, ToolOutput
+from ..context import Context, DepsT
+from ..exceptions import ToolNotFoundError
+from ..types import Jsonable
+from .tool_schema import ToolSchemaT
+from .tools import AsyncContextTool, AsyncTool, ContextTool, Tool
+
+ToolkitT = TypeVar(
+    "ToolkitT",
+    bound="Toolkit | AsyncToolkit | ContextToolkit[Any] | AsyncContextToolkit[Any]",
+    covariant=True,
+)
+
+
+@dataclass
+class BaseToolkit(Generic[ToolSchemaT]):
+    """Base class for tool collections.
+
+    Provides common functionality for managing collections of tools,
+    including name validation and tool lookup.
+    """
+
+    tools: Sequence[ToolSchemaT]
+    """The tools included in the toolkit."""
+
+    tools_dict: dict[str, ToolSchemaT]
+    """A mapping from tool names to tools in the toolkit."""
+
+    def __init__(self, tools: Sequence[ToolSchemaT] | None) -> None:
+        """Initialize the toolkit with a collection of tools.
+
+        Args:
+            tools: Sequence of tools to include in the toolkit.
+
+        Raises:
+            ValueError: If multiple tools have the same name.
+        """
+        self.tools = tools or []
+        self.tools_dict = {}
+        for tool in self.tools:
+            if tool.name in self.tools_dict:
+                raise ValueError(f"Multiple tools with name: {tool.name}")
+            self.tools_dict[tool.name] = tool
+
+    def get(self, tool_call: ToolCall) -> ToolSchemaT:
+        """Get a tool that can execute a specific tool call.
+
+        Args:
+            tool_call: The tool call containing the tool name to lookup.
+
+        Returns:
+            The tool whose name matches the tool call.
+
+        Raises:
+            ToolNotFoundError: If no tool with the given name exists.
+        """
+        tool = self.tools_dict.get(tool_call.name, None)
+        if not tool:
+            raise ToolNotFoundError(f"Tool not found in toolkit: {tool_call.name}")
+        return tool
+
+
+class Toolkit(BaseToolkit[Tool]):
+    """A collection of `Tool`s, with helpers for getting and executing specific tools."""
+
+    def execute(self, tool_call: ToolCall) -> ToolOutput[Jsonable]:
+        """Execute a `Tool` using the provided tool call.
+
+        Args:
+            tool_call: The tool call to execute.
+
+        Returns:
+            The output from executing the `Tool`.
+
+        Raises:
+            ToolNotFoundError: If the requested tool is not found.
+        """
+        tool = self.get(tool_call)
+        return tool.execute(tool_call)
+
+
+class AsyncToolkit(BaseToolkit[AsyncTool]):
+    """A collection of `AsyncTool`s, with helpers for getting and executing specific tools."""
+
+    async def execute(self, tool_call: ToolCall) -> ToolOutput[Jsonable]:
+        """Execute an `AsyncTool` using the provided tool call.
+
+        Args:
+            tool_call: The tool call to execute.
+
+        Returns:
+            The output from executing the `AsyncTool`.
+
+        Raises:
+            ToolNotFoundError: If the requested tool is not found.
+        """
+        tool = self.get(tool_call)
+        return await tool.execute(tool_call)
+
+
+class ContextToolkit(BaseToolkit[Tool | ContextTool[DepsT]], Generic[DepsT]):
+    """A collection of `ContextTool`s, with helpers for getting and executing specific tools."""
+
+    def execute(self, ctx: Context[DepsT], tool_call: ToolCall) -> ToolOutput[Jsonable]:
+        """Execute a `ContextTool` using the provided tool call.
+
+        Args:
+            ctx: The context containing dependencies that match the tool.
+            tool_call: The tool call to execute.
+
+        Returns:
+            The output from executing the `ContextTool`.
+
+        Raises:
+            ToolNotFoundError: If the requested tool is not found.
+        """
+        tool = self.get(tool_call)
+        if isinstance(tool, ContextTool):
+            return tool.execute(ctx, tool_call)
+        else:
+            return tool.execute(tool_call)
+
+
+class AsyncContextToolkit(
+    BaseToolkit[AsyncTool | AsyncContextTool[DepsT]], Generic[DepsT]
+):
+    """A collection of `AsyncContextTool`s, with helpers for getting and executing specific tools."""
+
+    async def execute(
+        self, ctx: Context[DepsT], tool_call: ToolCall
+    ) -> ToolOutput[Jsonable]:
+        """Execute an `AsyncContextTool` using the provided tool call.
+
+        Args:
+            ctx: The context containing dependencies that match the tool.
+            tool_call: The tool call to execute.
+
+        Returns:
+            The output from executing the `AsyncContextTool`.
+
+        Raises:
+            ToolNotFoundError: If the requested tool is not found.
+        """
+        tool = self.get(tool_call)
+        if isinstance(tool, AsyncContextTool):
+            return await tool.execute(ctx, tool_call)
+        else:
+            return await tool.execute(tool_call)

mirascope/llm/tools/tools.py

@@ -0,0 +1,169 @@
+"""The `Tool` class for defining tools that LLMs can request be called."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Awaitable
+from typing import Any, Generic, TypeVar, cast
+
+from ..content import ToolCall, ToolOutput
+from ..context import Context, DepsT
+from ..types import AnyP, JsonableCovariantT
+from .protocols import (
+    AsyncContextToolFn,
+    AsyncToolFn,
+    ContextToolFn,
+    ToolFn,
+    _AsyncJsonKwargsCallable,
+    _AsyncKwargsCallable,
+    _ContextKwargsCallable,
+    _KwargsCallable,
+)
+from .tool_schema import ToolSchema
+
+ToolT = TypeVar(
+    "ToolT",
+    bound="Tool | AsyncTool | ContextTool[Any] | AsyncContextTool[Any]",
+    covariant=True,
+)
+
+
+class Tool(
+    ToolSchema[ToolFn[AnyP, JsonableCovariantT]], Generic[AnyP, JsonableCovariantT]
+):
+    """A tool that can be used by LLMs.
+
+    A `Tool` represents a function that can be called by an LLM during a call.
+    It includes metadata like name, description, and parameter schema.
+
+    This class is not instantiated directly but created by the `@tool()` decorator.
+    """
+
+    def __init__(
+        self, fn: ToolFn[AnyP, JsonableCovariantT], *, strict: bool = False
+    ) -> None:
+        super().__init__(fn, strict=strict, is_context_tool=False)
+
+    def __call__(self, *args: AnyP.args, **kwargs: AnyP.kwargs) -> JsonableCovariantT:
+        """Call the underlying function directly."""
+        return self.fn(*args, **kwargs)
+
+    def execute(self, tool_call: ToolCall) -> ToolOutput[JsonableCovariantT]:
+        """Execute the tool using an LLM-provided `ToolCall`."""
+        kwargs_from_json = json.loads(tool_call.args)
+        kwargs_callable = cast(_KwargsCallable[JsonableCovariantT], self.fn)
+        result = kwargs_callable(**kwargs_from_json)
+        return ToolOutput(id=tool_call.id, value=result, name=self.name)
+
+
+class AsyncTool(
+    ToolSchema[AsyncToolFn[AnyP, JsonableCovariantT]],
+    Generic[AnyP, JsonableCovariantT],
+):
+    """An async tool that can be used by LLMs.
+
+    An `AsyncTool` represents an async function that can be called by an LLM during a call.
+    It includes metadata like name, description, and parameter schema.
+
+    This class is not instantiated directly but created by the `@tool()` decorator.
+    """
+
+    def __init__(
+        self, fn: AsyncToolFn[AnyP, JsonableCovariantT], *, strict: bool = False
+    ) -> None:
+        super().__init__(fn, strict=strict, is_context_tool=False)
+
+    def __call__(
+        self, *args: AnyP.args, **kwargs: AnyP.kwargs
+    ) -> Awaitable[JsonableCovariantT]:
+        """Call the underlying async function directly."""
+        return self.fn(*args, **kwargs)
+
+    async def execute(self, tool_call: ToolCall) -> ToolOutput[JsonableCovariantT]:
+        """Execute the async tool using an LLM-provided `ToolCall`."""
+        kwargs_from_json = json.loads(tool_call.args)
+        kwargs_callable = cast(_AsyncKwargsCallable[JsonableCovariantT], self.fn)
+        result = await kwargs_callable(**kwargs_from_json)
+        return ToolOutput(id=tool_call.id, value=result, name=self.name)
+
+
+class ContextTool(
+    ToolSchema[ContextToolFn[DepsT, AnyP, JsonableCovariantT]],
+    Generic[DepsT, AnyP, JsonableCovariantT],
+):
+    """Protocol defining a tool that can be used by LLMs.
+
+    A `ContextTool` represents a function that can be called by an LLM during a call.
+    It includes metadata like name, description, and parameter schema.
+
+    This class is not instantiated directly but created by the `@tool()` decorator.
+    """
+
+    def __init__(
+        self,
+        fn: ContextToolFn[DepsT, AnyP, JsonableCovariantT],
+        *,
+        strict: bool = False,
+    ) -> None:
+        super().__init__(fn, strict=strict, is_context_tool=True)
+
+    def __call__(
+        self,
+        ctx: Context[DepsT],
+        *args: AnyP.args,
+        **kwargs: AnyP.kwargs,
+    ) -> JsonableCovariantT:
+        """Call the underlying function directly with context."""
+        return self.fn(ctx, *args, **kwargs)
+
+    def execute(
+        self, ctx: Context[DepsT], tool_call: ToolCall
+    ) -> ToolOutput[JsonableCovariantT]:
+        """Execute the context tool using an LLM-provided `ToolCall`."""
+        kwargs_from_json = json.loads(tool_call.args)
+        kwargs_callable = cast(
+            _ContextKwargsCallable[DepsT, JsonableCovariantT], self.fn
+        )
+        result = kwargs_callable(ctx, **kwargs_from_json)
+        return ToolOutput(id=tool_call.id, value=result, name=self.name)
+
+
+class AsyncContextTool(
+    ToolSchema[AsyncContextToolFn[DepsT, AnyP, JsonableCovariantT]],
+    Generic[DepsT, AnyP, JsonableCovariantT],
+):
+    """Protocol defining an async tool that can be used by LLMs with context.
+
+    An `AsyncContextTool` represents an async function that can be called by an LLM during a call.
+    It includes metadata like name, description, and parameter schema.
+
+    This class is not instantiated directly but created by the `@tool()` decorator.
+    """
+
+    def __init__(
+        self,
+        fn: AsyncContextToolFn[DepsT, AnyP, JsonableCovariantT],
+        *,
+        strict: bool = False,
+    ) -> None:
+        super().__init__(fn, strict=strict, is_context_tool=True)
+
+    def __call__(
+        self,
+        ctx: Context[DepsT],
+        *args: AnyP.args,
+        **kwargs: AnyP.kwargs,
+    ) -> Awaitable[JsonableCovariantT]:
+        """Call the underlying async function directly with context."""
+        return self.fn(ctx, *args, **kwargs)
+
+    async def execute(
+        self, ctx: Context[DepsT], tool_call: ToolCall
+    ) -> ToolOutput[JsonableCovariantT]:
+        """Execute the async context tool using an LLM-provided `ToolCall`."""
+        kwargs_from_json = json.loads(tool_call.args)
+        kwargs_callable = cast(
+            _AsyncJsonKwargsCallable[DepsT, JsonableCovariantT], self.fn
+        )
+        result = await kwargs_callable(ctx, **kwargs_from_json)
+        return ToolOutput(id=tool_call.id, value=result, name=self.name)

mirascope/llm/types/__init__.py

@@ -0,0 +1,22 @@
+"""Types for the LLM module."""
+
+from .dataclass import Dataclass
+from .jsonable import (
+    Jsonable,
+    JsonableCovariantT,
+    JsonableT,
+)
+from .type_vars import AnyP, CovariantT, P
+
+NoneType = type(None)
+
+__all__ = [
+    "AnyP",
+    "CovariantT",
+    "Dataclass",
+    "Jsonable",
+    "JsonableCovariantT",
+    "JsonableT",
+    "NoneType",
+    "P",
+]

mirascope/llm/types/dataclass.py

@@ -0,0 +1,9 @@
+"""A protocol for defining a dataclass type."""
+
+from dataclasses import Field
+from typing import Any, ClassVar, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class Dataclass(Protocol):
+    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]

mirascope/llm/types/jsonable.py

@@ -0,0 +1,44 @@
+"""Typing for JSON serializable objects."""
+
+from collections.abc import Mapping, Sequence
+from typing import Protocol, TypeAlias
+from typing_extensions import TypeVar
+
+
+class JsonableObject(Protocol):
+    """Protocol for JSON-serializable objects.
+
+    This protocol defines the interface for objects that can be serialized to
+    JSON. It is used to annotate the `JsonableType` type alias.
+    """
+
+    def json(self) -> str:
+        """Serialize the object as a JSON string."""
+        ...
+
+
+Jsonable: TypeAlias = (
+    None
+    | str
+    | int
+    | float
+    | bool
+    | Sequence["Jsonable"]
+    | Mapping[str, "Jsonable"]
+    | JsonableObject
+)
+"""Simple type alias for JSON-serializable types."""
+
+
+JsonableT = TypeVar("JsonableT", bound=Jsonable)
+"""Type variable for tool output types.
+
+This TypeVar represents the return type of tool functions, which must be
+serializable to JSON (bound to Jsonable) for LLM consumption.
+"""
+
+
+JsonableCovariantT = TypeVar(
+    "JsonableCovariantT", covariant=True, bound=Jsonable, default=Jsonable
+)
+"""Type variable for covariant types that are Jsonable."""

mirascope/llm/types/type_vars.py

@@ -0,0 +1,19 @@
+"""Common TypeVar definitions for the LLM module."""
+
+from typing import TypeVar
+from typing_extensions import ParamSpec
+
+P = ParamSpec("P")
+"""Parameter specification for function signatures.
+
+This ParamSpec is used to preserve function parameter types and signatures
+when wrapping functions with decorators or creating generic callable types.
+It captures both positional and keyword arguments (*args, **kwargs) while
+maintaining their original types.
+"""
+
+AnyP = ParamSpec("AnyP", default=...)
+"""A parameter specification that defaults to ..."""
+
+CovariantT = TypeVar("CovariantT", covariant=True)
+"""Type variable for covariant types."""