mirascope 2.0.0a5__py3-none-any.whl → 2.0.1__py3-none-any.whl

mirascope/ops/_internal/instrumentation/llm/serialize.py

@@ -0,0 +1,300 @@
+"""Mirascope-specific serialization for span attributes."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, Protocol
+
+from .....llm.content import (
+    Audio,
+    Base64ImageSource,
+    Document,
+    Image,
+    Text,
+    Thought,
+    ToolCall,
+    ToolOutput,
+)
+from .....llm.content.document import Base64DocumentSource, TextDocumentSource
+from .....llm.messages import AssistantMessage, Message, SystemMessage, UserMessage
+from .....llm.responses.usage import Usage
+from ...utils import json_dumps
+from .cost import calculate_cost_async, calculate_cost_sync
+
+if TYPE_CHECKING:
+    from opentelemetry.util.types import AttributeValue
+
+    from .....llm.responses.root_response import RootResponse
+    from .....llm.types import Jsonable
+
+logger = logging.getLogger(__name__)
+
+
+class SpanProtocol(Protocol):
+    """Protocol for span objects that support setting attributes."""
+
+    def set(self, **attributes: AttributeValue) -> None:
+        """Set attributes on the span."""
+        ...
+
+
+def _serialize_content_part(
+    part: Text | ToolCall | Thought | Image | Audio | Document | ToolOutput[Jsonable],
+) -> dict[str, Any]:
+    """Serialize a single content part to a dict matching the Mirascope dataclass structure."""
+    if isinstance(part, Text):
+        return {"type": "text", "text": part.text}
+    elif isinstance(part, ToolCall):
+        return {
+            "type": "tool_call",
+            "id": part.id,
+            "name": part.name,
+            "args": part.args,
+        }
+    elif isinstance(part, Thought):
+        return {"type": "thought", "thought": part.thought}
+    elif isinstance(part, ToolOutput):
+        return {
+            "type": "tool_output",
+            "id": part.id,
+            "name": part.name,
+            "result": part.result,
+        }
+    elif isinstance(part, Image):
+        if isinstance(part.source, Base64ImageSource):
+            return {
+                "type": "image",
+                "source": {
+                    "type": "base64_image_source",
+                    "mime_type": part.source.mime_type,
+                    "data": part.source.data,
+                },
+            }
+        else:  # URLImageSource
+            return {
+                "type": "image",
+                "source": {"type": "url_image_source", "url": part.source.url},
+            }
+    elif isinstance(part, Audio):
+        return {
+            "type": "audio",
+            "source": {
+                "type": "base64_audio_source",
+                "mime_type": part.source.mime_type,
+                "data": part.source.data,
+            },
+        }
+    elif isinstance(part, Document):
+        # Document has multiple source types - serialize based on actual type
+        if isinstance(part.source, Base64DocumentSource):
+            return {
+                "type": "document",
+                "source": {
+                    "type": "base64_document_source",
+                    "data": part.source.data,
+                    "media_type": part.source.media_type,
+                },
+            }
+        elif isinstance(part.source, TextDocumentSource):
+            return {
+                "type": "document",
+                "source": {
+                    "type": "text_document_source",
+                    "data": part.source.data,
+                    "media_type": part.source.media_type,
+                },
+            }
+        else:  # URLDocumentSource
+            return {
+                "type": "document",
+                "source": {
+                    "type": "url_document_source",
+                    "url": part.source.url,
+                },
+            }
+    return {"type": "unknown"}  # pragma: no cover
+
+
+def _serialize_message(message: Message) -> dict[str, Any]:
+    """Serialize a Message to a dict matching the Mirascope dataclass structure."""
+    if isinstance(message, SystemMessage):
+        return {
+            "role": "system",
+            "content": _serialize_content_part(message.content),
+        }
+    elif isinstance(message, UserMessage):
+        return {
+            "role": "user",
+            "content": [_serialize_content_part(p) for p in message.content],
+            "name": message.name,
+        }
+    elif isinstance(message, AssistantMessage):
+        return {
+            "role": "assistant",
+            "content": [_serialize_content_part(p) for p in message.content],
+            "name": message.name,
+        }
+    return {"role": "unknown"}  # pragma: no cover
+
+
+def serialize_mirascope_messages(messages: Sequence[Message]) -> str:
+    """Serialize input messages to JSON for span attributes."""
+    return json_dumps([_serialize_message(m) for m in messages])
+
+
+def serialize_mirascope_content(
+    content: Sequence[Text | ToolCall | Thought],
+) -> str:
+    """Serialize response content to JSON for span attributes."""
+    return json_dumps([_serialize_content_part(p) for p in content])
+
+
+def serialize_mirascope_usage(usage: Usage | None) -> AttributeValue | None:
+    """Serialize response usage to JSON for span attributes. Returns None if usage is None."""
+    if usage is None:
+        return None
+    return json_dumps(
+        {
+            "input_tokens": usage.input_tokens,
+            "output_tokens": usage.output_tokens,
+            "cache_read_tokens": usage.cache_read_tokens,
+            "cache_write_tokens": usage.cache_write_tokens,
+            "reasoning_tokens": usage.reasoning_tokens,
+            "total_tokens": usage.total_tokens,
+        }
+    )
+
+
+def serialize_mirascope_cost(
+    input_cost: float,
+    output_cost: float,
+    total_cost: float,
+    cache_read_cost: float | None = None,
+    cache_write_cost: float | None = None,
+) -> str:
+    """Serialize cost to JSON for span attributes.
+
+    All costs are in centicents (1 centicent = $0.0001).
+    Consumers can divide by 10,000 to get dollar amounts.
+    """
+    return json_dumps(
+        {
+            "input_cost": input_cost,
+            "output_cost": output_cost,
+            "cache_read_cost": cache_read_cost,
+            "cache_write_cost": cache_write_cost,
+            "total_cost": total_cost,
+        }
+    )
+
+
+def attach_mirascope_response(
+    span: SpanProtocol, response: RootResponse[Any, Any]
+) -> None:
+    """Attach Mirascope-specific response attributes to a span.
+
+    Sets the following attributes:
+    - mirascope.trace.output: Pretty-printed response
+    - mirascope.response.messages: Serialized input messages (excluding final assistant message)
+    - mirascope.response.content: Serialized response content
+    - mirascope.response.usage: Serialized usage (if available)
+    - mirascope.response.cost: Serialized cost (if available)
+    """
+    span.set(
+        **{
+            "mirascope.response.provider_id": response.provider_id,
+            "mirascope.response.model_id": response.model_id,
+            "mirascope.trace.output": response.pretty(),
+            "mirascope.response.messages": serialize_mirascope_messages(
+                response.messages[:-1]
+            ),
+            "mirascope.response.content": serialize_mirascope_content(response.content),
+        }
+    )
+    if (usage_json := serialize_mirascope_usage(response.usage)) is not None:
+        span.set(**{"mirascope.response.usage": usage_json})
+        logger.debug("Attached usage to span")
+    else:
+        logger.debug("No usage available, skipping cost calculation")
+
+    # Calculate and attach cost if usage is available
+    if response.usage is not None:
+        logger.debug("Attempting cost calculation (sync)")
+        cost = calculate_cost_sync(
+            response.provider_id, response.model_id, response.usage
+        )
+        if cost is not None:
+            span.set(
+                **{
+                    "mirascope.response.cost": serialize_mirascope_cost(
+                        input_cost=cost.input_cost_centicents,
+                        output_cost=cost.output_cost_centicents,
+                        total_cost=cost.total_cost_centicents,
+                        cache_read_cost=cost.cache_read_cost_centicents,
+                        cache_write_cost=cost.cache_write_cost_centicents,
+                    )
+                }
+            )
+            logger.debug(
+                "Attached cost to span: total=%s centicents", cost.total_cost_centicents
+            )
+        else:
+            logger.debug("Cost calculation returned None, not attaching cost to span")
+
+
+async def attach_mirascope_response_async(
+    span: SpanProtocol, response: RootResponse[Any, Any]
+) -> None:
+    """Attach Mirascope-specific response attributes to a span (async version).
+
+    Sets the following attributes:
+    - mirascope.trace.output: Pretty-printed response
+    - mirascope.response.messages: Serialized input messages (excluding final assistant message)
+    - mirascope.response.content: Serialized response content
+    - mirascope.response.usage: Serialized usage (if available)
+    - mirascope.response.cost: Serialized cost (if available)
+    """
+    span.set(
+        **{
+            "mirascope.response.provider_id": response.provider_id,
+            "mirascope.response.model_id": response.model_id,
+            "mirascope.trace.output": response.pretty(),
+            "mirascope.response.messages": serialize_mirascope_messages(
+                response.messages[:-1]
+            ),
+            "mirascope.response.content": serialize_mirascope_content(response.content),
+        }
+    )
+    if (usage_json := serialize_mirascope_usage(response.usage)) is not None:
+        span.set(**{"mirascope.response.usage": usage_json})
+        logger.debug("Attached usage to span (async)")
+    else:
+        logger.debug("No usage available, skipping cost calculation (async)")
+
+    # Calculate and attach cost if usage is available (async)
+    if response.usage is not None:
+        logger.debug("Attempting cost calculation (async)")
+        cost = await calculate_cost_async(
+            response.provider_id, response.model_id, response.usage
+        )
+        if cost is not None:
+            span.set(
+                **{
+                    "mirascope.response.cost": serialize_mirascope_cost(
+                        input_cost=cost.input_cost_centicents,
+                        output_cost=cost.output_cost_centicents,
+                        total_cost=cost.total_cost_centicents,
+                        cache_read_cost=cost.cache_read_cost_centicents,
+                        cache_write_cost=cost.cache_write_cost_centicents,
+                    )
+                }
+            )
+            logger.debug(
+                "Attached cost to span (async): total=%s centicents",
+                cost.total_cost_centicents,
+            )
+        else:
+            logger.debug(
+                "Cost calculation returned None, not attaching cost to span (async)"
+            )

mirascope/ops/_internal/protocols.py

@@ -3,12 +3,15 @@
 from __future__ import annotations
 
 import inspect
-from typing import Protocol
+from typing import TYPE_CHECKING, Protocol
 from typing_extensions import TypeIs
 
 from ...llm.context import Context, DepsT
 from .types import P, R
 
+if TYPE_CHECKING:
+    from .spans import Span
+
 
 class SyncFunction(Protocol[P, R]):
     """Protocol for synchronous callable functions."""
@@ -26,6 +29,30 @@ class AsyncFunction(Protocol[P, R]):
         ...  # pragma: no cover
 
 
+class SyncSpanFunction(Protocol[P, R]):
+    """Protocol for synchronous functions that receive injected Span.
+
+    Functions matching this protocol have `trace_ctx: Span` as their first
+    parameter. The trace decorator will inject the span automatically.
+    """
+
+    def __call__(self, trace_ctx: Span, *args: P.args, **kwargs: P.kwargs) -> R:
+        """The function receives a Span as first parameter."""
+        ...  # pragma: no cover
+
+
+class AsyncSpanFunction(Protocol[P, R]):
+    """Protocol for asynchronous functions that receive injected Span.
+
+    Functions matching this protocol have `trace_ctx: Span` as their first
+    parameter. The trace decorator will inject the span automatically.
+    """
+
+    async def __call__(self, trace_ctx: Span, *args: P.args, **kwargs: P.kwargs) -> R:
+        """The function receives a Span as first parameter."""
+        ...  # pragma: no cover
+
+
 class SyncContextFunction(Protocol[P, DepsT, R]):
     """Protocol for synchronous callable functions with Context parameter."""
 
@@ -49,3 +76,58 @@ def fn_is_async(
 ) -> TypeIs[AsyncFunction[P, R]]:
     """Type check to determine if a given function is asynchronous."""
     return inspect.iscoroutinefunction(fn) or inspect.isasyncgenfunction(fn)
+
+
+def fn_wants_span(
+    fn: (
+        SyncFunction[P, R]
+        | AsyncFunction[P, R]
+        | SyncSpanFunction[P, R]
+        | AsyncSpanFunction[P, R]
+    ),
+) -> TypeIs[SyncSpanFunction[P, R] | AsyncSpanFunction[P, R]]:
+    """Check if function wants Span injection as first parameter.
+
+    Returns True if the function has a first parameter named `trace_ctx`
+    with type annotation `Span`.
+    """
+    # Import here to avoid circular imports
+    from .spans import Span
+
+    try:
+        sig = inspect.signature(fn)
+    except (ValueError, TypeError):
+        return False
+
+    params = list(sig.parameters.values())
+    if not params:
+        return False
+
+    first_param = params[0]
+    if first_param.name != "trace_ctx":
+        return False
+
+    # Check annotation
+    annotation = first_param.annotation
+    if annotation is inspect.Parameter.empty:
+        return False
+
+    # Handle string annotations (forward references)
+    # The annotation could be "Span", "ops.Span", "mirascope.ops.Span", etc.
+    if isinstance(annotation, str):
+        return annotation == "Span" or annotation.endswith(".Span")
+
+    # Check by identity first
+    if annotation is Span:
+        return True
+
+    # Fallback: check by class name and module for robustness
+    # This handles cases where the same class might have different identities
+    # due to module reloading or import issues in test environments
+    if isinstance(annotation, type):
+        return (
+            annotation.__name__ == "Span"
+            and annotation.__module__ == "mirascope.ops._internal.spans"
+        )
+
+    return False

mirascope/ops/_internal/traced_calls.py

@@ -22,8 +22,10 @@ from ...llm.responses import (
 from ...llm.types import P
 from .protocols import (
     AsyncFunction,
+    AsyncSpanFunction,
     R,
     SyncFunction,
+    SyncSpanFunction,
 )
 from .traced_functions import (
     AsyncTrace,
@@ -47,6 +49,8 @@ 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]

mirascope/ops/_internal/traced_functions.py

@@ -6,26 +6,31 @@ from abc import ABC, abstractmethod
 from collections.abc import Generator
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import (
-    Any,
-    Generic,
-    Literal,
-    TypeVar,
-)
+from typing import Any, Generic, Literal, TypeVar
 
 from opentelemetry.util.types import AttributeValue
 
+from ...api.client import get_async_client, get_sync_client
 from ...llm.context import Context, DepsT
 from ...llm.responses.root_response import RootResponse
+from .instrumentation.llm.serialize import attach_mirascope_response
 from .protocols import (
     AsyncContextFunction,
     AsyncFunction,
+    AsyncSpanFunction,
     SyncContextFunction,
     SyncFunction,
+    SyncSpanFunction,
 )
 from .spans import Span
 from .types import Jsonable, P, R
-from .utils import PrimitiveType, extract_arguments, get_qualified_name, json_dumps
+from .utils import (
+    PrimitiveType,
+    extract_arguments,
+    get_original_fn,
+    get_qualified_name,
+    json_dumps,
+)
 
 FunctionT = TypeVar(
     "FunctionT",
@@ -46,6 +51,7 @@ def record_result_to_span(span: Span, result: object) -> None:
         output: str | int | float | bool = result
     elif isinstance(result, RootResponse):
         output = result.pretty()
+        attach_mirascope_response(span, result)
     else:
         try:
             output = json_dumps(result)
@@ -91,7 +97,16 @@ class Trace(_BaseTrace[R]):
         metadata: dict[str, Jsonable] | None = None,
     ) -> None:
         """Annotates the current trace span."""
-        raise NotImplementedError("Trace.annotate not yet implemented")
+        if self.span.span_id is None or self.span.trace_id is None:
+            return
+
+        get_sync_client().annotations.create(
+            otel_span_id=self.span.span_id,
+            otel_trace_id=self.span.trace_id,
+            label=label,
+            reasoning=reasoning,
+            metadata=metadata,
+        )
 
     def tag(self, *tags: str) -> None:
         """Adds given tags to the current trace span."""
@@ -113,12 +128,21 @@ class AsyncTrace(_BaseTrace[R]):
     async def annotate(
         self,
         *,
-        label: str | None = None,
-        data: dict[str, Jsonable] | None = None,
+        label: Literal["pass", "fail"],
         reasoning: str | None = None,
+        metadata: dict[str, Jsonable] | None = None,
     ) -> None:
         """Annotates the current trace span."""
-        raise NotImplementedError("AsyncTrace.annotate not yet implemented")
+        if self.span.span_id is None or self.span.trace_id is None:
+            return
+
+        await get_async_client().annotations.create(
+            otel_span_id=self.span.span_id,
+            otel_trace_id=self.span.trace_id,
+            label=label,
+            reasoning=reasoning,
+            metadata=metadata,
+        )
 
     async def tag(self, *tags: str) -> None:
         """Adds given tags to the current trace span."""
@@ -154,7 +178,8 @@ class _BaseFunction(Generic[P, R, FunctionT], ABC):
     def __post_init__(self) -> None:
         """Initialize additional attributes after dataclass init."""
         self._qualified_name = get_qualified_name(self.fn)
-        self._module_name = getattr(self.fn, "__module__", "")
+        original_fn = get_original_fn(self.fn)
+        self._module_name = getattr(original_fn, "__module__", "")
 
 
 @dataclass(kw_only=True)
@@ -392,3 +417,107 @@ class AsyncTracedContextFunction(BaseTracedAsyncContextFunction[P, DepsT, R]):
             result = await self.fn(ctx, *args, **kwargs)
             record_result_to_span(span, result)
             return AsyncTrace(result=result, span=span)
+
+
+@dataclass(kw_only=True)
+class BaseSyncTracedSpanFunction(_BaseTracedFunction[P, R, SyncSpanFunction[P, R]]):
+    """Abstract base class for synchronous traced span function wrappers."""
+
+    _is_async: bool = field(default=False, init=False)
+    """Whether the wrapped function is asynchronous."""
+
+    @abstractmethod
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        """Returns the result of the traced function directly."""
+        ...
+
+    @abstractmethod
+    def wrapped(self, *args: P.args, **kwargs: P.kwargs) -> Trace[R]:
+        """Returns the trace containing the function result and span info."""
+        ...
+
+
+@dataclass(kw_only=True)
+class TracedSpanFunction(BaseSyncTracedSpanFunction[P, R]):
+    """Wrapper for synchronous functions that receive Span as first parameter.
+
+    The external interface does NOT include `trace_ctx` - it is injected
+    automatically by the decorator when calling the inner function.
+
+    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
+        result = my_fn("hello")  # Returns "HELLO"
+        ```
+    """
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        """Returns the result of the traced function directly."""
+        with self._span(*args, **kwargs) as span:
+            result = self.fn(span, *args, **kwargs)
+            record_result_to_span(span, result)
+            return result
+
+    def wrapped(self, *args: P.args, **kwargs: P.kwargs) -> Trace[R]:
+        """Returns the trace containing the function result and span info."""
+        with self._span(*args, **kwargs) as span:
+            result = self.fn(span, *args, **kwargs)
+            record_result_to_span(span, result)
+            return Trace(result=result, span=span)
+
+
+@dataclass(kw_only=True)
+class BaseAsyncTracedSpanFunction(_BaseTracedFunction[P, R, AsyncSpanFunction[P, R]]):
+    """Abstract base class for asynchronous traced span function wrappers."""
+
+    _is_async: bool = field(default=True, init=False)
+    """Whether the wrapped function is asynchronous."""
+
+    @abstractmethod
+    async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        """Returns the result of the traced function directly."""
+        ...
+
+    @abstractmethod
+    async def wrapped(self, *args: P.args, **kwargs: P.kwargs) -> AsyncTrace[R]:
+        """Returns the trace containing the function result and span info."""
+        ...
+
+
+@dataclass(kw_only=True)
+class AsyncTracedSpanFunction(BaseAsyncTracedSpanFunction[P, R]):
+    """Wrapper for asynchronous functions that receive Span as first parameter.
+
+    The external interface does NOT include `trace_ctx` - it is injected
+    automatically by the decorator when calling the inner function.
+
+    Example:
+        ```python
+        @ops.trace
+        async 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
+        result = await my_fn("hello")  # Returns "HELLO"
+        ```
+    """
+
+    async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        """Returns the result of the traced function directly."""
+        with self._span(*args, **kwargs) as span:
+            result = await self.fn(span, *args, **kwargs)
+            record_result_to_span(span, result)
+            return result
+
+    async def wrapped(self, *args: P.args, **kwargs: P.kwargs) -> AsyncTrace[R]:
+        """Returns the trace containing the function result and span info."""
+        with self._span(*args, **kwargs) as span:
+            result = await self.fn(span, *args, **kwargs)
+            record_result_to_span(span, result)
+            return AsyncTrace(result=result, span=span)