Add support for async calls in Anthropic and OpenAI integration

sentry_sdk/integrations/anthropic.py

@@ -15,14 +15,13 @@
 from typing import TYPE_CHECKING
 
 try:
-    from anthropic.resources import Messages
+    from anthropic.resources import Messages, AsyncMessages
 
     if TYPE_CHECKING:
         from anthropic.types import MessageStreamEvent
 except ImportError:
     raise DidNotEnable("Anthropic not installed")
 
-
 if TYPE_CHECKING:
     from typing import Any, Iterator
     from sentry_sdk.tracing import Span
@@ -48,6 +47,7 @@ def setup_once():
             raise DidNotEnable("anthropic 0.16 or newer required.")
 
         Messages.create = _wrap_message_create(Messages.create)
+        AsyncMessages.create = _wrap_async_message_create(AsyncMessages.create)
 
 
 def _capture_exception(exc):
@@ -75,105 +75,150 @@ def _calculate_token_usage(result, span):
     record_token_usage(span, input_tokens, output_tokens, total_tokens)
 
 
+def _sentry_patched_create_common(f, *args, **kwargs):
+    # type: (Any, *Any, **Any) -> Any
+    if "messages" not in kwargs:
+        return f(*args, **kwargs)
+
+    try:
+        iter(kwargs["messages"])
+    except TypeError:
+        return f(*args, **kwargs)
+
+    messages = list(kwargs["messages"])
+    model = kwargs.get("model")
+
+    span = sentry_sdk.start_span(
+        op=OP.ANTHROPIC_MESSAGES_CREATE,
+        description="Anthropic messages create",
+        origin=AnthropicIntegration.origin,
+    )
+    span.__enter__()
+
+    try:
+        result = yield f, args, kwargs
+    except Exception as exc:
+        _capture_exception(exc)
+        span.__exit__(None, None, None)
+        raise exc from None
+
+    integration = sentry_sdk.get_client().get_integration(AnthropicIntegration)
+
+    with capture_internal_exceptions():
+        span.set_data(SPANDATA.AI_MODEL_ID, model)
+        span.set_data(SPANDATA.AI_STREAMING, False)
+        if should_send_default_pii() and integration.include_prompts:
+            span.set_data(SPANDATA.AI_INPUT_MESSAGES, messages)
+        if hasattr(result, "content"):
+            if should_send_default_pii() and integration.include_prompts:
+                span.set_data(
+                    SPANDATA.AI_RESPONSES,
+                    list(
+                        map(
+                            lambda message: {
+                                "type": message.type,
+                                "text": message.text,
+                            },
+                            result.content,
+                        )
+                    ),
+                )
+            _calculate_token_usage(result, span)
+            span.__exit__(None, None, None)
+        elif hasattr(result, "_iterator"):
+            old_iterator = result._iterator
+
+            def new_iterator():
+                # type: () -> Iterator[MessageStreamEvent]
+                input_tokens = 0
+                output_tokens = 0
+                content_blocks = []
+                with capture_internal_exceptions():
+                    for event in old_iterator:
+                        if hasattr(event, "type"):
+                            if event.type == "message_start":
+                                usage = event.message.usage
+                                input_tokens += usage.input_tokens
+                                output_tokens += usage.output_tokens
+                            elif event.type == "content_block_start":
+                                pass
+                            elif event.type == "content_block_delta":
+                                content_blocks.append(event.delta.text)
+                            elif event.type == "content_block_stop":
+                                pass
+                            elif event.type == "message_delta":
+                                output_tokens += event.usage.output_tokens
+                            elif event.type == "message_stop":
+                                continue
+                        yield event
+
+                    if should_send_default_pii() and integration.include_prompts:
+                        complete_message = "".join(content_blocks)
+                        span.set_data(
+                            SPANDATA.AI_RESPONSES,
+                            [{"type": "text", "text": complete_message}],
+                        )
+                    total_tokens = input_tokens + output_tokens
+                    record_token_usage(span, input_tokens, output_tokens, total_tokens)
+                    span.set_data(SPANDATA.AI_STREAMING, True)
+                span.__exit__(None, None, None)
+
+            result._iterator = new_iterator()
+        else:
+            span.set_data("unknown_response", True)
+            span.__exit__(None, None, None)
+
+    return result
+
+
 def _wrap_message_create(f):
     # type: (Any) -> Any
+    def _execute_sync(f, *args, **kwargs):
+        # type: (Any, *Any, **Any) -> Any
+        gen = _sentry_patched_create_common(f, *args, **kwargs)
+
+        try:
+            f, args, kwargs = next(gen)
+        except StopIteration as e:
+            return e.value
+
+        try:
+            result = f(*args, **kwargs)
+            return gen.send(result)
+        except StopIteration as e:
+            return e.value
+
     @wraps(f)
     @ensure_integration_enabled(AnthropicIntegration, f)
-    def _sentry_patched_create(*args, **kwargs):
+    def _sentry_patched_create_sync(*args, **kwargs):
         # type: (*Any, **Any) -> Any
-        if "messages" not in kwargs:
-            return f(*args, **kwargs)
-
-        try:
-            iter(kwargs["messages"])
-        except TypeError:
-            return f(*args, **kwargs)
+        return _execute_sync(f, *args, **kwargs)
 
-        messages = list(kwargs["messages"])
-        model = kwargs.get("model")
+    return _sentry_patched_create_sync
 
-        span = sentry_sdk.start_span(
-            op=OP.ANTHROPIC_MESSAGES_CREATE,
-            description="Anthropic messages create",
-            origin=AnthropicIntegration.origin,
-        )
-        span.__enter__()
 
-        try:
-            result = f(*args, **kwargs)
-        except Exception as exc:
-            _capture_exception(exc)
-            span.__exit__(None, None, None)
-            raise exc from None
+def _wrap_async_message_create(f):
+    # type: (Any) -> Any
 
-        integration = sentry_sdk.get_client().get_integration(AnthropicIntegration)
+    async def _execute_async(f, *args, **kwargs):
+        # type: (Any, *Any, **Any) -> Any
+        gen = _sentry_patched_create_common(f, *args, **kwargs)
 
-        with capture_internal_exceptions():
-            span.set_data(SPANDATA.AI_MODEL_ID, model)
-            span.set_data(SPANDATA.AI_STREAMING, False)
-            if should_send_default_pii() and integration.include_prompts:
-                span.set_data(SPANDATA.AI_INPUT_MESSAGES, messages)
-            if hasattr(result, "content"):
-                if should_send_default_pii() and integration.include_prompts:
-                    span.set_data(
-                        SPANDATA.AI_RESPONSES,
-                        list(
-                            map(
-                                lambda message: {
-                                    "type": message.type,
-                                    "text": message.text,
-                                },
-                                result.content,
-                            )
-                        ),
-                    )
-                _calculate_token_usage(result, span)
-                span.__exit__(None, None, None)
-            elif hasattr(result, "_iterator"):
-                old_iterator = result._iterator
-
-                def new_iterator():
-                    # type: () -> Iterator[MessageStreamEvent]
-                    input_tokens = 0
-                    output_tokens = 0
-                    content_blocks = []
-                    with capture_internal_exceptions():
-                        for event in old_iterator:
-                            if hasattr(event, "type"):
-                                if event.type == "message_start":
-                                    usage = event.message.usage
-                                    input_tokens += usage.input_tokens
-                                    output_tokens += usage.output_tokens
-                                elif event.type == "content_block_start":
-                                    pass
-                                elif event.type == "content_block_delta":
-                                    content_blocks.append(event.delta.text)
-                                elif event.type == "content_block_stop":
-                                    pass
-                                elif event.type == "message_delta":
-                                    output_tokens += event.usage.output_tokens
-                                elif event.type == "message_stop":
-                                    continue
-                            yield event
-
-                        if should_send_default_pii() and integration.include_prompts:
-                            complete_message = "".join(content_blocks)
-                            span.set_data(
-                                SPANDATA.AI_RESPONSES,
-                                [{"type": "text", "text": complete_message}],
-                            )
-                        total_tokens = input_tokens + output_tokens
-                        record_token_usage(
-                            span, input_tokens, output_tokens, total_tokens
-                        )
-                        span.set_data(SPANDATA.AI_STREAMING, True)
-                    span.__exit__(None, None, None)
+        try:
+            f, args, kwargs = next(gen)
+        except StopIteration as e:
+            return await e.value
 
-                result._iterator = new_iterator()
-            else:
-                span.set_data("unknown_response", True)
-                span.__exit__(None, None, None)
+        try:
+            result = await f(*args, **kwargs)
+            return gen.send(result)
+        except StopIteration as e:
+            return e.value
 
-        return result
+    @wraps(f)
+    @ensure_integration_enabled(AnthropicIntegration, f)
+    async def _sentry_patched_create_async(*args, **kwargs):
+        # type: (*Any, **Any) -> Any
+        return await _execute_async(f, *args, **kwargs)
 
-    return _sentry_patched_create
+    return _sentry_patched_create_async