Add exec_sync and exec_async to correctly handle function execution in patched method

Author: vetyy

sentry_sdk/integrations/anthropic.py

@@ -96,7 +96,7 @@ def _sentry_patched_create_common(f, *args, **kwargs):
     span.__enter__()
 
     try:
-        result = f(*args, **kwargs)
+        result = yield f(*args, **kwargs)
     except Exception as exc:
         _capture_exception(exc)
         span.__exit__(None, None, None)
@@ -173,21 +173,34 @@ def new_iterator():
 
 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)
+        f, args, kwargs = next(gen)
+        return gen.send(f(*args, **kwargs))
+
     @wraps(f)
     @ensure_integration_enabled(AnthropicIntegration, f)
     def _sentry_patched_create_sync(*args, **kwargs):
         # type: (*Any, **Any) -> Any
-        return _sentry_patched_create_common(f, *args, **kwargs)
+        return _execute_sync(f, *args, **kwargs)
 
     return _sentry_patched_create_sync
 
 
 def _wrap_async_message_create(f):
     # type: (Any) -> Any
+
+    async def _execute_async(f, *args, **kwargs):
+        # type: (Any, *Any, **Any) -> Any
+        gen = _sentry_patched_create_common(f, *args, **kwargs)
+        f, args, kwargs = next(gen)
+        return gen.send(await f(*args, **kwargs))
+
     @wraps(f)
     @ensure_integration_enabled(AnthropicIntegration, f)
     async def _sentry_patched_create_async(*args, **kwargs):
         # type: (*Any, **Any) -> Any
-        return await _sentry_patched_create_common(f, *args, **kwargs)
+        return await _execute_async(f, *args, **kwargs)
 
     return _sentry_patched_create_async

sentry_sdk/integrations/openai.py

@@ -136,7 +136,7 @@ def _new_chat_completion_common(f, *args, **kwargs):
     )
     span.__enter__()
     try:
-        res = f(*args, **kwargs)
+        res = yield f(*args, **kwargs)
     except Exception as e:
         _capture_exception(e)
         span.__exit__(None, None, None)
@@ -209,22 +209,35 @@ def new_iterator():
 
 def _wrap_chat_completion_create(f):
     # type: (Callable[..., Any]) -> Callable[..., Any]
+    def _execute_sync(f, *args, **kwargs):
+        # type: (Any, *Any, **Any) -> Any
+        gen = _new_chat_completion_common(f, *args, **kwargs)
+        f, args, kwargs = next(gen)
+        return gen.send(f(*args, **kwargs))
+
     @wraps(f)
     @ensure_integration_enabled(OpenAIIntegration, f)
     def _sentry_patched_create_sync(*args, **kwargs):
         # type: (*Any, **Any) -> Any
-        return _new_chat_completion_common(f, *args, **kwargs)
+        return _execute_sync(f, *args, **kwargs)
 
     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
+        gen = _new_chat_completion_common(f, *args, **kwargs)
+        f, args, kwargs = next(gen)
+        return gen.send(await f(*args, **kwargs))
+
     @wraps(f)
     @ensure_integration_enabled(OpenAIIntegration, f)
     async def _sentry_patched_create_async(*args, **kwargs):
         # type: (*Any, **Any) -> Any
-        return await _new_chat_completion_common(f, *args, **kwargs)
+        return await _execute_async(f, *args, **kwargs)
 
     return _sentry_patched_create_async
 
@@ -251,7 +264,7 @@ def _new_embeddings_create_common(f, *args, **kwargs):
         if "model" in kwargs:
             set_data_normalized(span, "ai.model_id", kwargs["model"])
         try:
-            response = f(*args, **kwargs)
+            response = yield f(*args, **kwargs)
         except Exception as e:
             _capture_exception(e)
             raise e from None
@@ -278,21 +291,34 @@ def _new_embeddings_create_common(f, *args, **kwargs):
 
 def _wrap_embeddings_create(f):
     # type: (Any) -> Any
+
+    def _execute_sync(f, *args, **kwargs):
+        # type: (Any, *Any, **Any) -> Any
+        gen = _new_embeddings_create_common(f, *args, **kwargs)
+        f, args, kwargs = next(gen)
+        return gen.send(f(*args, **kwargs))
+
     @wraps(f)
     @ensure_integration_enabled(OpenAIIntegration, f)
     def _sentry_patched_create_sync(*args, **kwargs):
         # type: (*Any, **Any) -> Any
-        return _new_embeddings_create_common(f, *args, **kwargs)
+        return _execute_sync(f, *args, **kwargs)
 
     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
+        gen = _new_embeddings_create_common(f, *args, **kwargs)
+        f, args, kwargs = next(gen)
+        return gen.send(await f(*args, **kwargs))
+
     @wraps(f)
     @ensure_integration_enabled(OpenAIIntegration, f)
     async def _sentry_patched_create_async(*args, **kwargs):
         # type: (*Any, **Any) -> Any
-        return await _new_embeddings_create_common(f, *args, **kwargs)
+        return await _execute_async(f, *args, **kwargs)
 
     return _sentry_patched_create_async