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

mirascope/ops/_internal/spans.py

@@ -0,0 +1,232 @@
+"""Explicit span management utilities for `mirascope.ops`."""
+
+from __future__ import annotations
+
+import logging
+from contextvars import Token
+from typing import TYPE_CHECKING, Any
+
+from opentelemetry import context as otel_context, trace as otel_trace
+from opentelemetry.trace import (
+    SpanContext,
+    Status,
+    StatusCode,
+    format_span_id,
+    format_trace_id,
+)
+from opentelemetry.util.types import AttributeValue
+
+from .session import current_session
+from .utils import json_dumps
+
+if TYPE_CHECKING:
+    from opentelemetry.context import Context
+
+logger = logging.getLogger("mirascope.ops")
+_warned_noop = False
+
+
+class Span:
+    """Context-managed span for explicit tracing.
+
+    Creates a child span within the current trace context. Acts as a no-op
+    if tracing is not configured.
+    """
+
+    def __init__(self, name: str, **attributes: AttributeValue) -> None:
+        """Initialize a new span with the given name.
+
+        Args:
+            name: Name for the span.
+            **attributes: Initial attributes to set on the span.
+        """
+        self._name = name
+        self._initial_attributes = attributes
+        self._span: otel_trace.Span | None = None
+        self._token: Token[Context] | None = None
+        self._is_noop = True
+        self._finished = False
+
+    def __enter__(self) -> Span:
+        """Enter the span context.
+
+        Returns:
+            This span instance for use within the context.
+        """
+        tracer = otel_trace.get_tracer("mirascope.ops")
+        self._span = tracer.start_span(self._name)
+
+        if self._span.__class__.__name__ == "NonRecordingSpan":
+            self._is_noop = True
+            self._span = None
+            global _warned_noop
+            if not _warned_noop:
+                logger.warning(
+                    f"mirascope tracing is not configured; Span('{self._name}') is a no-op."
+                )
+                _warned_noop = True
+        else:
+            self._is_noop = False
+            self._span.set_attribute("mirascope.type", "trace")
+
+            session_ctx = current_session()
+            if session_ctx is not None:
+                self._span.set_attribute("mirascope.ops.session.id", session_ctx.id)
+                if session_ctx.attributes is not None:
+                    self._span.set_attribute(
+                        "mirascope.ops.session.attributes",
+                        json_dumps(session_ctx.attributes),
+                    )
+
+            self._token = otel_context.attach(
+                otel_trace.set_span_in_context(self._span)
+            )
+
+            if self._initial_attributes:
+                self.set(**self._initial_attributes)
+
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: Any,  # noqa: ANN401
+    ) -> None:
+        """Exit the span context.
+
+        Args:
+            exc_type: Exception type if an exception was raised.
+            exc: Exception instance if an exception was raised.
+            tb: Traceback if an exception was raised.
+        """
+        if self._span and not self._finished:
+            if exc is not None:
+                self._span.record_exception(exc)
+                self._span.set_status(Status(StatusCode.ERROR))
+            self.finish()
+
+    def set(self, **attributes: AttributeValue) -> None:
+        """Set attributes on the current span.
+
+        Args:
+            **attributes: Key-value pairs to set as span attributes.
+        """
+        if self._span and not self._finished:
+            for key, value in attributes.items():
+                self._span.set_attribute(key, value)
+
+    def event(self, name: str, **attributes: AttributeValue) -> None:
+        """Record an event within the span.
+
+        Args:
+            name: Name of the event.
+            **attributes: Event attributes as key-value pairs.
+        """
+        if self._span and not self._finished:
+            self._span.add_event(name, attributes=attributes)
+
+    def debug(self, message: str, **additional_attributes: AttributeValue) -> None:
+        """Log a debug message within the span.
+
+        Args:
+            message: Debug message text.
+            **additional_attributes: Additional structured attributes for the log entry.
+        """
+        self.event("debug", level="debug", message=message, **additional_attributes)
+
+    def info(self, message: str, **additional_attributes: AttributeValue) -> None:
+        """Log an info message within the span.
+
+        Args:
+            message: Info message text.
+            **additional_attributes: Additional structured attributes for the log entry.
+        """
+        self.event("info", level="info", message=message, **additional_attributes)
+
+    def warning(self, message: str, **additional_attributes: AttributeValue) -> None:
+        """Log a warning message within the span.
+
+        Args:
+            message: Warning message text.
+            **additional_attributes: Additional structured attributes for the log entry.
+        """
+        self.event("warning", level="warning", message=message, **additional_attributes)
+
+    def error(self, message: str, **additional_attributes: AttributeValue) -> None:
+        """Log an error message within the span.
+
+        Args:
+            message: Error message text.
+            **additional_attributes: Additional structured attributes for the log entry.
+        """
+        self.event("error", level="error", message=message, **additional_attributes)
+        if self._span and not self._finished:
+            self._span.set_status(Status(StatusCode.ERROR))
+
+    def critical(self, message: str, **additional_attributes: AttributeValue) -> None:
+        """Log a critical message within the span.
+
+        Args:
+            message: Critical message text.
+            **additional_attributes: Additional structured attributes for the log entry.
+        """
+        self.event("critical", level="error", message=message, **additional_attributes)
+        if self._span and not self._finished:
+            self._span.set_status(Status(StatusCode.ERROR))
+
+    def finish(self) -> None:
+        """Explicitly finish the span."""
+        if not self._finished:
+            self._finished = True
+            if self._span:
+                self._span.end()
+            if self._token:
+                otel_context.detach(self._token)
+                self._token = None
+
+    @property
+    def span_id(self) -> str | None:
+        """Get the span ID if available.
+
+        Returns:
+            The span ID or None if not available.
+        """
+        if span_context := self.span_context:
+            return format_span_id(span_context.span_id)
+        return None
+
+    @property
+    def trace_id(self) -> str | None:
+        if span_context := self.span_context:
+            return format_trace_id(span_context.trace_id)
+        return None
+
+    @property
+    def is_noop(self) -> bool:
+        """Check if this is a no-op span.
+
+        Returns:
+            True if tracing is disabled, False otherwise.
+        """
+        return self._is_noop
+
+    @property
+    def span_context(self) -> SpanContext | None:
+        """Return the span context if available."""
+        if otel_span := self._span:
+            return otel_span.get_span_context()
+        return None
+
+
+def span(name: str, **attributes: AttributeValue) -> Span:
+    """Create a new span context manager.
+
+    Args:
+        name: Name for the new span.
+        **attributes: Initial attributes to set on the span.
+
+    Returns:
+        A Span context manager that creates a child span when entered.
+    """
+    return Span(name, **attributes)

mirascope/ops/_internal/traced_calls.py

@@ -0,0 +1,389 @@
+"""Traced call wrappers for @ops.trace decorated @llm.call functions."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Generic, TypeVar
+from typing_extensions import TypeIs
+
+from ..._utils import copy_function_metadata
+from ...llm.calls import AsyncCall, AsyncContextCall, Call, ContextCall
+from ...llm.context import Context, DepsT
+from ...llm.formatting import FormattableT
+from ...llm.responses import (
+    AsyncContextResponse,
+    AsyncContextStreamResponse,
+    AsyncResponse,
+    AsyncStreamResponse,
+    ContextResponse,
+    ContextStreamResponse,
+    Response,
+    StreamResponse,
+)
+from ...llm.types import P
+from .protocols import (
+    AsyncFunction,
+    AsyncSpanFunction,
+    R,
+    SyncFunction,
+    SyncSpanFunction,
+)
+from .traced_functions import (
+    AsyncTrace,
+    AsyncTracedContextFunction,
+    AsyncTracedFunction,
+    Trace,
+    TracedContextFunction,
+    TracedFunction,
+)
+from .utils import get_original_fn
+
+CallT = TypeVar(
+    "CallT",
+    bound=Call[..., Any]
+    | AsyncCall[..., Any]
+    | ContextCall[..., Any, Any]
+    | AsyncContextCall[..., Any, Any],
+)
+
+
+def is_call_type(
+    fn: (
+        SyncFunction[P, R]
+        | AsyncFunction[P, R]
+        | SyncSpanFunction[P, R]
+        | AsyncSpanFunction[P, R]
+        | ContextCall[P, DepsT, FormattableT]
+        | AsyncContextCall[P, DepsT, FormattableT]
+        | Call[P, FormattableT]
+        | AsyncCall[P, FormattableT]
+    ),
+) -> TypeIs[
+    ContextCall[P, DepsT, FormattableT]
+    | AsyncContextCall[P, DepsT, FormattableT]
+    | Call[P, FormattableT]
+    | AsyncCall[P, FormattableT]
+]:
+    """Check if fn is any of the Call types."""
+    return isinstance(fn, Call | AsyncCall | ContextCall | AsyncContextCall)
+
+
+def wrap_call(
+    fn: (
+        ContextCall[P, DepsT, FormattableT]
+        | AsyncContextCall[P, DepsT, FormattableT]
+        | Call[P, FormattableT]
+        | AsyncCall[P, FormattableT]
+    ),
+    tags: tuple[str, ...],
+    metadata: dict[str, str] | None = None,
+) -> (
+    TracedContextCall[P, DepsT, FormattableT]
+    | TracedAsyncContextCall[P, DepsT, FormattableT]
+    | TracedCall[P, FormattableT]
+    | TracedAsyncCall[P, FormattableT]
+):
+    """Wrap a Call object with the appropriate TracedCall type."""
+    metadata = metadata or {}
+    if isinstance(fn, AsyncContextCall):
+        return TracedAsyncContextCall(_call=fn, tags=tags, metadata=metadata)
+    elif isinstance(fn, ContextCall):
+        return TracedContextCall(_call=fn, tags=tags, metadata=metadata)
+    elif isinstance(fn, AsyncCall):
+        return TracedAsyncCall(_call=fn, tags=tags, metadata=metadata)
+    else:
+        return TracedCall(_call=fn, tags=tags, metadata=metadata)
+
+
+@dataclass(kw_only=True)
+class _BaseTracedCall(Generic[CallT]):
+    """Wrapper for traced Call objects."""
+
+    _call: CallT
+    """The original unwrapped Call object."""
+
+    tags: tuple[str, ...]
+    """Tags to be associated with traced calls."""
+
+    metadata: dict[str, str] = field(default_factory=dict)
+    """Arbitrary key-value pairs for additional metadata."""
+
+    __name__: str = field(init=False, repr=False, default="")
+    """The name of the underlying function (preserved for decorator stacking)."""
+
+    def __post_init__(self) -> None:
+        """Preserve standard function attributes for decorator stacking."""
+        original_fn = get_original_fn(self._call.prompt.fn)
+        copy_function_metadata(self, original_fn)
+
+
+@dataclass(kw_only=True)
+class TracedCall(_BaseTracedCall[Call[P, FormattableT]]):
+    """Traced wrapper for traced synchronous Call objects.
+
+    When @ops.trace decorates an @llm.call, it returns a TracedCall that wraps
+    the call and stream methods with tracing capabilities.
+
+    Example:
+        ```python
+        @ops.trace
+        @llm.call("gpt-4o-mini")
+        def recommend_book(genre: str):
+            return f"Recommend a {genre} book"
+
+        # Returns Response directly (but execution is traced)
+        response = recommend_book("fantasy")
+        print(response.content)
+
+        # Same as __call__
+        response = recommend_book.call("fantasy")
+
+        # Streaming returns StreamResponse (traced)
+        stream = recommend_book.stream("fantasy")
+        for chunk in stream:
+            print(chunk)
+
+        # Use wrapped to get Trace[Response] with span info
+        trace = recommend_book.wrapped("fantasy")
+        print(trace.result.content)
+        print(trace.span_id)
+
+        # Use wrapped_stream to get Trace[StreamResponse]
+        trace = recommend_book.wrapped_stream("fantasy")
+        ```
+    """
+
+    call: TracedFunction[P, Response | Response[FormattableT]] = field(init=False)
+    """TracedFunction wrapping the call method."""
+
+    stream: TracedFunction[P, StreamResponse | StreamResponse[FormattableT]] = field(
+        init=False
+    )
+    """TracedFunction wrapping the stream method."""
+
+    def __post_init__(self) -> None:
+        """Initialize TracedFunction wrappers for call and stream methods."""
+        super().__post_init__()
+        self.call = TracedFunction(
+            fn=self._call.call, tags=self.tags, metadata=self.metadata
+        )
+        self.stream = TracedFunction(
+            fn=self._call.stream, tags=self.tags, metadata=self.metadata
+        )
+
+    def __call__(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> Response | Response[FormattableT]:
+        """Call the traced function and return Response directly."""
+        return self.call(*args, **kwargs)
+
+    def wrapped(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> Trace[Response | Response[FormattableT]]:
+        """Call the traced function and return a wrapped Response."""
+        return self.call.wrapped(*args, **kwargs)
+
+    def wrapped_stream(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> Trace[StreamResponse | StreamResponse[FormattableT]]:
+        """Stream the traced function and return a wrapped StreamResponse."""
+        return self.stream.wrapped(*args, **kwargs)
+
+
+@dataclass(kw_only=True)
+class TracedAsyncCall(_BaseTracedCall[AsyncCall[P, FormattableT]]):
+    """Traced wrapper for traced asynchronous AsyncCall objects.
+
+    Example:
+        ```python
+        @ops.trace
+        @llm.call("gpt-4o-mini")
+        async def recommend_book(genre: str):
+            return f"Recommend a {genre} book"
+
+        # Returns AsyncResponse directly (but execution is traced)
+        response = await recommend_book("fantasy")
+        print(response.content)
+
+        # Use wrapped to get AsyncTrace[AsyncResponse] with span info
+        trace = await recommend_book.wrapped("fantasy")
+        print(trace.result.content)
+        print(trace.span_id)
+        ```
+    """
+
+    call: AsyncTracedFunction[P, AsyncResponse | AsyncResponse[FormattableT]] = field(
+        init=False
+    )
+    """AsyncTracedFunction wrapping the call method."""
+
+    stream: AsyncTracedFunction[
+        P, AsyncStreamResponse | AsyncStreamResponse[FormattableT]
+    ] = field(init=False)
+    """AsyncTracedFunction wrapping the stream method."""
+
+    def __post_init__(self) -> None:
+        """Initialize AsyncTracedFunction wrappers for call and stream methods."""
+        super().__post_init__()
+        self.call = AsyncTracedFunction(
+            fn=self._call.call, tags=self.tags, metadata=self.metadata
+        )
+        self.stream = AsyncTracedFunction(
+            fn=self._call.stream, tags=self.tags, metadata=self.metadata
+        )
+
+    async def __call__(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncResponse | AsyncResponse[FormattableT]:
+        """Call the traced function and return AsyncResponse directly."""
+        return await self.call(*args, **kwargs)
+
+    async def wrapped(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncTrace[AsyncResponse | AsyncResponse[FormattableT]]:
+        """Call the traced function and return a wrapped Response."""
+        return await self.call.wrapped(*args, **kwargs)
+
+    async def wrapped_stream(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncTrace[AsyncStreamResponse | AsyncStreamResponse[FormattableT]]:
+        """Stream the traced function and return a wrapped StreamResponse."""
+        return await self.stream.wrapped(*args, **kwargs)
+
+
+@dataclass(kw_only=True)
+class TracedContextCall(_BaseTracedCall[ContextCall[P, DepsT, FormattableT]]):
+    """Traced wrapper for traced synchronous ContextCall objects.
+
+    Example:
+        ```python
+        @ops.trace
+        @llm.call("gpt-4o-mini")
+        def recommend_book(ctx: llm.Context[str], genre: str):
+            return f"{ctx.deps} Recommend a {genre} book"
+
+        ctx = llm.Context(deps="As a librarian,")
+
+        # Returns ContextResponse directly (but execution is traced)
+        response = recommend_book(ctx, "fantasy")
+        print(response.content)
+
+        # Use wrapped to get Trace[ContextResponse] with span info
+        trace = recommend_book.wrapped(ctx, "fantasy")
+        print(trace.result.content)
+        ```
+    """
+
+    call: TracedContextFunction[
+        P, DepsT, ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]
+    ] = field(init=False)
+    """TracedContextFunction wrapping the call method."""
+
+    stream: TracedContextFunction[
+        P,
+        DepsT,
+        ContextStreamResponse[DepsT, None] | ContextStreamResponse[DepsT, FormattableT],
+    ] = field(init=False)
+    """TracedContextFunction wrapping the stream method."""
+
+    def __post_init__(self) -> None:
+        """Initialize TracedContextFunction wrappers for call and stream methods."""
+        super().__post_init__()
+        self.call = TracedContextFunction(
+            fn=self._call.call, tags=self.tags, metadata=self.metadata
+        )
+        self.stream = TracedContextFunction(
+            fn=self._call.stream, tags=self.tags, metadata=self.metadata
+        )
+
+    def __call__(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]:
+        """Call the traced function and return ContextResponse directly."""
+        return self.call(ctx, *args, **kwargs)
+
+    def wrapped(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> Trace[ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]]:
+        """Call the traced function and return a wrapped Response."""
+        return self.call.wrapped(ctx, *args, **kwargs)
+
+    def wrapped_stream(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> Trace[
+        ContextStreamResponse[DepsT, None] | ContextStreamResponse[DepsT, FormattableT]
+    ]:
+        """Stream the traced function and return a wrapped StreamResponse."""
+        return self.stream.wrapped(ctx, *args, **kwargs)
+
+
+@dataclass(kw_only=True)
+class TracedAsyncContextCall(_BaseTracedCall[AsyncContextCall[P, DepsT, FormattableT]]):
+    """Traced wrapper for traced asynchronous AsyncContextCall objects.
+
+    Example:
+        ```python
+        @ops.trace
+        @llm.call("gpt-4o-mini")
+        async def recommend_book(ctx: llm.Context[str], genre: str):
+            return f"{ctx.deps} Recommend a {genre} book"
+
+        ctx = llm.Context(deps="As a librarian,")
+
+        # Returns AsyncContextResponse directly (but execution is traced)
+        response = await recommend_book(ctx, "fantasy")
+        print(response.content)
+
+        # Use wrapped to get AsyncTrace[AsyncContextResponse] with span info
+        trace = await recommend_book.wrapped(ctx, "fantasy")
+        print(trace.result.content)
+        ```
+    """
+
+    call: AsyncTracedContextFunction[
+        P,
+        DepsT,
+        AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT],
+    ] = field(init=False)
+    """AsyncTracedContextFunction wrapping the call method."""
+
+    stream: AsyncTracedContextFunction[
+        P,
+        DepsT,
+        AsyncContextStreamResponse[DepsT, None]
+        | AsyncContextStreamResponse[DepsT, FormattableT],
+    ] = field(init=False)
+    """AsyncTracedContextFunction wrapping the stream method."""
+
+    def __post_init__(self) -> None:
+        """Initialize AsyncTracedContextFunction wrappers for call and stream methods."""
+        super().__post_init__()
+        self.call = AsyncTracedContextFunction(
+            fn=self._call.call, tags=self.tags, metadata=self.metadata
+        )
+        self.stream = AsyncTracedContextFunction(
+            fn=self._call.stream, tags=self.tags, metadata=self.metadata
+        )
+
+    async def __call__(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT]:
+        """Call the traced function and return AsyncContextResponse directly."""
+        return await self.call(ctx, *args, **kwargs)
+
+    async def wrapped(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncTrace[
+        AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT]
+    ]:
+        """Call the traced function and return a wrapped Response."""
+        return await self.call.wrapped(ctx, *args, **kwargs)
+
+    async def wrapped_stream(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> AsyncTrace[
+        AsyncContextStreamResponse[DepsT, None]
+        | AsyncContextStreamResponse[DepsT, FormattableT]
+    ]:
+        """Stream the traced function and return a wrapped StreamResponse."""
+        return await self.stream.wrapped(ctx, *args, **kwargs)