mirascope 2.0.0a1__py3-none-any.whl → 2.0.0a3__py3-none-any.whl

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/decorator.py

@@ -20,14 +20,14 @@ class ToolDecorator:
     @overload
     def __call__(  # pyright:ignore[reportOverlappingOverload]
         self, fn: ContextToolFn[DepsT, P, JsonableCovariantT]
-    ) -> ContextTool[DepsT, P, JsonableCovariantT]:
+    ) -> ContextTool[DepsT, JsonableCovariantT, P]:
         """Call the decorator with a context function."""
         ...
 
     @overload
     def __call__(  # pyright:ignore[reportOverlappingOverload]
         self, fn: AsyncContextToolFn[DepsT, P, JsonableCovariantT]
-    ) -> AsyncContextTool[DepsT, P, JsonableCovariantT]:
+    ) -> AsyncContextTool[DepsT, JsonableCovariantT, P]:
         """Call the decorator with an async context function."""
         ...
 
@@ -52,8 +52,8 @@ class ToolDecorator:
         | ToolFn[P, JsonableCovariantT]
         | AsyncToolFn[P, JsonableCovariantT],
     ) -> (
-        ContextTool[DepsT, P, JsonableCovariantT]
-        | AsyncContextTool[DepsT, P, JsonableCovariantT]
+        ContextTool[DepsT, JsonableCovariantT, P]
+        | AsyncContextTool[DepsT, JsonableCovariantT, P]
         | Tool[P, JsonableCovariantT]
         | AsyncTool[P, JsonableCovariantT]
     ):
@@ -62,11 +62,11 @@ class ToolDecorator:
         is_async = _tool_utils.is_async_tool_fn(fn)
 
         if is_context and is_async:
-            return AsyncContextTool[DepsT, P, JsonableCovariantT](
+            return AsyncContextTool[DepsT, JsonableCovariantT, P](
                 fn, strict=self.strict
             )
         elif is_context:
-            return ContextTool[DepsT, P, JsonableCovariantT](fn, strict=self.strict)
+            return ContextTool[DepsT, JsonableCovariantT, P](fn, strict=self.strict)
         elif is_async:
             return AsyncTool[P, JsonableCovariantT](fn, strict=self.strict)
         else:
@@ -76,7 +76,7 @@ class ToolDecorator:
 @overload
 def tool(  # pyright:ignore[reportOverlappingOverload]
     __fn: AsyncContextToolFn[DepsT, P, JsonableCovariantT],
-) -> AsyncContextTool[DepsT, P, JsonableCovariantT]:
+) -> AsyncContextTool[DepsT, JsonableCovariantT, P]:
     """Overload for async context tool functions."""
     ...
 
@@ -84,7 +84,7 @@ def tool(  # pyright:ignore[reportOverlappingOverload]
 @overload
 def tool(  # pyright:ignore[reportOverlappingOverload]
     __fn: ContextToolFn[DepsT, P, JsonableCovariantT],
-) -> ContextTool[DepsT, P, JsonableCovariantT]:
+) -> ContextTool[DepsT, JsonableCovariantT, P]:
     """Overload for context tool functions."""
     ...
 
@@ -116,8 +116,8 @@ def tool(
     *,
     strict: bool = False,
 ) -> (
-    ContextTool[DepsT, P, JsonableCovariantT]
-    | AsyncContextTool[DepsT, P, JsonableCovariantT]
+    ContextTool[DepsT, JsonableCovariantT, P]
+    | AsyncContextTool[DepsT, JsonableCovariantT, P]
     | Tool[P, JsonableCovariantT]
     | AsyncTool[P, JsonableCovariantT]
     | ToolDecorator

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,14 +83,14 @@ 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)
 
 
 class ContextTool(
     ToolSchema[ContextToolFn[DepsT, AnyP, JsonableCovariantT]],
-    Generic[DepsT, AnyP, JsonableCovariantT],
+    Generic[DepsT, JsonableCovariantT, AnyP],
 ):
     """Protocol defining a tool that can be used by LLMs.
 
@@ -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)
@@ -130,7 +131,7 @@ class ContextTool(
 
 class AsyncContextTool(
     ToolSchema[AsyncContextToolFn[DepsT, AnyP, JsonableCovariantT]],
-    Generic[DepsT, AnyP, JsonableCovariantT],
+    Generic[DepsT, JsonableCovariantT, AnyP],
 ):
     """Protocol defining an async tool that can be used by LLMs with context.
 
@@ -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.
+"""