mirascope 2.0.0a2__py3-none-any.whl → 2.0.0a4__py3-none-any.whl

mirascope/llm/responses/usage.py

@@ -0,0 +1,95 @@
+"""Provider-agnostic usage statistics for LLM API calls."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal
+
+
+@dataclass(kw_only=True)
+class UsageDeltaChunk:
+    """A chunk containing incremental token usage information from a streaming response.
+
+    This represents a delta/increment in usage statistics as they arrive during streaming.
+    Multiple UsageDeltaChunks are accumulated to produce the final Usage object.
+    """
+
+    type: Literal["usage_delta_chunk"] = "usage_delta_chunk"
+
+    input_tokens: int = 0
+    """Delta in input tokens."""
+
+    output_tokens: int = 0
+    """Delta in output tokens."""
+
+    cache_read_tokens: int = 0
+    """Delta in cache read tokens."""
+
+    cache_write_tokens: int = 0
+    """Delta in cache write tokens."""
+
+    reasoning_tokens: int = 0
+    """Delta in reasoning/thinking tokens."""
+
+
+@dataclass(kw_only=True)
+class Usage:
+    """Token usage statistics from an LLM API call.
+
+    This abstraction captures common usage metrics across providers while preserving
+    access to the raw provider-specific usage data.
+    """
+
+    input_tokens: int = 0
+    """The number of input tokens used.
+
+    This includes ALL input tokens, including cache read and write tokens.
+
+    Will be 0 if not reported by the provider.
+    """
+
+    output_tokens: int = 0
+    """The number of output tokens used.
+
+    This includes ALL output tokens, including `reasoning_tokens` that may not be
+    in the user's visible output, or other "hidden" tokens.
+
+    Will be 0 if not reported by the provider.
+    """
+
+    cache_read_tokens: int = 0
+    """The number of tokens read from cache (prompt caching).
+
+    These are input tokens that were read from cache. Cache read tokens are generally
+    much less expensive than regular input tokens.
+
+    Will be 0 if not reported by the provider or if caching was not used.
+    """
+
+    cache_write_tokens: int = 0
+    """The number of tokens written to cache (cache creation).
+
+    These are input tokens that were written to cache, for future reuse and retrieval.
+    Cache write tokens are generally more expensive than uncached input tokens,
+    but may lead to cost savings down the line when they are re-read as cache_read_tokens.
+
+    Will be 0 if not reported by the provider or if caching was not used.
+    """
+
+    reasoning_tokens: int = 0
+    """The number of tokens used for reasoning/thinking.
+
+    Reasoning tokens are a subset of output_tokens that were generated as part of the model's
+    interior reasoning process. They are billed as output tokens, though they are generally
+    not shown to the user.
+
+    Will be 0 if not reported by the provider or if the model does not support reasoning.
+    """
+
+    raw: Any = None
+    """The raw usage object from the provider."""
+
+    @property
+    def total_tokens(self) -> int:
+        """The total number of tokens used (input + output)."""
+        return self.input_tokens + self.output_tokens

mirascope/llm/tools/__init__.py

@@ -1,9 +1,11 @@
 """The Tools module for LLMs."""
 
-from . import protocols
 from .decorator import ToolDecorator, tool
+from .protocols import AsyncContextToolFn, AsyncToolFn, ContextToolFn, ToolFn
 from .tool_schema import (
     FORMAT_TOOL_NAME,
+    AnyToolFn,
+    AnyToolSchema,
     ToolParameterSchema,
     ToolSchema,
     ToolSchemaT,
@@ -20,21 +22,26 @@ from .tools import AsyncContextTool, AsyncTool, ContextTool, Tool, ToolT
 
 __all__ = [
     "FORMAT_TOOL_NAME",
+    "AnyToolFn",
+    "AnyToolSchema",
     "AsyncContextTool",
+    "AsyncContextToolFn",
     "AsyncContextToolkit",
     "AsyncTool",
+    "AsyncToolFn",
     "AsyncToolkit",
     "BaseToolkit",
     "ContextTool",
+    "ContextToolFn",
     "ContextToolkit",
     "Tool",
     "ToolDecorator",
+    "ToolFn",
     "ToolParameterSchema",
     "ToolSchema",
     "ToolSchemaT",
     "ToolT",
     "Toolkit",
     "ToolkitT",
-    "protocols",
     "tool",
 ]

mirascope/llm/tools/_utils.py

@@ -9,7 +9,10 @@ from .protocols import AsyncContextToolFn, AsyncToolFn, ContextToolFn, ToolFn
 
 
 def is_context_tool_fn(
-    fn: ToolFn | AsyncToolFn | ContextToolFn | AsyncContextToolFn,
+    fn: ContextToolFn[DepsT, P, JsonableCovariantT]
+    | AsyncContextToolFn[DepsT, P, JsonableCovariantT]
+    | ToolFn[P, JsonableCovariantT]
+    | AsyncToolFn[P, JsonableCovariantT],
 ) -> TypeIs[
     ContextToolFn[DepsT, P, JsonableCovariantT]
     | AsyncContextToolFn[DepsT, P, JsonableCovariantT]
@@ -19,7 +22,13 @@ def is_context_tool_fn(
 
 
 def is_async_tool_fn(
-    fn: ToolFn | AsyncToolFn | ContextToolFn | AsyncContextToolFn,
-) -> TypeIs[AsyncToolFn | AsyncContextToolFn]:
+    fn: ContextToolFn[DepsT, P, JsonableCovariantT]
+    | AsyncContextToolFn[DepsT, P, JsonableCovariantT]
+    | ToolFn[P, JsonableCovariantT]
+    | AsyncToolFn[P, JsonableCovariantT],
+) -> TypeIs[
+    AsyncContextToolFn[DepsT, P, JsonableCovariantT]
+    | AsyncToolFn[P, JsonableCovariantT]
+]:
     """Type guard to check if a function is an async tool function."""
     return inspect.iscoroutinefunction(fn)

mirascope/llm/tools/protocols.py

@@ -52,7 +52,7 @@ class AsyncContextToolFn(Protocol[DepsT, P, JsonableCovariantT]):
         raise NotImplementedError()
 
 
-class _KwargsCallable(Protocol[JsonableCovariantT]):
+class KwargsCallable(Protocol[JsonableCovariantT]):
     """Protocol for functions that can be called with `Any`-typed kwargs.
 
     Used internally to type-cast tool functions for compatibility with
@@ -62,7 +62,7 @@ class _KwargsCallable(Protocol[JsonableCovariantT]):
     def __call__(self, **kwargs: dict[str, Any]) -> JsonableCovariantT: ...
 
 
-class _AsyncKwargsCallable(Protocol[JsonableCovariantT]):
+class AsyncKwargsCallable(Protocol[JsonableCovariantT]):
     """Protocol for async functions that can be called with `Any`-typed kwargs.
 
     Used internally to type-cast async tool functions for compatibility with
@@ -72,7 +72,7 @@ class _AsyncKwargsCallable(Protocol[JsonableCovariantT]):
     async def __call__(self, **kwargs: dict[str, Any]) -> JsonableCovariantT: ...
 
 
-class _ContextKwargsCallable(Protocol[DepsT, JsonableCovariantT]):
+class ContextKwargsCallable(Protocol[DepsT, JsonableCovariantT]):
     """Protocol for context functions that can be called with `Any`-typed kwargs.
 
     Used internally to type-cast context tool functions for compatibility with
@@ -84,7 +84,7 @@ class _ContextKwargsCallable(Protocol[DepsT, JsonableCovariantT]):
     ) -> JsonableCovariantT: ...
 
 
-class _AsyncJsonKwargsCallable(Protocol[DepsT, JsonableCovariantT]):
+class AsyncJsonKwargsCallable(Protocol[DepsT, JsonableCovariantT]):
     """Protocol for async context functions that can be called with `Any`-typed kwargs.
 
     Used internally to type-cast async context tool functions for compatibility with

mirascope/llm/tools/tool_schema.py

@@ -3,13 +3,16 @@
 from __future__ import annotations
 
 import inspect
-from collections import namedtuple
 from dataclasses import dataclass
 from typing import (
     Annotated,
     Any,
     Generic,
+    NamedTuple,
+    TypeAlias,
+    TypedDict,
     TypeVar,
+    cast,
     get_args,
     get_origin,
     get_type_hints,
@@ -20,16 +23,48 @@ from pydantic import BaseModel, Field, create_model
 from pydantic.fields import FieldInfo
 
 from ..content import ToolCall
+from ..types import Jsonable
 from .protocols import AsyncContextToolFn, AsyncToolFn, ContextToolFn, ToolFn
 
+AnyToolFn: TypeAlias = (
+    ToolFn[..., Jsonable]
+    | AsyncToolFn[..., Jsonable]
+    | ContextToolFn[Any, ..., Jsonable]
+    | AsyncContextToolFn[Any, ..., Jsonable]
+)
+
 ToolFnT = TypeVar(
     "ToolFnT",
-    bound=ToolFn | AsyncToolFn | ContextToolFn | AsyncContextToolFn,
+    bound=AnyToolFn,
     covariant=True,
 )
-ToolSchemaT = TypeVar("ToolSchemaT", bound="ToolSchema")
 
-DocstringArg = namedtuple("DocstringArg", ["name", "description"])
+AnyToolSchema: TypeAlias = "ToolSchema[AnyToolFn]"
+ToolSchemaT = TypeVar("ToolSchemaT", bound=AnyToolSchema, covariant=True)
+
+
+ModelJsonSchema = TypedDict(
+    "ModelJsonSchema",
+    {
+        "properties": dict[str, dict[str, Any]],
+        "required": list[str],
+        "$defs": dict[str, dict[str, Any]],
+    },
+    total=False,
+)
+"""Type for Pydantic's model_json_schema() output.
+
+This TypedDict defines the structure of JSON schemas returned by Pydantic models,
+allowing us to avoid type casts when extracting schema components.
+"""
+
+
+class DocstringArg(NamedTuple):
+    """A parameter from a docstring."""
+
+    name: str
+    description: str
+
 
 FORMAT_TOOL_NAME = "__mirascope_formatted_output_tool__"
 """Reserved name of the formatted output tool. 
@@ -63,7 +98,7 @@ def _parse_docstring_params(docstring: str | None) -> ParsedDocstring:
         return ParsedDocstring(args=[])
 
     parsed = parse(docstring)
-    args = []
+    args: list[DocstringArg] = []
 
     for param in parsed.params:
         if param.description:
@@ -170,7 +205,7 @@ class ToolSchema(Generic[ToolFnT]):
 
         param_descriptions = _parse_docstring_params(fn.__doc__)
 
-        field_definitions = {}
+        field_definitions: dict[str, tuple[Any, Any]] = {}
         hints = get_type_hints(fn, include_extras=True)
 
         context_param_skipped = False
@@ -203,7 +238,7 @@ class ToolSchema(Generic[ToolFnT]):
                     description=field_info.description,
                 )
             else:
-                docstring_description = None
+                docstring_description: str | None = None
                 for arg in param_descriptions.args:
                     if arg.name == param.name:
                         docstring_description = arg.description
@@ -217,9 +252,9 @@ class ToolSchema(Generic[ToolFnT]):
 
             field_definitions[param.name] = (param_type, field_value)
 
-        TempModel = create_model("TempModel", **field_definitions)
+        TempModel = create_model("TempModel", **cast(dict[str, Any], field_definitions))
 
-        schema = TempModel.model_json_schema()
+        schema = cast(ModelJsonSchema, TempModel.model_json_schema())
 
         parameters = ToolParameterSchema(
             properties=schema.get("properties", {}),

mirascope/llm/tools/tools.py

@@ -4,20 +4,21 @@ from __future__ import annotations
 
 import json
 from collections.abc import Awaitable
-from typing import Any, Generic, TypeVar, cast
+from typing import Any, Generic, cast
+from typing_extensions import TypeVar
 
 from ..content import ToolCall, ToolOutput
 from ..context import Context, DepsT
 from ..types import AnyP, JsonableCovariantT
 from .protocols import (
     AsyncContextToolFn,
+    AsyncJsonKwargsCallable,
+    AsyncKwargsCallable,
     AsyncToolFn,
+    ContextKwargsCallable,
     ContextToolFn,
+    KwargsCallable,
     ToolFn,
-    _AsyncJsonKwargsCallable,
-    _AsyncKwargsCallable,
-    _ContextKwargsCallable,
-    _KwargsCallable,
 )
 from .tool_schema import ToolSchema
 
@@ -51,7 +52,7 @@ class Tool(
     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)
+        kwargs_callable = cast(KwargsCallable[JsonableCovariantT], self.fn)
         result = kwargs_callable(**kwargs_from_json)
         return ToolOutput(id=tool_call.id, value=result, name=self.name)
 
@@ -82,7 +83,7 @@ class AsyncTool(
     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)
+        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)
 
@@ -122,7 +123,7 @@ class ContextTool(
         """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
+            ContextKwargsCallable[DepsT, JsonableCovariantT], self.fn
         )
         result = kwargs_callable(ctx, **kwargs_from_json)
         return ToolOutput(id=tool_call.id, value=result, name=self.name)
@@ -163,7 +164,7 @@ class AsyncContextTool(
         """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
+            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/ops/__init__.py

@@ -0,0 +1,156 @@
+"""Operational helpers (tracing, sessions, future ops.* utilities)."""
+
+from __future__ import annotations
+
+try:
+    from ._internal.configuration import configure, tracer_context
+    from ._internal.context import propagated_context
+    from ._internal.instrumentation.llm import (
+        instrument_llm,
+        uninstrument_llm,
+    )
+    from ._internal.propagation import (
+        ContextPropagator,
+        PropagatorFormat,
+        extract_context,
+        get_propagator,
+        inject_context,
+        reset_propagator,
+    )
+    from ._internal.session import (
+        SESSION_HEADER_NAME,
+        SessionContext,
+        current_session,
+        extract_session_id,
+        session,
+    )
+    from ._internal.spans import Span, span
+    from ._internal.traced_calls import (
+        TracedAsyncCall,
+        TracedAsyncContextCall,
+        TracedCall,
+        TracedContextCall,
+    )
+    from ._internal.traced_functions import (
+        AsyncTrace,
+        AsyncTracedFunction,
+        Trace,
+        TracedFunction,
+    )
+    from ._internal.tracing import (
+        TraceDecorator,
+        trace,
+    )
+    from ._internal.versioned_calls import (
+        VersionedAsyncCall,
+        VersionedAsyncContextCall,
+        VersionedCall,
+        VersionedContextCall,
+    )
+    from ._internal.versioned_functions import (
+        AsyncVersionedFunction,
+        VersionedFunction,
+        VersionInfo,
+    )
+    from ._internal.versioning import (
+        VersionDecorator,
+        version,
+    )
+    from .exceptions import ClosureComputationError
+
+except ImportError:  # pragma: no cover
+    # TODO: refactor alongside other import error handling improvements
+    from collections.abc import Callable
+
+    def _create_otel_import_error_stub(name: str) -> Callable[..., None]:
+        """Create a stub that raises ImportError with helpful message."""
+
+        def _raise_not_installed() -> None:
+            raise ImportError(
+                f"The 'opentelemetry' packages are required to use {name}. "
+                "Install them with: `uv add 'mirascope[otel]'`."
+            )
+
+        return _raise_not_installed
+
+    propagated_context = _create_otel_import_error_stub("propagated_context")
+    instrument_llm = _create_otel_import_error_stub("instrument_llm")
+    uninstrument_llm = _create_otel_import_error_stub("uninstrument_llm")
+    ContextPropagator = _create_otel_import_error_stub("ContextPropagator")
+    PropagatorFormat = str
+    extract_context = _create_otel_import_error_stub("extract_context")
+    get_propagator = _create_otel_import_error_stub("get_propagator")
+    inject_context = _create_otel_import_error_stub("inject_context")
+    reset_propagator = _create_otel_import_error_stub("reset_propagator")
+    SESSION_HEADER_NAME = "X-Mirascope-Session-ID"
+    SessionContext = _create_otel_import_error_stub("SessionContext")
+    current_session = _create_otel_import_error_stub("current_session")
+    extract_session_id = _create_otel_import_error_stub("extract_session_id")
+    session = _create_otel_import_error_stub("session")
+    Span = _create_otel_import_error_stub("Span")
+    span = _create_otel_import_error_stub("span")
+    AsyncTrace = _create_otel_import_error_stub("AsyncTrace")
+    AsyncTracedFunction = _create_otel_import_error_stub("AsyncTracedFunction")
+    Trace = _create_otel_import_error_stub("Trace")
+    TraceDecorator = _create_otel_import_error_stub("TraceDecorator")
+    TracedAsyncCall = _create_otel_import_error_stub("TracedAsyncCall")
+    TracedAsyncContextCall = _create_otel_import_error_stub("TracedAsyncContextCall")
+    TracedCall = _create_otel_import_error_stub("TracedCall")
+    TracedContextCall = _create_otel_import_error_stub("TracedContextCall")
+    TracedFunction = _create_otel_import_error_stub("TracedFunction")
+    trace = _create_otel_import_error_stub("trace")
+    configure = _create_otel_import_error_stub("configure")
+    tracer_context = _create_otel_import_error_stub("tracer_context")
+    AsyncVersionedFunction = _create_otel_import_error_stub("AsyncVersionedFunction")
+    VersionDecorator = _create_otel_import_error_stub("VersionDecorator")
+    VersionedAsyncCall = _create_otel_import_error_stub("VersionedAsyncCall")
+    VersionedAsyncContextCall = _create_otel_import_error_stub(
+        "VersionedAsyncContextCall"
+    )
+    VersionedCall = _create_otel_import_error_stub("VersionedCall")
+    VersionedContextCall = _create_otel_import_error_stub("VersionedContextCall")
+    VersionedFunction = _create_otel_import_error_stub("VersionedFunction")
+    VersionInfo = _create_otel_import_error_stub("VersionInfo")
+    version = _create_otel_import_error_stub("version")
+
+
+__all__ = [
+    "SESSION_HEADER_NAME",
+    "AsyncTrace",
+    "AsyncTracedFunction",
+    "AsyncVersionedFunction",
+    "ClosureComputationError",
+    "ContextPropagator",
+    "PropagatorFormat",
+    "SessionContext",
+    "Span",
+    "Trace",
+    "TraceDecorator",
+    "TracedAsyncCall",
+    "TracedAsyncContextCall",
+    "TracedCall",
+    "TracedContextCall",
+    "TracedFunction",
+    "VersionDecorator",
+    "VersionInfo",
+    "VersionedAsyncCall",
+    "VersionedAsyncContextCall",
+    "VersionedCall",
+    "VersionedContextCall",
+    "VersionedFunction",
+    "configure",
+    "current_session",
+    "extract_context",
+    "extract_session_id",
+    "get_propagator",
+    "inject_context",
+    "instrument_llm",
+    "propagated_context",
+    "reset_propagator",
+    "session",
+    "span",
+    "trace",
+    "tracer_context",
+    "uninstrument_llm",
+    "version",
+]

mirascope/ops/_internal/__init__.py

@@ -0,0 +1,5 @@
+"""Private internal module for Mirascope.
+
+All modules within this directory are considered private implementation details
+and should not be imported directly by users.
+"""