mirascope 1.0.5__py3-none-any.whl → 2.1.1__py3-none-any.whl

mirascope/ops/_internal/tracing.py

@@ -0,0 +1,353 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from typing import (
+    TYPE_CHECKING,
+    overload,
+)
+
+from ...llm.calls import AsyncCall, AsyncContextCall, Call, ContextCall
+from ...llm.context import DepsT
+from .protocols import (
+    AsyncFunction,
+    AsyncSpanFunction,
+    SyncFunction,
+    SyncSpanFunction,
+    fn_is_async,
+    fn_wants_span,
+)
+from .traced_calls import (
+    TracedAsyncCall,
+    TracedAsyncContextCall,
+    TracedCall,
+    TracedContextCall,
+    is_call_type,
+    wrap_call,
+)
+from .traced_functions import (
+    AsyncTracedFunction,
+    AsyncTracedSpanFunction,
+    TracedFunction,
+    TracedSpanFunction,
+)
+from .types import P, R
+
+if TYPE_CHECKING:
+    from ...llm.formatting import FormattableT
+
+
+@dataclass(kw_only=True)
+class TraceDecorator:
+    """Decorator implementation for adding tracing capabilities to functions."""
+
+    tags: tuple[str, ...] = ()
+    """Tags to be associated with traced function calls."""
+
+    metadata: dict[str, str] = field(default_factory=dict)
+    """Arbitrary key-value pairs for additional metadata."""
+
+    # IMPORTANT: The order of these overloads matters for type inference.
+    # Call type overloads come first, then span function overloads, then regular functions.
+    @overload
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
+        self,
+        fn: AsyncContextCall[P, DepsT, FormattableT],
+    ) -> TracedAsyncContextCall[P, DepsT, FormattableT]:
+        """Overload for applying decorator to an AsyncContextCall."""
+        ...
+
+    @overload
+    def __call__(
+        self,
+        fn: ContextCall[P, DepsT, FormattableT],
+    ) -> TracedContextCall[P, DepsT, FormattableT]:
+        """Overload for applying decorator to a ContextCall."""
+        ...
+
+    @overload
+    def __call__(
+        self,
+        fn: AsyncCall[P, FormattableT],
+    ) -> TracedAsyncCall[P, FormattableT]:
+        """Overload for applying decorator to an AsyncCall."""
+        ...
+
+    @overload
+    def __call__(
+        self,
+        fn: Call[P, FormattableT],
+    ) -> TracedCall[P, FormattableT]:
+        """Overload for applying decorator to a Call."""
+        ...
+
+    @overload
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
+        self,
+        fn: AsyncSpanFunction[P, R],
+    ) -> AsyncTracedSpanFunction[P, R]:
+        """Overload for applying decorator to an async function with span injection."""
+        ...
+
+    @overload
+    def __call__(
+        self,
+        fn: SyncSpanFunction[P, R],
+    ) -> TracedSpanFunction[P, R]:
+        """Overload for applying decorator to a sync function with span injection."""
+        ...
+
+    @overload
+    def __call__(
+        self,
+        fn: AsyncFunction[P, R],
+    ) -> AsyncTracedFunction[P, R]:
+        """Overload for applying decorator to an async function."""
+        ...
+
+    @overload
+    def __call__(
+        self,
+        fn: SyncFunction[P, R],
+    ) -> TracedFunction[P, R]:
+        """Overload for applying decorator to a sync function."""
+        ...
+
+    def __call__(  # pyright: ignore[reportGeneralTypeIssues]
+        self,
+        fn: (
+            AsyncContextCall[P, DepsT, FormattableT]
+            | ContextCall[P, DepsT, FormattableT]
+            | AsyncCall[P, FormattableT]
+            | Call[P, FormattableT]
+            | AsyncSpanFunction[P, R]
+            | SyncSpanFunction[P, R]
+            | AsyncFunction[P, R]
+            | SyncFunction[P, R]
+        ),
+    ) -> (
+        TracedAsyncContextCall[P, DepsT, FormattableT]
+        | TracedContextCall[P, DepsT, FormattableT]
+        | TracedAsyncCall[P, FormattableT]
+        | TracedCall[P, FormattableT]
+        | AsyncTracedSpanFunction[P, R]
+        | TracedSpanFunction[P, R]
+        | AsyncTracedFunction[P, R]
+        | TracedFunction[P, R]
+    ):
+        """Applies the decorator to the given function or Call object."""
+        if is_call_type(fn):
+            return wrap_call(fn=fn, tags=self.tags, metadata=self.metadata)
+        elif fn_wants_span(fn):
+            if fn_is_async(fn):
+                return AsyncTracedSpanFunction(
+                    fn=fn, tags=self.tags, metadata=self.metadata
+                )
+            else:
+                return TracedSpanFunction(fn=fn, tags=self.tags, metadata=self.metadata)
+        elif fn_is_async(fn):
+            return AsyncTracedFunction(fn=fn, tags=self.tags, metadata=self.metadata)
+        else:
+            return TracedFunction(fn=fn, tags=self.tags, metadata=self.metadata)
+
+
+@overload
+def trace(
+    __fn: None = None,
+    *,
+    tags: list[str] | None = None,
+    metadata: dict[str, str] | None = None,
+) -> TraceDecorator:
+    """Overload for providing kwargs before decorating (e.g. tags)."""
+    ...
+
+
+@overload
+def trace(  # pyright: ignore[reportOverlappingOverload]
+    __fn: AsyncContextCall[P, DepsT, FormattableT],
+    *,
+    tags: None = None,
+    metadata: None = None,
+) -> TracedAsyncContextCall[P, DepsT, FormattableT]:
+    """Overload for directly decorating an AsyncContextCall."""
+    ...
+
+
+@overload
+def trace(
+    __fn: ContextCall[P, DepsT, FormattableT],
+    *,
+    tags: None = None,
+    metadata: None = None,
+) -> TracedContextCall[P, DepsT, FormattableT]:
+    """Overload for directly decorating a ContextCall."""
+    ...
+
+
+@overload
+def trace(
+    __fn: AsyncCall[P, FormattableT],
+    *,
+    tags: None = None,
+    metadata: None = None,
+) -> TracedAsyncCall[P, FormattableT]:
+    """Overload for directly decorating an AsyncCall."""
+    ...
+
+
+@overload
+def trace(
+    __fn: Call[P, FormattableT],
+    *,
+    tags: None = None,
+    metadata: None = None,
+) -> TracedCall[P, FormattableT]:
+    """Overload for directly decorating a Call."""
+    ...
+
+
+@overload
+def trace(  # pyright: ignore[reportOverlappingOverload]
+    __fn: AsyncSpanFunction[P, R],
+    *,
+    tags: None = None,
+    metadata: None = None,
+) -> AsyncTracedSpanFunction[P, R]:
+    """Overload for directly decorating an async function with span injection."""
+    ...
+
+
+@overload
+def trace(
+    __fn: SyncSpanFunction[P, R],
+    *,
+    tags: None = None,
+    metadata: None = None,
+) -> TracedSpanFunction[P, R]:
+    """Overload for directly decorating a sync function with span injection."""
+    ...
+
+
+@overload
+def trace(
+    __fn: AsyncFunction[P, R],
+    *,
+    tags: None = None,
+    metadata: None = None,
+) -> AsyncTracedFunction[P, R]:
+    """Overload for directly (no argument) decorating an asynchronous function"""
+    ...
+
+
+@overload
+def trace(
+    __fn: SyncFunction[P, R],
+    *,
+    tags: None = None,
+    metadata: None = None,
+) -> TracedFunction[P, R]:
+    """Overload for directly (no argument) decorating a synchronous function"""
+    ...
+
+
+def trace(  # pyright: ignore[reportGeneralTypeIssues]
+    __fn: (
+        AsyncContextCall[P, DepsT, FormattableT]
+        | ContextCall[P, DepsT, FormattableT]
+        | AsyncCall[P, FormattableT]
+        | Call[P, FormattableT]
+        | AsyncSpanFunction[P, R]
+        | SyncSpanFunction[P, R]
+        | AsyncFunction[P, R]
+        | SyncFunction[P, R]
+        | None
+    ) = None,
+    *,
+    tags: Sequence[str] | None = None,
+    metadata: dict[str, str] | None = None,
+) -> (
+    TracedAsyncContextCall[P, DepsT, FormattableT]
+    | TracedContextCall[P, DepsT, FormattableT]
+    | TracedAsyncCall[P, FormattableT]
+    | TracedCall[P, FormattableT]
+    | AsyncTracedSpanFunction[P, R]
+    | TracedSpanFunction[P, R]
+    | AsyncTracedFunction[P, R]
+    | TracedFunction[P, R]
+    | TraceDecorator
+):
+    """Decorator for adding tracing capabilities to functions and LLM calls.
+
+    Creates a wrapper that enables distributed tracing, performance monitoring,
+    and execution tracking for decorated functions. When called, the decorated
+    function returns a Trace containing both the result and span info.
+
+    When decorating an @llm.call function, returns a TracedCall that wraps both
+    the call and stream methods with tracing capabilities.
+
+    If the decorated function has `trace_ctx: Span` as its first parameter,
+    the span will be injected automatically and callers should NOT pass it.
+
+    Args:
+        __fn: The function or Call object to decorate.
+        tags: Optional list of string tags to associate with traced executions.
+        metadata: Arbitrary key-value pairs for additional metadata.
+
+    Returns:
+        A decorator that wraps functions with tracing capabilities.
+
+    Example:
+        ```python
+        @ops.trace
+        def process_data(data: dict) -> dict:
+            return {"processed": data}
+
+        traced_result = process_data({"key": "value"})
+        print(traced_result.result)    # {"processed": {"key": "value"}}
+        print(traced_result.span_id)   # Access span ID
+        ```
+
+    Example:
+        ```python
+        @ops.trace
+        @llm.call("gpt-4o-mini")
+        def recommend_book(genre: str):
+            return f"Recommend a {genre} book"
+
+        # Returns Response directly (execution is still traced)
+        response = recommend_book("fantasy")
+        print(response.content)
+
+        # Use .wrapped() to get Trace[Response] with span info
+        trace = recommend_book.wrapped("fantasy")
+        print(trace.result.content)
+        print(trace.span_id)
+        ```
+
+    Example:
+        ```python
+        @ops.trace
+        def my_fn(trace_ctx: Span, arg: str) -> str:
+            trace_ctx.info(f"Processing: {arg}")
+            return arg.upper()
+
+        # Call without trace_ctx - it's injected automatically
+        result = my_fn("hello")  # Returns "HELLO"
+        ```
+    """
+    tags = tuple(sorted(set(tags or [])))
+    metadata = metadata or {}
+    if __fn is None:
+        return TraceDecorator(tags=tags, metadata=metadata)
+
+    if is_call_type(__fn):
+        return wrap_call(fn=__fn, tags=tags, metadata=metadata)
+    elif fn_wants_span(__fn):
+        if fn_is_async(__fn):
+            return AsyncTracedSpanFunction(fn=__fn, tags=tags, metadata=metadata)
+        else:
+            return TracedSpanFunction(fn=__fn, tags=tags, metadata=metadata)
+    elif fn_is_async(__fn):
+        return AsyncTracedFunction(fn=__fn, tags=tags, metadata=metadata)
+    else:
+        return TracedFunction(fn=__fn, tags=tags, metadata=metadata)

mirascope/ops/_internal/types.py

@@ -0,0 +1,13 @@
+"""Type definitions for Mirascope tracing and capabilities."""
+
+from collections.abc import Mapping, Sequence
+from typing import ParamSpec, TypeAlias
+from typing_extensions import TypeVar
+
+P = ParamSpec("P")
+R = TypeVar("R", infer_variance=True)
+
+Jsonable: TypeAlias = (
+    None | str | int | float | bool | Sequence["Jsonable"] | Mapping[str, "Jsonable"]
+)
+"""Simple type alias for JSON-serializable types."""

mirascope/ops/_internal/utils.py

@@ -0,0 +1,131 @@
+"""Internal utility functions for the Mirascope."""
+
+import inspect
+from collections.abc import Callable
+from typing import Any, TypeAlias
+
+import orjson
+
+from .protocols import P, fn_wants_span
+
+ORJSON_OPTS = (
+    orjson.OPT_NON_STR_KEYS
+    | orjson.OPT_NAIVE_UTC
+    | orjson.OPT_SERIALIZE_NUMPY
+    | orjson.OPT_SERIALIZE_DATACLASS
+    | orjson.OPT_SERIALIZE_UUID
+)
+
+PrimitiveType: TypeAlias = str | int | float | bool
+
+
+def json_dumps(obj: Any) -> str:  # noqa: ANN401
+    """Serialize Python objects to JSON using orjson."""
+    # json should be utf-8 encoded, json key should be str
+    return orjson.dumps(obj, option=ORJSON_OPTS).decode("utf-8")
+
+
+def _is_call_method(fn: Callable[..., Any]) -> bool:
+    """Check if fn is a bound method of a Call object (e.g., Call.call or Call.stream)."""
+    return (
+        hasattr(fn, "__self__") and hasattr(fn.__self__, "prompt")  # pyright: ignore[reportFunctionMemberAccess]
+    )
+
+
+def get_original_fn(fn: Callable[..., Any]) -> Callable[..., Any]:
+    """Get the original unwrapped function.
+
+    Follows the __wrapped__ chain set by copy_function_metadata() to find the
+    original function. Falls back to checking for Call methods for bound methods.
+    """
+    # Follow __wrapped__ chain if available (set by copy_function_metadata)
+    # In practice, get_original_fn is called with original functions or bound methods,
+    # so this loop is defensive code for edge cases like @functools.wraps chains.
+    while hasattr(fn, "__wrapped__"):
+        fn = fn.__wrapped__  # pyright: ignore[reportFunctionMemberAccess] # pragma: no cover
+
+    # Handle bound methods of Call objects (e.g., Call.call or Call.stream)
+    if _is_call_method(fn):
+        prompt = fn.__self__.prompt  # pyright: ignore[reportFunctionMemberAccess]
+        if hasattr(prompt, "fn"):
+            return get_original_fn(prompt.fn)
+
+    return fn
+
+
+def get_qualified_name(fn: Callable[..., Any]) -> str:
+    """Return the simplified qualified name of a function.
+
+    If the function is a bound method of a Call object (e.g., Call.call or Call.stream),
+    returns the qualified name of the original decorated function from prompt.fn,
+    suffixed with the method name (e.g., "recommend.call" or "recommend.stream").
+
+    If the function is defined locally, return the name after '<locals>.'; otherwise,
+    return the last non-empty part after splitting by '.'.
+    """
+    # Check if this is a Call method and capture the method name
+    method_suffix = f".{fn.__name__}" if _is_call_method(fn) else ""
+
+    fn = get_original_fn(fn)
+    qualified_name = fn.__qualname__
+    if "<locals>." in qualified_name:
+        base_name = qualified_name.split("<locals>.")[-1]
+    else:
+        parts = [part for part in qualified_name.split(".") if part]
+        base_name = parts[-1] if parts else qualified_name
+
+    return f"{base_name}{method_suffix}"
+
+
+def extract_arguments(
+    fn: Callable[P, Any],
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> tuple[dict[str, str], dict[str, Any]]:
+    """Returns a tuple of (arg_types, arg_values) dictionaries from function call.
+
+    If the function is a Call method (e.g., Call.call), uses the original
+    decorated function's signature to get proper parameter names and types.
+
+    If the function has `trace_ctx: Span` as first parameter (detected via
+    fn_wants_span), that parameter is skipped since it's injected by the
+    decorator and shouldn't be recorded in span attributes.
+    """
+    # Use original function signature for Call methods
+    fn = get_original_fn(fn)
+    signature = inspect.signature(fn)
+
+    # If function wants span injection, skip the first parameter
+    if fn_wants_span(fn):
+        params = list(signature.parameters.values())
+        if params:
+            remaining_params = params[1:]
+            signature = signature.replace(parameters=remaining_params)
+
+    bound_arguments = signature.bind(*args, **kwargs)
+    bound_arguments.apply_defaults()
+
+    arg_types: dict[str, str] = {}
+    arg_values: dict[str, Any] = {}
+
+    for param_name, param_value in bound_arguments.arguments.items():
+        parameter = signature.parameters[param_name]
+        if parameter.annotation != inspect.Parameter.empty:
+            if hasattr(parameter.annotation, "__name__"):
+                type_str = parameter.annotation.__name__
+            else:
+                type_str = str(parameter.annotation)
+            arg_types[param_name] = type_str
+        else:
+            arg_types[param_name] = type(param_value).__name__
+
+        if isinstance(param_value, PrimitiveType) or param_value is None:
+            arg_values[param_name] = param_value
+        else:
+            try:
+                json_dumps(param_value)
+                arg_values[param_name] = param_value
+            except (TypeError, ValueError):
+                arg_values[param_name] = repr(param_value)
+
+    return arg_types, arg_values