Add exception handling to AI monitoring

getsentry · Apr 25, 2024 · 27afb81 · 27afb81
1 parent 26a3eca
commit 27afb81
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 4 deletions.
diff --git a/sentry_sdk/ai_analytics.py → sentry_sdk/ai_monitoring.py b/sentry_sdk/ai_analytics.py → sentry_sdk/ai_monitoring.py
@@ -1,5 +1,6 @@
 from functools import wraps
 
+import sentry_sdk.utils
 from sentry_sdk import start_span
 from sentry_sdk.tracing import Span
 from sentry_sdk.utils import ContextVar
@@ -36,8 +37,18 @@ def wrapped(*args, **kwargs):
                     return f(*args, **kwargs)
                 else:
                     _ai_pipeline_name.set(description)
-                    res = f(*args, **kwargs)
-                    _ai_pipeline_name.set(None)
+                    try:
+                        res = f(*args, **kwargs)
+                    except Exception as e:
+                        event, hint = sentry_sdk.utils.event_from_exception(
+                            e,
+                            client_options=sentry_sdk.get_client().options,
+                            mechanism={"type": "ai_monitoring", "handled": False},
+                        )
+                        sentry_sdk.capture_event(event, hint=hint)
+                        raise e from None
+                    finally:
+                        _ai_pipeline_name.set(None)
                     return res
 
         return wrapped

diff --git a/sentry_sdk/integrations/langchain.py b/sentry_sdk/integrations/langchain.py
@@ -3,7 +3,7 @@
 
 import sentry_sdk
 from sentry_sdk._types import TYPE_CHECKING
-from sentry_sdk.ai_analytics import set_ai_pipeline_name, record_token_usage
+from sentry_sdk.ai_monitoring import set_ai_pipeline_name, record_token_usage
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations._ai_common import set_data_normalized
 from sentry_sdk.scope import should_send_default_pii

diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py
@@ -2,7 +2,7 @@
 
 from sentry_sdk import consts
 from sentry_sdk._types import TYPE_CHECKING
-from sentry_sdk.ai_analytics import record_token_usage
+from sentry_sdk.ai_monitoring import record_token_usage
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations._ai_common import set_data_normalized