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

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

@@ -1,1238 +1,68 @@
-"""OpenTelemetry GenAI instrumentation for `mirascope.llm`."""
+"""OpenTelemetry GenAI instrumentation for `mirascope.llm`.
+
+This module provides the public instrument_llm() and uninstrument_llm() functions
+that enable/disable instrumentation for Model calls and Response.resume methods.
+"""
 
 from __future__ import annotations
 
 import os
-import weakref
-from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
-from contextlib import AbstractContextManager, contextmanager
-from dataclasses import dataclass
-from functools import wraps
-from types import TracebackType
-from typing import TYPE_CHECKING, Protocol, TypeAlias, overload, runtime_checkable
-from typing_extensions import TypeIs
-
-from opentelemetry.semconv._incubating.attributes import (
-    gen_ai_attributes as GenAIAttributes,
-)
-from opentelemetry.semconv.attributes import (
-    error_attributes as ErrorAttributes,
-)
 
-from .....llm.context import Context, DepsT
-from .....llm.formatting import Format, FormattableT
-from .....llm.messages import Message
-from .....llm.models import Model
-from .....llm.providers import Params, ProviderId
-from .....llm.providers.model_id import ModelId
-from .....llm.responses import (
-    AsyncContextResponse,
-    AsyncContextStreamResponse,
-    AsyncResponse,
-    ContextResponse,
-    ContextStreamResponse,
-    Response,
-    StreamResponse,
-    StreamResponseChunk,
-)
-from .....llm.responses.root_response import RootResponse
-from .....llm.tools import (
-    AnyToolFn,
-    AnyToolSchema,
-    AsyncContextTool,
-    AsyncContextToolkit,
-    AsyncTool,
-    AsyncToolkit,
-    ContextTool,
-    ContextToolkit,
-    Tool,
-    Toolkit,
-)
-from .....llm.tools.tool_schema import ToolSchema
-from .....llm.tools.toolkit import BaseToolkit, ToolkitT
-from .....llm.types import Jsonable
-from ...configuration import (
-    get_tracer,
+from opentelemetry import trace as otel_trace
+
+from ...configuration import get_tracer
+from .model import (
+    unwrap_model_call,
+    unwrap_model_call_async,
+    unwrap_model_context_call,
+    unwrap_model_context_call_async,
+    unwrap_model_context_resume,
+    unwrap_model_context_resume_async,
+    unwrap_model_context_resume_stream,
+    unwrap_model_context_resume_stream_async,
+    unwrap_model_context_stream,
+    unwrap_model_context_stream_async,
+    unwrap_model_resume,
+    unwrap_model_resume_async,
+    unwrap_model_resume_stream,
+    unwrap_model_resume_stream_async,
+    unwrap_model_stream,
+    unwrap_model_stream_async,
+    wrap_model_call,
+    wrap_model_call_async,
+    wrap_model_context_call,
+    wrap_model_context_call_async,
+    wrap_model_context_resume,
+    wrap_model_context_resume_async,
+    wrap_model_context_resume_stream,
+    wrap_model_context_resume_stream_async,
+    wrap_model_context_stream,
+    wrap_model_context_stream_async,
+    wrap_model_resume,
+    wrap_model_resume_async,
+    wrap_model_resume_stream,
+    wrap_model_resume_stream_async,
+    wrap_model_stream,
+    wrap_model_stream_async,
 )
-from ...utils import json_dumps
-from .encode import (
-    map_finish_reason,
-    snapshot_from_root_response,
-    split_request_messages,
-)
-
-# TODO: refactor alongside all other import error handling to provide nice error messages
-try:
-    from opentelemetry import trace as otel_trace
-    from opentelemetry.trace import SpanKind, Status, StatusCode
-except ImportError:  # pragma: no cover
-    if not TYPE_CHECKING:
-        otel_trace = None
-        SpanKind = None
-        StatusCode = None
-        Status = None
-
-if TYPE_CHECKING:
-    from opentelemetry import trace as otel_trace
-    from opentelemetry.trace import (
-        Span,
-        SpanKind,
-        Status,
-        StatusCode,
-        Tracer,
-    )
-    from opentelemetry.util.types import AttributeValue
-
-    from . import gen_ai_types
-else:
-    AttributeValue = None
-    Span = None
-    Tracer = None
-
-
-ToolsParam: TypeAlias = (
-    Sequence[ToolSchema[AnyToolFn]] | BaseToolkit[AnyToolSchema] | None
+from .response import (
+    unwrap_async_context_response_resume,
+    unwrap_async_context_stream_response_resume,
+    unwrap_async_response_resume,
+    unwrap_async_stream_response_resume,
+    unwrap_context_response_resume,
+    unwrap_context_stream_response_resume,
+    unwrap_response_resume,
+    unwrap_stream_response_resume,
+    wrap_async_context_response_resume,
+    wrap_async_context_stream_response_resume,
+    wrap_async_response_resume,
+    wrap_async_stream_response_resume,
+    wrap_context_response_resume,
+    wrap_context_stream_response_resume,
+    wrap_response_resume,
+    wrap_stream_response_resume,
 )
-FormatParam: TypeAlias = Format[FormattableT] | None
-ParamsDict: TypeAlias = Mapping[str, str | int | float | bool | Sequence[str] | None]
-SpanAttributes: TypeAlias = Mapping[str, AttributeValue]
-AttributeSetter: TypeAlias = Callable[[str, AttributeValue], None]
-ParamsValue = str | int | float | bool | Sequence[str] | None
-
-
-@dataclass(slots=True)
-class SpanContext:
-    """Container for a GenAI span and its associated dropped parameters."""
-
-    span: Span | None
-    """The active span, if any."""
-
-    dropped_params: dict[str, Jsonable]
-    """Parameters that could not be recorded as span attributes."""
-
-
-@runtime_checkable
-class Identifiable(Protocol):
-    """Protocol for objects with an optional ID attribute."""
-
-    id: str | None
-    """Optional ID attribute."""
-
-
-_PARAM_ATTRIBUTE_MAP: Mapping[str, str] = {
-    "temperature": GenAIAttributes.GEN_AI_REQUEST_TEMPERATURE,
-    "max_tokens": GenAIAttributes.GEN_AI_REQUEST_MAX_TOKENS,
-    "max_output_tokens": GenAIAttributes.GEN_AI_REQUEST_MAX_TOKENS,
-    "max_completion_tokens": GenAIAttributes.GEN_AI_REQUEST_MAX_TOKENS,
-    "top_p": GenAIAttributes.GEN_AI_REQUEST_TOP_P,
-    "top_k": GenAIAttributes.GEN_AI_REQUEST_TOP_K,
-    "frequency_penalty": GenAIAttributes.GEN_AI_REQUEST_FREQUENCY_PENALTY,
-    "presence_penalty": GenAIAttributes.GEN_AI_REQUEST_PRESENCE_PENALTY,
-    "seed": GenAIAttributes.GEN_AI_REQUEST_SEED,
-    "stop_sequences": GenAIAttributes.GEN_AI_REQUEST_STOP_SEQUENCES,
-    "stop": GenAIAttributes.GEN_AI_REQUEST_STOP_SEQUENCES,
-    "n": GenAIAttributes.GEN_AI_REQUEST_CHOICE_COUNT,
-    "choice_count": GenAIAttributes.GEN_AI_REQUEST_CHOICE_COUNT,
-}
-
-
-def _record_exception(span: Span, exc: Exception) -> None:
-    """Record exception details on span following OpenTelemetry semantic conventions."""
-    span.record_exception(exc)
-    span.set_attribute(ErrorAttributes.ERROR_TYPE, exc.__class__.__name__)
-    error_message = str(exc)
-    if error_message:
-        span.set_attribute("error.message", error_message)
-    span.set_status(Status(StatusCode.ERROR, error_message))
-
-
-def _infer_output_type(format_obj: FormatParam) -> str:
-    """Infer the GenAI output type from the format parameter."""
-    if format_obj is None:
-        return GenAIAttributes.GenAiOutputTypeValues.TEXT.value
-    return GenAIAttributes.GenAiOutputTypeValues.JSON.value
-
-
-def _apply_param_attributes(
-    attrs: dict[str, AttributeValue], params: ParamsDict
-) -> None:
-    """Apply model parameters as span attributes."""
-    if not params:
-        return
-
-    for key, attr_key in _PARAM_ATTRIBUTE_MAP.items():
-        if key not in params:
-            continue
-        value = params[key]
-        if value is None:
-            continue
-        if key in {"stop", "stop_sequences"} and isinstance(value, str):
-            value = [value]
-        attrs[attr_key] = value
-
-
-def _set_json_attribute(
-    setter: AttributeSetter,
-    *,
-    key: str,
-    payload: (
-        gen_ai_types.SystemInstructions
-        | gen_ai_types.InputMessages
-        | gen_ai_types.OutputMessages
-    ),
-) -> None:
-    """Assign a JSON-serialized attribute to a span."""
-    if not payload:
-        return
-    setter(key, json_dumps(payload))
-
-
-def _assign_request_message_attributes(
-    setter: AttributeSetter,
-    *,
-    messages: Sequence[Message],
-) -> None:
-    """Assign request message attributes to a span."""
-    system_payload, input_payload = split_request_messages(messages)
-    _set_json_attribute(
-        setter,
-        key=GenAIAttributes.GEN_AI_SYSTEM_INSTRUCTIONS,
-        payload=system_payload,
-    )
-    _set_json_attribute(
-        setter,
-        key=GenAIAttributes.GEN_AI_INPUT_MESSAGES,
-        payload=input_payload,
-    )
-
-
-def _collect_tool_schemas(
-    tools: Sequence[ToolSchema[AnyToolFn]] | BaseToolkit[AnyToolSchema],
-) -> list[ToolSchema[AnyToolFn]]:
-    """Collect ToolSchema instances from a tools parameter."""
-    iterable = list(tools.tools) if isinstance(tools, BaseToolkit) else list(tools)
-    schemas: list[ToolSchema[AnyToolFn]] = []
-    for tool in iterable:
-        if isinstance(tool, ToolSchema):
-            schemas.append(tool)
-    return schemas
-
-
-def _serialize_tool_definitions(
-    tools: ToolsParam,
-    format: FormatParam = None,
-) -> str | None:
-    """Serialize tool definitions to JSON for span attributes."""
-    if tools is None:
-        tool_schemas: list[ToolSchema[AnyToolFn]] = []
-    else:
-        tool_schemas = _collect_tool_schemas(tools)
-
-    if isinstance(format, Format) and format.mode == "tool":
-        tool_schemas.append(format.create_tool_schema())
-
-    if not tool_schemas:
-        return None
-    definitions: list[dict[str, str | int | bool | dict[str, str | int | bool]]] = []
-    for tool in tool_schemas:
-        definitions.append(
-            {
-                "name": tool.name,
-                "description": tool.description,
-                "strict": tool.strict,
-                "parameters": tool.parameters.model_dump(by_alias=True, mode="json"),
-            }
-        )
-    return json_dumps(definitions)
-
-
-def _build_request_attributes(
-    *,
-    operation: str,
-    provider: ProviderId,
-    model_id: ModelId,
-    messages: Sequence[Message],
-    tools: ToolsParam,
-    format: FormatParam,
-    params: ParamsDict,
-) -> dict[str, AttributeValue]:
-    """Build GenAI request attributes for a span."""
-    attrs: dict[str, AttributeValue] = {
-        GenAIAttributes.GEN_AI_OPERATION_NAME: operation,
-        GenAIAttributes.GEN_AI_PROVIDER_NAME: provider,
-        GenAIAttributes.GEN_AI_REQUEST_MODEL: model_id,
-        GenAIAttributes.GEN_AI_OUTPUT_TYPE: _infer_output_type(format),
-    }
-    _apply_param_attributes(attrs, params)
-
-    _assign_request_message_attributes(
-        attrs.__setitem__,
-        messages=messages,
-    )
-
-    tool_payload = _serialize_tool_definitions(tools, format=format)
-    if tool_payload:
-        # The incubating semconv module does not yet expose a constant for this key.
-        attrs["gen_ai.tool.definitions"] = tool_payload
-
-    return attrs
-
-
-def _extract_response_id(
-    raw: dict[str, str | int] | str | Identifiable | None,
-) -> str | None:
-    """Extract response ID from raw response data."""
-    if isinstance(raw, dict):
-        for key in ("id", "response_id", "responseId"):
-            value = raw.get(key)
-            if isinstance(value, str):
-                return value
-    elif isinstance(raw, Identifiable):
-        return raw.id
-    return None
-
-
-def _attach_response(
-    span: Span,
-    response: RootResponse[ToolkitT, FormattableT | None],
-    *,
-    request_messages: Sequence[Message],
-) -> None:
-    """Attach response attributes to a GenAI span."""
-    span.set_attribute(GenAIAttributes.GEN_AI_RESPONSE_MODEL, response.model_id)
-    span.set_attribute(
-        GenAIAttributes.GEN_AI_RESPONSE_FINISH_REASONS,
-        [map_finish_reason(response.finish_reason)],
-    )
-    response_id = _extract_response_id(getattr(response, "raw", None))
-    if response_id:
-        span.set_attribute(GenAIAttributes.GEN_AI_RESPONSE_ID, response_id)
-
-    snapshot = snapshot_from_root_response(
-        response,
-        request_messages=request_messages,
-    )
-    _set_json_attribute(
-        span.set_attribute,
-        key=GenAIAttributes.GEN_AI_SYSTEM_INSTRUCTIONS,
-        payload=snapshot.system_instructions,
-    )
-    _set_json_attribute(
-        span.set_attribute,
-        key=GenAIAttributes.GEN_AI_INPUT_MESSAGES,
-        payload=snapshot.inputs,
-    )
-    _set_json_attribute(
-        span.set_attribute,
-        key=GenAIAttributes.GEN_AI_OUTPUT_MESSAGES,
-        payload=snapshot.outputs,
-    )
-    # TODO: Emit gen_ai.usage metrics once Response exposes provider-agnostic usage fields.
-
-
-_ORIGINAL_MODEL_CALL = Model.call
-_MODEL_CALL_WRAPPED = False
-_ORIGINAL_MODEL_CALL_ASYNC = Model.call_async
-_MODEL_CALL_ASYNC_WRAPPED = False
-_ORIGINAL_MODEL_CONTEXT_CALL = Model.context_call
-_MODEL_CONTEXT_CALL_WRAPPED = False
-_ORIGINAL_MODEL_CONTEXT_CALL_ASYNC = Model.context_call_async
-_MODEL_CONTEXT_CALL_ASYNC_WRAPPED = False
-_ORIGINAL_MODEL_STREAM = Model.stream
-_MODEL_STREAM_WRAPPED = False
-_ORIGINAL_MODEL_CONTEXT_STREAM = Model.context_stream
-_MODEL_CONTEXT_STREAM_WRAPPED = False
-_ORIGINAL_MODEL_CONTEXT_STREAM_ASYNC = Model.context_stream_async
-_MODEL_CONTEXT_STREAM_ASYNC_WRAPPED = False
-
-
-def _is_supported_param_value(value: object) -> TypeIs[ParamsValue]:
-    """Returns True if the value can be exported as an OTEL attribute."""
-    if isinstance(value, str | int | float | bool) or value is None:
-        return True
-    if isinstance(value, Sequence) and not isinstance(value, str | bytes):
-        return all(isinstance(item, str) for item in value)
-    return False
-
-
-def _normalize_dropped_value(value: object) -> Jsonable:
-    """Returns a JSON-safe representation for unsupported param values."""
-    if isinstance(value, str | int | float | bool) or value is None:
-        return value
-    if isinstance(value, Mapping):
-        normalized: dict[str, Jsonable] = {}
-        for key, item in value.items():
-            normalized[str(key)] = _normalize_dropped_value(item)
-        return normalized
-    if isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray):
-        return [_normalize_dropped_value(item) for item in value]
-    try:
-        return str(value)
-    except Exception:  # pragma: no cover
-        return f"<{type(value).__name__}>"
-
-
-def _params_as_mapping(params: Params) -> tuple[ParamsDict, dict[str, Jsonable]]:
-    """Returns supported params and a mapping of dropped params."""
-    filtered: dict[str, ParamsValue] = {}
-    dropped: dict[str, Jsonable] = {}
-    for key, value in params.items():
-        if _is_supported_param_value(value):
-            filtered[key] = value
-        else:
-            dropped[key] = _normalize_dropped_value(value)
-    return filtered, dropped
-
-
-def _record_dropped_params(
-    span: Span,
-    dropped_params: Mapping[str, Jsonable],
-) -> None:
-    """Emit an event with JSON-encoded params that cannot become attributes.
-
-    See https://opentelemetry.io/docs/specs/otel/common/ for the attribute type limits,
-    https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-events/ for the GenAI
-    guidance on recording richer payloads via events, and
-    https://opentelemetry.io/blog/2025/complex-attribute-types/ for the recommendation
-    to serialize unsupported complex types instead of dropping them outright.
-    """
-    if not dropped_params:
-        return None
-    payload = json_dumps(dropped_params)
-    span.add_event(
-        "gen_ai.request.params.untracked",
-        attributes={
-            "gen_ai.untracked_params.count": len(dropped_params),
-            "gen_ai.untracked_params.keys": list(dropped_params.keys()),
-            "gen_ai.untracked_params.json": payload,
-        },
-    )
-    return None
-
-
-@contextmanager
-def _start_model_span(
-    model: Model,
-    *,
-    messages: Sequence[Message],
-    tools: ToolsParam,
-    format: FormatParam,
-    activate: bool = True,
-) -> Iterator[SpanContext]:
-    """Context manager that yields a SpanContext for a model call."""
-    params, dropped_params = _params_as_mapping(model.params)
-    tracer = get_tracer()
-
-    if tracer is None or otel_trace is None:
-        yield SpanContext(None, dropped_params)
-        return
-
-    operation = GenAIAttributes.GenAiOperationNameValues.CHAT.value
-    attrs = _build_request_attributes(
-        operation=operation,
-        provider=model.provider_id,
-        model_id=model.model_id,
-        messages=messages,
-        tools=tools,
-        format=format,
-        params=params,
-    )
-    span_name = f"{operation} {model.model_id}"
-
-    if activate:
-        with tracer.start_as_current_span(
-            name=span_name,
-            kind=SpanKind.CLIENT,
-        ) as active_span:
-            for key, value in attrs.items():
-                active_span.set_attribute(key, value)
-            try:
-                yield SpanContext(active_span, dropped_params)
-            except Exception as exc:
-                _record_exception(active_span, exc)
-                raise
-        return
-
-    span = tracer.start_span(
-        name=span_name,
-        kind=SpanKind.CLIENT,
-    )
-    for key, value in attrs.items():
-        span.set_attribute(key, value)
-    try:
-        yield SpanContext(span, dropped_params)
-    except Exception as exc:
-        _record_exception(span, exc)
-        raise
-    finally:
-        span.end()
-
-
-@overload
-def _instrumented_model_call(
-    self: Model,
-    *,
-    messages: Sequence[Message],
-    tools: Sequence[Tool] | Toolkit | None = None,
-    format: None = None,
-) -> Response: ...
-
-
-@overload
-def _instrumented_model_call(
-    self: Model,
-    *,
-    messages: Sequence[Message],
-    tools: Sequence[Tool] | Toolkit | None = None,
-    format: type[FormattableT] | Format[FormattableT],
-) -> Response[FormattableT]: ...
-
-
-@overload
-def _instrumented_model_call(
-    self: Model,
-    *,
-    messages: Sequence[Message],
-    tools: Sequence[Tool] | Toolkit | None = None,
-    format: type[FormattableT] | Format[FormattableT] | None = None,
-) -> Response | Response[FormattableT]: ...
-
-
-@wraps(_ORIGINAL_MODEL_CALL)
-def _instrumented_model_call(
-    self: Model,
-    *,
-    messages: Sequence[Message],
-    tools: Sequence[Tool] | Toolkit | None = None,
-    format: FormatParam = None,
-) -> Response | Response[FormattableT]:
-    """Returns a GenAI-instrumented result of `Model.call`."""
-    with _start_model_span(
-        self,
-        messages=messages,
-        tools=tools,
-        format=format,
-    ) as span_ctx:
-        response = _ORIGINAL_MODEL_CALL(
-            self,
-            messages=messages,
-            tools=tools,
-            format=format,
-        )
-        if span_ctx.span is not None:
-            _attach_response(
-                span_ctx.span,
-                response,
-                request_messages=messages,
-            )
-            _record_dropped_params(span_ctx.span, span_ctx.dropped_params)
-        return response
-
-
-def _wrap_model_call() -> None:
-    """Returns None. Replaces `Model.call` with the instrumented wrapper."""
-    global _MODEL_CALL_WRAPPED
-    if _MODEL_CALL_WRAPPED:
-        return
-    Model.call = _instrumented_model_call
-    _MODEL_CALL_WRAPPED = True
-
-
-def _unwrap_model_call() -> None:
-    """Returns None. Restores the original `Model.call` implementation."""
-    global _MODEL_CALL_WRAPPED
-    if not _MODEL_CALL_WRAPPED:
-        return
-    Model.call = _ORIGINAL_MODEL_CALL
-    _MODEL_CALL_WRAPPED = False
-
-
-@overload
-async def _instrumented_model_call_async(
-    self: Model,
-    *,
-    messages: Sequence[Message],
-    tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
-    format: None = None,
-) -> AsyncResponse: ...
-
-
-@overload
-async def _instrumented_model_call_async(
-    self: Model,
-    *,
-    messages: Sequence[Message],
-    tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
-    format: type[FormattableT] | Format[FormattableT],
-) -> AsyncResponse[FormattableT]: ...
-
-
-@overload
-async def _instrumented_model_call_async(
-    self: Model,
-    *,
-    messages: Sequence[Message],
-    tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
-    format: type[FormattableT] | Format[FormattableT] | None = None,
-) -> AsyncResponse | AsyncResponse[FormattableT]: ...
-
-
-@wraps(_ORIGINAL_MODEL_CALL_ASYNC)
-async def _instrumented_model_call_async(
-    self: Model,
-    *,
-    messages: Sequence[Message],
-    tools: Sequence[AsyncTool] | AsyncToolkit | None = None,
-    format: FormatParam = None,
-) -> AsyncResponse | AsyncResponse[FormattableT]:
-    """Returns a GenAI-instrumented result of `Model.call_async`."""
-    with _start_model_span(
-        self,
-        messages=messages,
-        tools=tools,
-        format=format,
-        activate=True,
-    ) as span_ctx:
-        response = await _ORIGINAL_MODEL_CALL_ASYNC(
-            self,
-            messages=messages,
-            tools=tools,
-            format=format,
-        )
-        if span_ctx.span is not None:
-            _attach_response(
-                span_ctx.span,
-                response,
-                request_messages=messages,
-            )
-            _record_dropped_params(span_ctx.span, span_ctx.dropped_params)
-        return response
-
-
-def _wrap_model_call_async() -> None:
-    """Returns None. Replaces `Model.call_async` with the instrumented wrapper."""
-    global _MODEL_CALL_ASYNC_WRAPPED
-    if _MODEL_CALL_ASYNC_WRAPPED:
-        return
-    Model.call_async = _instrumented_model_call_async
-    _MODEL_CALL_ASYNC_WRAPPED = True
-
-
-def _unwrap_model_call_async() -> None:
-    """Returns None. Restores the original `Model.call_async` implementation."""
-    global _MODEL_CALL_ASYNC_WRAPPED
-    if not _MODEL_CALL_ASYNC_WRAPPED:
-        return
-    Model.call_async = _ORIGINAL_MODEL_CALL_ASYNC
-    _MODEL_CALL_ASYNC_WRAPPED = False
-
-
-@overload
-def _instrumented_model_context_call(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[Tool | ContextTool[DepsT]] | ContextToolkit[DepsT] | None = None,
-    format: None = None,
-) -> ContextResponse[DepsT, None]: ...
-
-
-@overload
-def _instrumented_model_context_call(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[Tool | ContextTool[DepsT]] | ContextToolkit[DepsT] | None = None,
-    format: type[FormattableT] | Format[FormattableT],
-) -> ContextResponse[DepsT, FormattableT]: ...
-
-
-@overload
-def _instrumented_model_context_call(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[Tool | ContextTool[DepsT]] | ContextToolkit[DepsT] | None = None,
-    format: type[FormattableT] | Format[FormattableT] | None = None,
-) -> ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]: ...
-
-
-@wraps(_ORIGINAL_MODEL_CONTEXT_CALL)
-def _instrumented_model_context_call(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[Tool | ContextTool[DepsT]] | ContextToolkit[DepsT] | None = None,
-    format: FormatParam = None,
-) -> ContextResponse[DepsT, None] | ContextResponse[DepsT, FormattableT]:
-    """Returns a GenAI-instrumented result of `Model.context_call`."""
-    with _start_model_span(
-        self,
-        messages=messages,
-        tools=tools,
-        format=format,
-        activate=True,
-    ) as span_ctx:
-        response = _ORIGINAL_MODEL_CONTEXT_CALL(
-            self,
-            ctx=ctx,
-            messages=messages,
-            tools=tools,
-            format=format,
-        )
-        if span_ctx.span is not None:
-            _attach_response(
-                span_ctx.span,
-                response,
-                request_messages=messages,
-            )
-            _record_dropped_params(span_ctx.span, span_ctx.dropped_params)
-        return response
-
-
-def _wrap_model_context_call() -> None:
-    """Returns None. Replaces `Model.context_call` with the instrumented wrapper."""
-    global _MODEL_CONTEXT_CALL_WRAPPED
-    if _MODEL_CONTEXT_CALL_WRAPPED:
-        return
-    Model.context_call = _instrumented_model_context_call
-    _MODEL_CONTEXT_CALL_WRAPPED = True
-
-
-def _unwrap_model_context_call() -> None:
-    """Returns None. Restores the original `Model.context_call` implementation."""
-    global _MODEL_CONTEXT_CALL_WRAPPED
-    if not _MODEL_CONTEXT_CALL_WRAPPED:
-        return
-    Model.context_call = _ORIGINAL_MODEL_CONTEXT_CALL
-    _MODEL_CONTEXT_CALL_WRAPPED = False
-
-
-@overload
-async def _instrumented_model_context_call_async(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
-    | AsyncContextToolkit[DepsT]
-    | None = None,
-    format: None = None,
-) -> AsyncContextResponse[DepsT, None]: ...
-
-
-@overload
-async def _instrumented_model_context_call_async(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
-    | AsyncContextToolkit[DepsT]
-    | None = None,
-    format: type[FormattableT] | Format[FormattableT],
-) -> AsyncContextResponse[DepsT, FormattableT]: ...
-
-
-@overload
-async def _instrumented_model_context_call_async(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
-    | AsyncContextToolkit[DepsT]
-    | None = None,
-    format: type[FormattableT] | Format[FormattableT] | None = None,
-) -> AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT]: ...
-
-
-@wraps(_ORIGINAL_MODEL_CONTEXT_CALL_ASYNC)
-async def _instrumented_model_context_call_async(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
-    | AsyncContextToolkit[DepsT]
-    | None = None,
-    format: FormatParam = None,
-) -> AsyncContextResponse[DepsT, None] | AsyncContextResponse[DepsT, FormattableT]:
-    """Returns a GenAI-instrumented result of `Model.context_call_async`."""
-    with _start_model_span(
-        self,
-        messages=messages,
-        tools=tools,
-        format=format,
-        activate=True,
-    ) as span_ctx:
-        response = await _ORIGINAL_MODEL_CONTEXT_CALL_ASYNC(
-            self,
-            ctx=ctx,
-            messages=messages,
-            tools=tools,
-            format=format,
-        )
-        if span_ctx.span is not None:
-            _attach_response(
-                span_ctx.span,
-                response,
-                request_messages=messages,
-            )
-            _record_dropped_params(span_ctx.span, span_ctx.dropped_params)
-        return response
-
-
-def _wrap_model_context_call_async() -> None:
-    """Returns None. Replaces `Model.context_call_async` with the instrumented wrapper."""
-    global _MODEL_CONTEXT_CALL_ASYNC_WRAPPED
-    if _MODEL_CONTEXT_CALL_ASYNC_WRAPPED:
-        return
-    Model.context_call_async = _instrumented_model_context_call_async
-    _MODEL_CONTEXT_CALL_ASYNC_WRAPPED = True
-
-
-def _unwrap_model_context_call_async() -> None:
-    """Returns None. Restores the original `Model.context_call_async` implementation."""
-    global _MODEL_CONTEXT_CALL_ASYNC_WRAPPED
-    if not _MODEL_CONTEXT_CALL_ASYNC_WRAPPED:
-        return
-    Model.context_call_async = _ORIGINAL_MODEL_CONTEXT_CALL_ASYNC
-    _MODEL_CONTEXT_CALL_ASYNC_WRAPPED = False
-
-
-@overload
-def _instrumented_model_stream(
-    self: Model,
-    *,
-    messages: Sequence[Message],
-    tools: Sequence[Tool] | Toolkit | None = None,
-    format: None = None,
-) -> StreamResponse: ...
-
-
-@overload
-def _instrumented_model_stream(
-    self: Model,
-    *,
-    messages: Sequence[Message],
-    tools: Sequence[Tool] | Toolkit | None = None,
-    format: type[FormattableT] | Format[FormattableT],
-) -> StreamResponse[FormattableT]: ...
-
-
-@overload
-def _instrumented_model_stream(
-    self: Model,
-    *,
-    messages: Sequence[Message],
-    tools: Sequence[Tool] | Toolkit | None = None,
-    format: type[FormattableT] | Format[FormattableT] | None = None,
-) -> StreamResponse | StreamResponse[FormattableT]: ...
-
-
-@wraps(_ORIGINAL_MODEL_STREAM)
-def _instrumented_model_stream(
-    self: Model,
-    *,
-    messages: Sequence[Message],
-    tools: Sequence[Tool] | Toolkit | None = None,
-    format: FormatParam = None,
-) -> StreamResponse | StreamResponse[FormattableT]:
-    """Returns a GenAI-instrumented result of `Model.stream`."""
-    span_cm = _start_model_span(
-        self,
-        messages=messages,
-        tools=tools,
-        format=format,
-        activate=False,
-    )
-    span_ctx = span_cm.__enter__()
-    if span_ctx.span is None:
-        response = _ORIGINAL_MODEL_STREAM(
-            self,
-            messages=messages,
-            tools=tools,
-            format=format,
-        )
-        span_cm.__exit__(None, None, None)
-        return response
-
-    try:
-        with otel_trace.use_span(span_ctx.span, end_on_exit=False):
-            response = _ORIGINAL_MODEL_STREAM(
-                self,
-                messages=messages,
-                tools=tools,
-                format=format,
-            )
-    except Exception as exc:
-        span_cm.__exit__(type(exc), exc, exc.__traceback__)
-        raise
-
-    _record_dropped_params(span_ctx.span, span_ctx.dropped_params)
-
-    try:
-        _attach_stream_span_handlers(
-            response=response,
-            span_cm=span_cm,
-            span=span_ctx.span,
-            request_messages=messages,
-        )
-    except Exception as exc:  # pragma: no cover
-        span_cm.__exit__(type(exc), exc, exc.__traceback__)
-        raise
-
-    return response
-
-
-def _attach_stream_span_handlers(
-    *,
-    response: ContextStreamResponse[DepsT, FormattableT | None]
-    | StreamResponse[FormattableT | None],
-    span_cm: AbstractContextManager[SpanContext],
-    span: Span,
-    request_messages: Sequence[Message],
-) -> None:
-    """Returns None. Closes the span when streaming completes."""
-    chunk_iterator: Iterator[StreamResponseChunk] = response._chunk_iterator
-
-    response_ref = weakref.ref(response)
-    closed = False
-
-    def _close_span(
-        exc_type: type[BaseException] | None,
-        exc: BaseException | None,
-        tb: TracebackType | None,
-    ) -> None:
-        nonlocal closed
-        if closed:
-            return
-        closed = True
-        response_obj = response_ref()
-        if response_obj is not None:
-            _attach_response(
-                span,
-                response_obj,
-                request_messages=request_messages,
-            )
-        span_cm.__exit__(exc_type, exc, tb)
-
-    def _wrapped_iterator() -> Iterator[StreamResponseChunk]:
-        with otel_trace.use_span(span, end_on_exit=False):
-            try:
-                yield from chunk_iterator
-            except Exception as exc:  # noqa: BLE001
-                _close_span(type(exc), exc, exc.__traceback__)
-                raise
-            else:
-                _close_span(None, None, None)
-            finally:
-                _close_span(None, None, None)
-
-    response._chunk_iterator = _wrapped_iterator()
-
-
-def _wrap_model_stream() -> None:
-    """Returns None. Replaces `Model.stream` with the instrumented wrapper."""
-    global _MODEL_STREAM_WRAPPED
-    if _MODEL_STREAM_WRAPPED:
-        return
-    Model.stream = _instrumented_model_stream
-    _MODEL_STREAM_WRAPPED = True
-
-
-def _unwrap_model_stream() -> None:
-    """Returns None. Restores the original `Model.stream` implementation."""
-    global _MODEL_STREAM_WRAPPED
-    if not _MODEL_STREAM_WRAPPED:
-        return
-    Model.stream = _ORIGINAL_MODEL_STREAM
-    _MODEL_STREAM_WRAPPED = False
-
-
-@overload
-def _instrumented_model_context_stream(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[Tool | ContextTool[DepsT]] | ContextToolkit[DepsT] | None = None,
-    format: None = None,
-) -> ContextStreamResponse[DepsT, None]: ...
-
-
-@overload
-def _instrumented_model_context_stream(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[Tool | ContextTool[DepsT]] | ContextToolkit[DepsT] | None = None,
-    format: type[FormattableT] | Format[FormattableT],
-) -> ContextStreamResponse[DepsT, FormattableT]: ...
-
-
-@overload
-def _instrumented_model_context_stream(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[Tool | ContextTool[DepsT]] | ContextToolkit[DepsT] | None = None,
-    format: type[FormattableT] | Format[FormattableT] | None = None,
-) -> (
-    ContextStreamResponse[DepsT, None] | ContextStreamResponse[DepsT, FormattableT]
-): ...
-
-
-@wraps(_ORIGINAL_MODEL_CONTEXT_STREAM)
-def _instrumented_model_context_stream(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[Tool | ContextTool[DepsT]] | ContextToolkit[DepsT] | None = None,
-    format: FormatParam = None,
-) -> ContextStreamResponse[DepsT, None] | ContextStreamResponse[DepsT, FormattableT]:
-    """Returns a GenAI-instrumented result of `Model.context_stream`."""
-    span_cm = _start_model_span(
-        self,
-        messages=messages,
-        tools=tools,
-        format=format,
-        activate=False,
-    )
-    span_ctx = span_cm.__enter__()
-    if span_ctx.span is None:
-        response = _ORIGINAL_MODEL_CONTEXT_STREAM(
-            self,
-            ctx=ctx,
-            messages=messages,
-            tools=tools,
-            format=format,
-        )
-        span_cm.__exit__(None, None, None)
-        return response
-
-    try:
-        with otel_trace.use_span(span_ctx.span, end_on_exit=False):
-            response = _ORIGINAL_MODEL_CONTEXT_STREAM(
-                self,
-                ctx=ctx,
-                messages=messages,
-                tools=tools,
-                format=format,
-            )
-    except Exception as exc:
-        span_cm.__exit__(type(exc), exc, exc.__traceback__)
-        raise
-
-    _record_dropped_params(span_ctx.span, span_ctx.dropped_params)
-
-    try:
-        _attach_stream_span_handlers(
-            response=response,
-            span_cm=span_cm,
-            span=span_ctx.span,
-            request_messages=messages,
-        )
-    except Exception as exc:  # pragma: no cover
-        span_cm.__exit__(type(exc), exc, exc.__traceback__)
-        raise
-
-    return response
-
-
-def _wrap_model_context_stream() -> None:
-    """Returns None. Replaces `Model.context_stream` with the instrumented wrapper."""
-    global _MODEL_CONTEXT_STREAM_WRAPPED
-    if _MODEL_CONTEXT_STREAM_WRAPPED:
-        return
-    Model.context_stream = _instrumented_model_context_stream
-    _MODEL_CONTEXT_STREAM_WRAPPED = True
-
-
-def _unwrap_model_context_stream() -> None:
-    """Returns None. Restores the original `Model.context_stream` implementation."""
-    global _MODEL_CONTEXT_STREAM_WRAPPED
-    if not _MODEL_CONTEXT_STREAM_WRAPPED:
-        return
-    Model.context_stream = _ORIGINAL_MODEL_CONTEXT_STREAM
-    _MODEL_CONTEXT_STREAM_WRAPPED = False
-
-
-def _attach_async_stream_span_handlers(
-    *,
-    response: AsyncContextStreamResponse[DepsT, FormattableT | None],
-    span_cm: AbstractContextManager[SpanContext],
-    span: Span,
-    request_messages: Sequence[Message],
-) -> None:
-    """Returns None. Closes the span when async streaming completes."""
-    chunk_iterator: AsyncIterator[StreamResponseChunk] = response._chunk_iterator
-
-    response_ref = weakref.ref(response)
-    closed = False
-
-    def _close_span(
-        exc_type: type[BaseException] | None,
-        exc: BaseException | None,
-        tb: TracebackType | None,
-    ) -> None:
-        nonlocal closed
-        if closed:
-            return
-        closed = True
-        response_obj = response_ref()
-        if response_obj is not None:
-            _attach_response(
-                span,
-                response_obj,
-                request_messages=request_messages,
-            )
-        span_cm.__exit__(exc_type, exc, tb)
-
-    async def _wrapped_iterator() -> AsyncIterator[StreamResponseChunk]:
-        try:
-            async for chunk in chunk_iterator:
-                yield chunk
-        except Exception as exc:  # noqa: BLE001
-            _close_span(type(exc), exc, exc.__traceback__)
-            raise
-        else:
-            _close_span(None, None, None)
-        finally:
-            _close_span(None, None, None)
-
-    response._chunk_iterator = _wrapped_iterator()
-
-
-@overload
-async def _instrumented_model_context_stream_async(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
-    | AsyncContextToolkit[DepsT]
-    | None = None,
-    format: None = None,
-) -> AsyncContextStreamResponse[DepsT, None]: ...
-
-
-@overload
-async def _instrumented_model_context_stream_async(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
-    | AsyncContextToolkit[DepsT]
-    | None = None,
-    format: type[FormattableT] | Format[FormattableT],
-) -> AsyncContextStreamResponse[DepsT, FormattableT]: ...
-
-
-@overload
-async def _instrumented_model_context_stream_async(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
-    | AsyncContextToolkit[DepsT]
-    | None = None,
-    format: type[FormattableT] | Format[FormattableT] | None = None,
-) -> (
-    AsyncContextStreamResponse[DepsT, None]
-    | AsyncContextStreamResponse[DepsT, FormattableT]
-): ...
-
-
-@wraps(_ORIGINAL_MODEL_CONTEXT_STREAM_ASYNC)
-async def _instrumented_model_context_stream_async(
-    self: Model,
-    *,
-    ctx: Context[DepsT],
-    messages: Sequence[Message],
-    tools: Sequence[AsyncTool | AsyncContextTool[DepsT]]
-    | AsyncContextToolkit[DepsT]
-    | None = None,
-    format: FormatParam = None,
-) -> (
-    AsyncContextStreamResponse[DepsT, None]
-    | AsyncContextStreamResponse[DepsT, FormattableT]
-):
-    """Returns a GenAI-instrumented result of `Model.context_stream_async`."""
-    span_cm = _start_model_span(
-        self,
-        messages=messages,
-        tools=tools,
-        format=format,
-        activate=False,
-    )
-    span_ctx = span_cm.__enter__()
-    if span_ctx.span is None:
-        response = await _ORIGINAL_MODEL_CONTEXT_STREAM_ASYNC(
-            self,
-            ctx=ctx,
-            messages=messages,
-            tools=tools,
-            format=format,
-        )
-        span_cm.__exit__(None, None, None)
-        return response
-
-    try:
-        with otel_trace.use_span(span_ctx.span, end_on_exit=False):
-            response = await _ORIGINAL_MODEL_CONTEXT_STREAM_ASYNC(
-                self,
-                ctx=ctx,
-                messages=messages,
-                tools=tools,
-                format=format,
-            )
-    except Exception as exc:
-        span_cm.__exit__(type(exc), exc, exc.__traceback__)
-        raise
-
-    _record_dropped_params(span_ctx.span, span_ctx.dropped_params)
-
-    try:
-        _attach_async_stream_span_handlers(
-            response=response,
-            span_cm=span_cm,
-            span=span_ctx.span,
-            request_messages=messages,
-        )
-    except Exception as exc:  # pragma: no cover
-        span_cm.__exit__(type(exc), exc, exc.__traceback__)
-        raise
-
-    return response
-
-
-def _wrap_model_context_stream_async() -> None:
-    """Returns None. Replaces `Model.context_stream_async` with the instrumented wrapper."""
-    global _MODEL_CONTEXT_STREAM_ASYNC_WRAPPED
-    if _MODEL_CONTEXT_STREAM_ASYNC_WRAPPED:
-        return
-    Model.context_stream_async = _instrumented_model_context_stream_async
-    _MODEL_CONTEXT_STREAM_ASYNC_WRAPPED = True
-
-
-def _unwrap_model_context_stream_async() -> None:
-    """Returns None. Restores the original `Model.context_stream_async` implementation."""
-    global _MODEL_CONTEXT_STREAM_ASYNC_WRAPPED
-    if not _MODEL_CONTEXT_STREAM_ASYNC_WRAPPED:
-        return
-    Model.context_stream_async = _ORIGINAL_MODEL_CONTEXT_STREAM_ASYNC
-    _MODEL_CONTEXT_STREAM_ASYNC_WRAPPED = False
 
 
 def instrument_llm() -> None:
@@ -1267,21 +97,65 @@ def instrument_llm() -> None:
             "You must call `configure()` before calling `instrument_llm()`."
         )
 
-    _wrap_model_call()
-    _wrap_model_call_async()
-    _wrap_model_context_call()
-    _wrap_model_context_call_async()
-    _wrap_model_stream()
-    _wrap_model_context_stream()
-    _wrap_model_context_stream_async()
+    # Model call methods
+    wrap_model_call()
+    wrap_model_call_async()
+    wrap_model_context_call()
+    wrap_model_context_call_async()
+    wrap_model_stream()
+    wrap_model_stream_async()
+    wrap_model_context_stream()
+    wrap_model_context_stream_async()
+
+    # Model resume methods
+    wrap_model_resume()
+    wrap_model_resume_async()
+    wrap_model_context_resume()
+    wrap_model_context_resume_async()
+    wrap_model_resume_stream()
+    wrap_model_resume_stream_async()
+    wrap_model_context_resume_stream()
+    wrap_model_context_resume_stream_async()
+
+    # Response resume methods (Mirascope-specific spans)
+    wrap_response_resume()
+    wrap_async_response_resume()
+    wrap_context_response_resume()
+    wrap_async_context_response_resume()
+    wrap_stream_response_resume()
+    wrap_async_stream_response_resume()
+    wrap_context_stream_response_resume()
+    wrap_async_context_stream_response_resume()
 
 
 def uninstrument_llm() -> None:
     """Disable previously configured instrumentation."""
-    _unwrap_model_call()
-    _unwrap_model_call_async()
-    _unwrap_model_context_call()
-    _unwrap_model_context_call_async()
-    _unwrap_model_stream()
-    _unwrap_model_context_stream()
-    _unwrap_model_context_stream_async()
+    # Model call methods
+    unwrap_model_call()
+    unwrap_model_call_async()
+    unwrap_model_context_call()
+    unwrap_model_context_call_async()
+    unwrap_model_stream()
+    unwrap_model_stream_async()
+    unwrap_model_context_stream()
+    unwrap_model_context_stream_async()
+
+    # Model resume methods
+    unwrap_model_resume()
+    unwrap_model_resume_async()
+    unwrap_model_context_resume()
+    unwrap_model_context_resume_async()
+    unwrap_model_resume_stream()
+    unwrap_model_resume_stream_async()
+    unwrap_model_context_resume_stream()
+    unwrap_model_context_resume_stream_async()
+
+    # Response resume methods (Mirascope-specific spans)
+    unwrap_response_resume()
+    unwrap_async_response_resume()
+    unwrap_context_response_resume()
+    unwrap_async_context_response_resume()
+    unwrap_stream_response_resume()
+    unwrap_async_stream_response_resume()
+    unwrap_context_stream_response_resume()
+    unwrap_async_context_stream_response_resume()