From da92544f47f66280f97c11cb4acbbd43eb55a84b Mon Sep 17 00:00:00 2001
From: Colin Chartier <colin.chartier@sentry.io>
Date: Thu, 25 Apr 2024 10:40:35 -0400
Subject: [PATCH] Add exception handling to AI monitoring

---
 sentry_sdk/{ai_analytics.py => ai_monitoring.py} | 15 +++++++++++++--
 sentry_sdk/integrations/langchain.py             |  2 +-
 sentry_sdk/integrations/openai.py                |  2 +-
 3 files changed, 15 insertions(+), 4 deletions(-)
 rename sentry_sdk/{ai_analytics.py => ai_monitoring.py} (77%)

diff --git a/sentry_sdk/ai_analytics.py b/sentry_sdk/ai_monitoring.py
similarity index 77%
rename from sentry_sdk/ai_analytics.py
rename to sentry_sdk/ai_monitoring.py
index 2fd39a4f98..f5f9cd7aad 100644
--- a/sentry_sdk/ai_analytics.py
+++ b/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
index f3058fe087..3e8ed89fff 100644
--- 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
index ae5c9e70ac..5af79c4d93 100644
--- 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