sentry-sdk 3.0.0a1__py2.py3-none-any.whl → 3.0.0a3__py2.py3-none-any.whl

sentry_sdk/integrations/modules.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import sentry_sdk
 from sentry_sdk.integrations import Integration
 from sentry_sdk.scope import add_global_event_processor
@@ -14,11 +15,9 @@ class ModulesIntegration(Integration):
     identifier = "modules"
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         @add_global_event_processor
-        def processor(event, hint):
-            # type: (Event, Any) -> Event
+        def processor(event: Event, hint: Any) -> Event:
             if event.get("type") == "transaction":
                 return event
 

sentry_sdk/integrations/openai.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from functools import wraps
 
 import sentry_sdk
@@ -32,8 +33,11 @@ class OpenAIIntegration(Integration):
     identifier = "openai"
     origin = f"auto.ai.{identifier}"
 
-    def __init__(self, include_prompts=True, tiktoken_encoding_name=None):
-        # type: (OpenAIIntegration, bool, Optional[str]) -> None
+    def __init__(
+        self: OpenAIIntegration,
+        include_prompts: bool = True,
+        tiktoken_encoding_name: Optional[str] = None,
+    ) -> None:
         self.include_prompts = include_prompts
 
         self.tiktoken_encoding = None
@@ -43,8 +47,7 @@ class OpenAIIntegration(Integration):
             self.tiktoken_encoding = tiktoken.get_encoding(tiktoken_encoding_name)
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         Completions.create = _wrap_chat_completion_create(Completions.create)
         Embeddings.create = _wrap_embeddings_create(Embeddings.create)
 
@@ -53,15 +56,13 @@ class OpenAIIntegration(Integration):
         )
         AsyncEmbeddings.create = _wrap_async_embeddings_create(AsyncEmbeddings.create)
 
-    def count_tokens(self, s):
-        # type: (OpenAIIntegration, str) -> int
+    def count_tokens(self: OpenAIIntegration, s: str) -> int:
         if self.tiktoken_encoding is not None:
             return len(self.tiktoken_encoding.encode_ordinary(s))
         return 0
 
 
-def _capture_exception(exc):
-    # type: (Any) -> None
+def _capture_exception(exc: Any) -> None:
     event, hint = event_from_exception(
         exc,
         client_options=sentry_sdk.get_client().options,
@@ -70,52 +71,78 @@ def _capture_exception(exc):
     sentry_sdk.capture_event(event, hint=hint)
 
 
-def _calculate_chat_completion_usage(
-    messages, response, span, streaming_message_responses, count_tokens
-):
-    # type: (Iterable[ChatCompletionMessageParam], Any, Span, Optional[List[str]], Callable[..., Any]) -> None
-    completion_tokens = 0  # type: Optional[int]
-    prompt_tokens = 0  # type: Optional[int]
-    total_tokens = 0  # type: Optional[int]
+def _get_usage(usage: Any, names: List[str]) -> int:
+    for name in names:
+        if hasattr(usage, name) and isinstance(getattr(usage, name), int):
+            return getattr(usage, name)
+    return 0
+
+
+def _calculate_token_usage(
+    messages: Iterable[ChatCompletionMessageParam],
+    response: Any,
+    span: Span,
+    streaming_message_responses: Optional[List[str]],
+    count_tokens: Callable[..., Any],
+) -> None:
+    input_tokens: Optional[int] = 0
+    input_tokens_cached: Optional[int] = 0
+    output_tokens: Optional[int] = 0
+    output_tokens_reasoning: Optional[int] = 0
+    total_tokens: Optional[int] = 0
+
     if hasattr(response, "usage"):
-        if hasattr(response.usage, "completion_tokens") and isinstance(
-            response.usage.completion_tokens, int
-        ):
-            completion_tokens = response.usage.completion_tokens
-        if hasattr(response.usage, "prompt_tokens") and isinstance(
-            response.usage.prompt_tokens, int
-        ):
-            prompt_tokens = response.usage.prompt_tokens
-        if hasattr(response.usage, "total_tokens") and isinstance(
-            response.usage.total_tokens, int
-        ):
-            total_tokens = response.usage.total_tokens
+        input_tokens = _get_usage(response.usage, ["input_tokens", "prompt_tokens"])
+        if hasattr(response.usage, "input_tokens_details"):
+            input_tokens_cached = _get_usage(
+                response.usage.input_tokens_details, ["cached_tokens"]
+            )
 
-    if prompt_tokens == 0:
+        output_tokens = _get_usage(
+            response.usage, ["output_tokens", "completion_tokens"]
+        )
+        if hasattr(response.usage, "output_tokens_details"):
+            output_tokens_reasoning = _get_usage(
+                response.usage.output_tokens_details, ["reasoning_tokens"]
+            )
+
+        total_tokens = _get_usage(response.usage, ["total_tokens"])
+
+    # Manually count tokens
+    # TODO: when implementing responses API, check for responses API
+    if input_tokens == 0:
         for message in messages:
             if "content" in message:
-                prompt_tokens += count_tokens(message["content"])
+                input_tokens += count_tokens(message["content"])
 
-    if completion_tokens == 0:
+    # TODO: when implementing responses API, check for responses API
+    if output_tokens == 0:
         if streaming_message_responses is not None:
             for message in streaming_message_responses:
-                completion_tokens += count_tokens(message)
+                output_tokens += count_tokens(message)
         elif hasattr(response, "choices"):
             for choice in response.choices:
                 if hasattr(choice, "message"):
-                    completion_tokens += count_tokens(choice.message)
-
-    if prompt_tokens == 0:
-        prompt_tokens = None
-    if completion_tokens == 0:
-        completion_tokens = None
-    if total_tokens == 0:
-        total_tokens = None
-    record_token_usage(span, prompt_tokens, completion_tokens, total_tokens)
+                    output_tokens += count_tokens(choice.message)
+
+    # Do not set token data if it is 0
+    input_tokens = input_tokens or None
+    input_tokens_cached = input_tokens_cached or None
+    output_tokens = output_tokens or None
+    output_tokens_reasoning = output_tokens_reasoning or None
+    total_tokens = total_tokens or None
+
+    record_token_usage(
+        span,
+        input_tokens=input_tokens,
+        input_tokens_cached=input_tokens_cached,
+        output_tokens=output_tokens,
+        output_tokens_reasoning=output_tokens_reasoning,
+        total_tokens=total_tokens,
+    )
 
 
-def _new_chat_completion_common(f, *args, **kwargs):
-    # type: (Any, *Any, **Any) -> Any
+def _new_chat_completion_common(f: Any, *args: Any, **kwargs: Any) -> Any:
     integration = sentry_sdk.get_client().get_integration(OpenAIIntegration)
     if integration is None:
         return f(*args, **kwargs)
@@ -139,7 +166,6 @@ def _new_chat_completion_common(f, *args, **kwargs):
         op=consts.OP.OPENAI_CHAT_COMPLETIONS_CREATE,
         name="Chat Completion",
         origin=OpenAIIntegration.origin,
-        only_if_parent=True,
     )
     span.__enter__()
 
@@ -156,20 +182,17 @@ def _new_chat_completion_common(f, *args, **kwargs):
             if should_send_default_pii() and integration.include_prompts:
                 set_data_normalized(
                     span,
-                    "ai.responses",
+                    SPANDATA.AI_RESPONSES,
                     list(map(lambda x: x.message, res.choices)),
                 )
-            _calculate_chat_completion_usage(
-                messages, res, span, None, integration.count_tokens
-            )
+            _calculate_token_usage(messages, res, span, None, integration.count_tokens)
             span.__exit__(None, None, None)
         elif hasattr(res, "_iterator"):
             data_buf: list[list[str]] = []  # one for each choice
 
             old_iterator = res._iterator
 
-            def new_iterator():
-                # type: () -> Iterator[ChatCompletionChunk]
+            def new_iterator() -> Iterator[ChatCompletionChunk]:
                 with capture_internal_exceptions():
                     for x in old_iterator:
                         if hasattr(x, "choices"):
@@ -192,7 +215,7 @@ def _new_chat_completion_common(f, *args, **kwargs):
                             set_data_normalized(
                                 span, SPANDATA.AI_RESPONSES, all_responses
                             )
-                        _calculate_chat_completion_usage(
+                        _calculate_token_usage(
                             messages,
                             res,
                             span,
@@ -201,8 +224,7 @@ def _new_chat_completion_common(f, *args, **kwargs):
                         )
                 span.__exit__(None, None, None)
 
-            async def new_iterator_async():
-                # type: () -> AsyncIterator[ChatCompletionChunk]
+            async def new_iterator_async() -> AsyncIterator[ChatCompletionChunk]:
                 with capture_internal_exceptions():
                     async for x in old_iterator:
                         if hasattr(x, "choices"):
@@ -225,7 +247,7 @@ def _new_chat_completion_common(f, *args, **kwargs):
                             set_data_normalized(
                                 span, SPANDATA.AI_RESPONSES, all_responses
                             )
-                        _calculate_chat_completion_usage(
+                        _calculate_token_usage(
                             messages,
                             res,
                             span,
@@ -245,10 +267,8 @@ def _new_chat_completion_common(f, *args, **kwargs):
     return res
 
 
-def _wrap_chat_completion_create(f):
-    # type: (Callable[..., Any]) -> Callable[..., Any]
-    def _execute_sync(f, *args, **kwargs):
-        # type: (Any, *Any, **Any) -> Any
+def _wrap_chat_completion_create(f: Callable[..., Any]) -> Callable[..., Any]:
+    def _execute_sync(f: Any, *args: Any, **kwargs: Any) -> Any:
         gen = _new_chat_completion_common(f, *args, **kwargs)
 
         try:
@@ -268,8 +288,7 @@ def _wrap_chat_completion_create(f):
             return e.value
 
     @wraps(f)
-    def _sentry_patched_create_sync(*args, **kwargs):
-        # type: (*Any, **Any) -> Any
+    def _sentry_patched_create_sync(*args: Any, **kwargs: Any) -> Any:
         integration = sentry_sdk.get_client().get_integration(OpenAIIntegration)
         if integration is None or "messages" not in kwargs:
             # no "messages" means invalid call (in all versions of openai), let it return error
@@ -280,10 +299,8 @@ def _wrap_chat_completion_create(f):
     return _sentry_patched_create_sync
 
 
-def _wrap_async_chat_completion_create(f):
-    # type: (Callable[..., Any]) -> Callable[..., Any]
-    async def _execute_async(f, *args, **kwargs):
-        # type: (Any, *Any, **Any) -> Any
+def _wrap_async_chat_completion_create(f: Callable[..., Any]) -> Callable[..., Any]:
+    async def _execute_async(f: Any, *args: Any, **kwargs: Any) -> Any:
         gen = _new_chat_completion_common(f, *args, **kwargs)
 
         try:
@@ -303,8 +320,7 @@ def _wrap_async_chat_completion_create(f):
             return e.value
 
     @wraps(f)
-    async def _sentry_patched_create_async(*args, **kwargs):
-        # type: (*Any, **Any) -> Any
+    async def _sentry_patched_create_async(*args: Any, **kwargs: Any) -> Any:
         integration = sentry_sdk.get_client().get_integration(OpenAIIntegration)
         if integration is None or "messages" not in kwargs:
             # no "messages" means invalid call (in all versions of openai), let it return error
@@ -315,8 +331,7 @@ def _wrap_async_chat_completion_create(f):
     return _sentry_patched_create_async
 
 
-def _new_embeddings_create_common(f, *args, **kwargs):
-    # type: (Any, *Any, **Any) -> Any
+def _new_embeddings_create_common(f: Any, *args: Any, **kwargs: Any) -> Any:
     integration = sentry_sdk.get_client().get_integration(OpenAIIntegration)
     if integration is None:
         return f(*args, **kwargs)
@@ -325,48 +340,49 @@ def _new_embeddings_create_common(f, *args, **kwargs):
         op=consts.OP.OPENAI_EMBEDDINGS_CREATE,
         description="OpenAI Embedding Creation",
         origin=OpenAIIntegration.origin,
-        only_if_parent=True,
     ) as span:
         if "input" in kwargs and (
             should_send_default_pii() and integration.include_prompts
         ):
             if isinstance(kwargs["input"], str):
-                set_data_normalized(span, "ai.input_messages", [kwargs["input"]])
+                set_data_normalized(span, SPANDATA.AI_INPUT_MESSAGES, [kwargs["input"]])
             elif (
                 isinstance(kwargs["input"], list)
                 and len(kwargs["input"]) > 0
                 and isinstance(kwargs["input"][0], str)
             ):
-                set_data_normalized(span, "ai.input_messages", kwargs["input"])
+                set_data_normalized(span, SPANDATA.AI_INPUT_MESSAGES, kwargs["input"])
         if "model" in kwargs:
-            set_data_normalized(span, "ai.model_id", kwargs["model"])
+            set_data_normalized(span, SPANDATA.AI_MODEL_ID, kwargs["model"])
 
         response = yield f, args, kwargs
 
-        prompt_tokens = 0
+        input_tokens = 0
         total_tokens = 0
         if hasattr(response, "usage"):
             if hasattr(response.usage, "prompt_tokens") and isinstance(
                 response.usage.prompt_tokens, int
             ):
-                prompt_tokens = response.usage.prompt_tokens
+                input_tokens = response.usage.prompt_tokens
             if hasattr(response.usage, "total_tokens") and isinstance(
                 response.usage.total_tokens, int
             ):
                 total_tokens = response.usage.total_tokens
 
-        if prompt_tokens == 0:
-            prompt_tokens = integration.count_tokens(kwargs["input"] or "")
+        if input_tokens == 0:
+            input_tokens = integration.count_tokens(kwargs["input"] or "")
 
-        record_token_usage(span, prompt_tokens, None, total_tokens or prompt_tokens)
+        record_token_usage(
+            span,
+            input_tokens=input_tokens,
+            total_tokens=total_tokens or input_tokens,
+        )
 
         return response
 
 
-def _wrap_embeddings_create(f):
-    # type: (Any) -> Any
-    def _execute_sync(f, *args, **kwargs):
-        # type: (Any, *Any, **Any) -> Any
+def _wrap_embeddings_create(f: Any) -> Any:
+    def _execute_sync(f: Any, *args: Any, **kwargs: Any) -> Any:
         gen = _new_embeddings_create_common(f, *args, **kwargs)
 
         try:
@@ -386,8 +402,7 @@ def _wrap_embeddings_create(f):
             return e.value
 
     @wraps(f)
-    def _sentry_patched_create_sync(*args, **kwargs):
-        # type: (*Any, **Any) -> Any
+    def _sentry_patched_create_sync(*args: Any, **kwargs: Any) -> Any:
         integration = sentry_sdk.get_client().get_integration(OpenAIIntegration)
         if integration is None:
             return f(*args, **kwargs)
@@ -397,10 +412,8 @@ def _wrap_embeddings_create(f):
     return _sentry_patched_create_sync
 
 
-def _wrap_async_embeddings_create(f):
-    # type: (Any) -> Any
-    async def _execute_async(f, *args, **kwargs):
-        # type: (Any, *Any, **Any) -> Any
+def _wrap_async_embeddings_create(f: Any) -> Any:
+    async def _execute_async(f: Any, *args: Any, **kwargs: Any) -> Any:
         gen = _new_embeddings_create_common(f, *args, **kwargs)
 
         try:
@@ -420,8 +433,7 @@ def _wrap_async_embeddings_create(f):
             return e.value
 
     @wraps(f)
-    async def _sentry_patched_create_async(*args, **kwargs):
-        # type: (*Any, **Any) -> Any
+    async def _sentry_patched_create_async(*args: Any, **kwargs: Any) -> Any:
         integration = sentry_sdk.get_client().get_integration(OpenAIIntegration)
         if integration is None:
             return await f(*args, **kwargs)

sentry_sdk/integrations/openai_agents/__init__.py

@@ -0,0 +1,49 @@
+from sentry_sdk.integrations import DidNotEnable, Integration
+
+from .patches import (
+    _create_get_model_wrapper,
+    _create_get_all_tools_wrapper,
+    _create_run_wrapper,
+    _patch_agent_run,
+)
+
+try:
+    import agents
+
+except ImportError:
+    raise DidNotEnable("OpenAI Agents not installed")
+
+
+def _patch_runner() -> None:
+    # Create the root span for one full agent run (including eventual handoffs)
+    # Note agents.run.DEFAULT_AGENT_RUNNER.run_sync is a wrapper around
+    # agents.run.DEFAULT_AGENT_RUNNER.run. It does not need to be wrapped separately.
+    # TODO-anton: Also patch streaming runner: agents.Runner.run_streamed
+    agents.run.DEFAULT_AGENT_RUNNER.run = _create_run_wrapper(
+        agents.run.DEFAULT_AGENT_RUNNER.run
+    )
+
+    # Creating the actual spans for each agent run.
+    _patch_agent_run()
+
+
+def _patch_model() -> None:
+    agents.run.AgentRunner._get_model = classmethod(
+        _create_get_model_wrapper(agents.run.AgentRunner._get_model),
+    )
+
+
+def _patch_tools() -> None:
+    agents.run.AgentRunner._get_all_tools = classmethod(
+        _create_get_all_tools_wrapper(agents.run.AgentRunner._get_all_tools),
+    )
+
+
+class OpenAIAgentsIntegration(Integration):
+    identifier = "openai_agents"
+
+    @staticmethod
+    def setup_once() -> None:
+        _patch_tools()
+        _patch_model()
+        _patch_runner()

sentry_sdk/integrations/openai_agents/consts.py

@@ -0,0 +1 @@
+SPAN_ORIGIN = "auto.ai.openai_agents"

sentry_sdk/integrations/openai_agents/patches/__init__.py

@@ -0,0 +1,4 @@
+from .models import _create_get_model_wrapper  # noqa: F401
+from .tools import _create_get_all_tools_wrapper  # noqa: F401
+from .runner import _create_run_wrapper  # noqa: F401
+from .agent_run import _patch_agent_run  # noqa: F401

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+from functools import wraps
+
+from sentry_sdk.integrations import DidNotEnable
+
+from ..spans import invoke_agent_span, update_invoke_agent_span, handoff_span
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Optional
+
+
+try:
+    import agents
+except ImportError:
+    raise DidNotEnable("OpenAI Agents not installed")
+
+
+def _patch_agent_run() -> None:
+    """
+    Patches AgentRunner methods to create agent invocation spans.
+    This directly patches the execution flow to track when agents start and stop.
+    """
+
+    # Store original methods
+    original_run_single_turn = agents.run.AgentRunner._run_single_turn
+    original_execute_handoffs = agents._run_impl.RunImpl.execute_handoffs
+    original_execute_final_output = agents._run_impl.RunImpl.execute_final_output
+
+    def _start_invoke_agent_span(
+        context_wrapper: agents.RunContextWrapper, agent: agents.Agent
+    ) -> None:
+        """Start an agent invocation span"""
+        # Store the agent on the context wrapper so we can access it later
+        context_wrapper._sentry_current_agent = agent
+        invoke_agent_span(context_wrapper, agent)
+
+    def _end_invoke_agent_span(
+        context_wrapper: agents.RunContextWrapper,
+        agent: agents.Agent,
+        output: Optional[Any] = None,
+    ) -> None:
+        """End the agent invocation span"""
+        # Clear the stored agent
+        if hasattr(context_wrapper, "_sentry_current_agent"):
+            delattr(context_wrapper, "_sentry_current_agent")
+
+        update_invoke_agent_span(context_wrapper, agent, output)
+
+    def _has_active_agent_span(context_wrapper: agents.RunContextWrapper) -> bool:
+        """Check if there's an active agent span for this context"""
+        return getattr(context_wrapper, "_sentry_current_agent", None) is not None
+
+    def _get_current_agent(
+        context_wrapper: agents.RunContextWrapper,
+    ) -> Optional[agents.Agent]:
+        """Get the current agent from context wrapper"""
+        return getattr(context_wrapper, "_sentry_current_agent", None)
+
+    @wraps(
+        original_run_single_turn.__func__
+        if hasattr(original_run_single_turn, "__func__")
+        else original_run_single_turn
+    )
+    async def patched_run_single_turn(
+        cls: agents.Runner, *args: Any, **kwargs: Any
+    ) -> Any:
+        """Patched _run_single_turn that creates agent invocation spans"""
+        agent = kwargs.get("agent")
+        context_wrapper = kwargs.get("context_wrapper")
+        should_run_agent_start_hooks = kwargs.get("should_run_agent_start_hooks")
+
+        # Start agent span when agent starts (but only once per agent)
+        if should_run_agent_start_hooks and agent and context_wrapper:
+            # End any existing span for a different agent
+            if _has_active_agent_span(context_wrapper):
+                current_agent = _get_current_agent(context_wrapper)
+                if current_agent and current_agent != agent:
+                    _end_invoke_agent_span(context_wrapper, current_agent)
+
+            _start_invoke_agent_span(context_wrapper, agent)
+
+        # Call original method with all the correct parameters
+        try:
+            result = await original_run_single_turn(*args, **kwargs)
+        finally:
+            if agent and context_wrapper and _has_active_agent_span(context_wrapper):
+                _end_invoke_agent_span(context_wrapper, agent)
+
+        return result
+
+    @wraps(
+        original_execute_handoffs.__func__
+        if hasattr(original_execute_handoffs, "__func__")
+        else original_execute_handoffs
+    )
+    async def patched_execute_handoffs(
+        cls: agents.Runner, *args: Any, **kwargs: Any
+    ) -> Any:
+        """Patched execute_handoffs that creates handoff spans and ends agent span for handoffs"""
+        context_wrapper = kwargs.get("context_wrapper")
+        run_handoffs = kwargs.get("run_handoffs")
+        agent = kwargs.get("agent")
+
+        # Create Sentry handoff span for the first handoff (agents library only processes the first one)
+        if run_handoffs:
+            first_handoff = run_handoffs[0]
+            handoff_agent_name = first_handoff.handoff.agent_name
+            handoff_span(context_wrapper, agent, handoff_agent_name)
+
+        # Call original method with all parameters
+        try:
+            result = await original_execute_handoffs(*args, **kwargs)
+
+        finally:
+            # End span for current agent after handoff processing is complete
+            if agent and context_wrapper and _has_active_agent_span(context_wrapper):
+                _end_invoke_agent_span(context_wrapper, agent)
+
+        return result
+
+    @wraps(
+        original_execute_final_output.__func__
+        if hasattr(original_execute_final_output, "__func__")
+        else original_execute_final_output
+    )
+    async def patched_execute_final_output(
+        cls: agents.Runner, *args: Any, **kwargs: Any
+    ) -> Any:
+        """Patched execute_final_output that ends agent span for final outputs"""
+        agent = kwargs.get("agent")
+        context_wrapper = kwargs.get("context_wrapper")
+        final_output = kwargs.get("final_output")
+
+        # Call original method with all parameters
+        try:
+            result = await original_execute_final_output(*args, **kwargs)
+        finally:
+            # End span for current agent after final output processing is complete
+            if agent and context_wrapper and _has_active_agent_span(context_wrapper):
+                _end_invoke_agent_span(context_wrapper, agent, final_output)
+
+        return result
+
+    # Apply patches
+    agents.run.AgentRunner._run_single_turn = classmethod(patched_run_single_turn)
+    agents._run_impl.RunImpl.execute_handoffs = classmethod(patched_execute_handoffs)
+    agents._run_impl.RunImpl.execute_final_output = classmethod(
+        patched_execute_final_output
+    )

sentry_sdk/integrations/openai_agents/patches/models.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from functools import wraps
+
+from sentry_sdk.integrations import DidNotEnable
+
+from ..spans import ai_client_span, update_ai_client_span
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Callable
+
+
+try:
+    import agents
+except ImportError:
+    raise DidNotEnable("OpenAI Agents not installed")
+
+
+def _create_get_model_wrapper(
+    original_get_model: Callable[..., Any],
+) -> Callable[..., Any]:
+    """
+    Wraps the agents.Runner._get_model method to wrap the get_response method of the model to create a AI client span.
+    """
+
+    @wraps(
+        original_get_model.__func__
+        if hasattr(original_get_model, "__func__")
+        else original_get_model
+    )
+    def wrapped_get_model(
+        cls: agents.Runner, agent: agents.Agent, run_config: agents.RunConfig
+    ) -> agents.Model:
+        model = original_get_model(agent, run_config)
+        original_get_response = model.get_response
+
+        @wraps(original_get_response)
+        async def wrapped_get_response(*args: Any, **kwargs: Any) -> Any:
+            with ai_client_span(agent, kwargs) as span:
+                result = await original_get_response(*args, **kwargs)
+
+                update_ai_client_span(span, agent, kwargs, result)
+
+            return result
+
+        model.get_response = wrapped_get_response
+
+        return model
+
+    return wrapped_get_model