fix(AI): Make agents integrations set the span status in case of error

sentry_sdk/consts.py

@@ -765,18 +765,20 @@ class SPANSTATUS:
     CANCELLED = "cancelled"
     DATA_LOSS = "data_loss"
     DEADLINE_EXCEEDED = "deadline_exceeded"
+    ERROR = "error"  # OTel status code: https://opentelemetry.io/docs/concepts/signals/traces/#span-status
     FAILED_PRECONDITION = "failed_precondition"
     INTERNAL_ERROR = "internal_error"
     INVALID_ARGUMENT = "invalid_argument"
     NOT_FOUND = "not_found"
-    OK = "ok"
+    OK = "ok"  # HTTP 200 and OTel status code: https://opentelemetry.io/docs/concepts/signals/traces/#span-status
     OUT_OF_RANGE = "out_of_range"
     PERMISSION_DENIED = "permission_denied"
     RESOURCE_EXHAUSTED = "resource_exhausted"
     UNAUTHENTICATED = "unauthenticated"
     UNAVAILABLE = "unavailable"
     UNIMPLEMENTED = "unimplemented"
     UNKNOWN_ERROR = "unknown_error"
+    UNSET = "unset"  # OTel status code: https://opentelemetry.io/docs/concepts/signals/traces/#span-status
 
 
 class OP:

sentry_sdk/integrations/anthropic.py

@@ -4,9 +4,10 @@
 import sentry_sdk
 from sentry_sdk.ai.monitoring import record_token_usage
 from sentry_sdk.ai.utils import set_data_normalized, get_start_span_function
-from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS
 from sentry_sdk.integrations import _check_minimum_version, DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     event_from_exception,
@@ -52,6 +53,8 @@ def setup_once():
 
 def _capture_exception(exc):
     # type: (Any) -> None
+    set_span_errored()
+
     event, hint = event_from_exception(
         exc,
         client_options=sentry_sdk.get_client().options,
@@ -357,7 +360,13 @@ def _sentry_patched_create_sync(*args, **kwargs):
         integration = sentry_sdk.get_client().get_integration(AnthropicIntegration)
         kwargs["integration"] = integration
 
-        return _execute_sync(f, *args, **kwargs)
+        try:
+            return _execute_sync(f, *args, **kwargs)
+        finally:
+            span = sentry_sdk.get_current_span()
+            if span is not None and span.status == SPANSTATUS.ERROR:
+                with capture_internal_exceptions():
+                    span.__exit__(None, None, None)
 
     return _sentry_patched_create_sync
 
@@ -390,6 +399,12 @@ async def _sentry_patched_create_async(*args, **kwargs):
         integration = sentry_sdk.get_client().get_integration(AnthropicIntegration)
         kwargs["integration"] = integration
 
-        return await _execute_async(f, *args, **kwargs)
+        try:
+            return await _execute_async(f, *args, **kwargs)
+        finally:
+            span = sentry_sdk.get_current_span()
+            if span is not None and span.status == SPANSTATUS.ERROR:
+                with capture_internal_exceptions():
+                    span.__exit__(None, None, None)
 
     return _sentry_patched_create_async

sentry_sdk/integrations/cohere.py

@@ -7,6 +7,8 @@
 
 from typing import TYPE_CHECKING
 
+from sentry_sdk.tracing_utils import set_span_errored
+
 if TYPE_CHECKING:
     from typing import Any, Callable, Iterator
     from sentry_sdk.tracing import Span
@@ -84,6 +86,8 @@ def setup_once():
 
 def _capture_exception(exc):
     # type: (Any) -> None
+    set_span_errored()
+
     event, hint = event_from_exception(
         exc,
         client_options=sentry_sdk.get_client().options,

sentry_sdk/integrations/huggingface_hub.py

@@ -7,6 +7,7 @@
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     event_from_exception,
@@ -52,6 +53,8 @@ def setup_once():
 
 def _capture_exception(exc):
     # type: (Any) -> None
+    set_span_errored()
+
     event, hint = event_from_exception(
         exc,
         client_options=sentry_sdk.get_client().options,
@@ -127,8 +130,6 @@ def new_huggingface_task(*args, **kwargs):
         try:
             res = f(*args, **kwargs)
         except Exception as e:
-            # Error Handling
-            span.set_status("error")
             _capture_exception(e)
             span.__exit__(None, None, None)
             raise e from None

sentry_sdk/integrations/langchain.py

@@ -5,7 +5,7 @@
 import sentry_sdk
 from sentry_sdk.ai.monitoring import set_ai_pipeline_name
 from sentry_sdk.ai.utils import set_data_normalized, get_start_span_function
-from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.tracing import Span
@@ -116,7 +116,7 @@ def _handle_error(self, run_id, error):
 
             span_data = self.span_map[run_id]
             span = span_data.span
-            span.set_status("unknown")
+            span.set_status(SPANSTATUS.ERROR)
 
             sentry_sdk.capture_exception(error, span.scope)
 

sentry_sdk/integrations/openai.py

@@ -7,6 +7,7 @@
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     event_from_exception,
@@ -83,6 +84,8 @@ def _capture_exception(exc, manual_span_cleanup=True):
     # Close an eventually open span
     # We need to do this by hand because we are not using the start_span context manager
     current_span = sentry_sdk.get_current_span()
+    set_span_errored(current_span)
+
     if manual_span_cleanup and current_span is not None:
         current_span.__exit__(None, None, None)
 

sentry_sdk/integrations/openai_agents/spans/execute_tool.py

@@ -42,7 +42,7 @@ def update_execute_tool_span(span, agent, tool, result):
     if isinstance(result, str) and result.startswith(
         "An error occurred while running the tool"
     ):
-        span.set_status(SPANSTATUS.INTERNAL_ERROR)
+        span.set_status(SPANSTATUS.ERROR)
 
     if should_send_default_pii():
         span.set_data(SPANDATA.GEN_AI_TOOL_OUTPUT, result)

sentry_sdk/integrations/openai_agents/utils.py

@@ -3,6 +3,7 @@
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import event_from_exception, safe_serialize
 
 from typing import TYPE_CHECKING
@@ -20,6 +21,8 @@
 
 def _capture_exception(exc):
     # type: (Any) -> None
+    set_span_errored()
+
     event, hint = event_from_exception(
         exc,
         client_options=sentry_sdk.get_client().options,

sentry_sdk/tracing_utils.py

@@ -11,7 +11,7 @@
 import uuid
 
 import sentry_sdk
-from sentry_sdk.consts import OP, SPANDATA, SPANTEMPLATE
+from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS, SPANTEMPLATE
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     filename_for_module,
@@ -892,6 +892,19 @@ def get_current_span(scope=None):
     return current_span
 
 
+def set_span_errored(span=None):
+    # type: (Optional[Span]) -> None
+    """
+    Set the status of the current or given span to ERROR.
+    Also sets the status of the transaction (root span) to ERROR.
+    """
+    span = span or get_current_span()
+    if span is not None:
+        span.set_status(SPANSTATUS.ERROR)
+        if span.containing_transaction is not None:
+            span.containing_transaction.set_status(SPANSTATUS.ERROR)
+
+
 def _generate_sample_rand(
     trace_id,  # type: Optional[str]
     *,

tests/integrations/anthropic/test_anthropic.py

@@ -698,8 +698,54 @@ def test_exception_message_create(sentry_init, capture_events):
             max_tokens=1024,
         )
 
-    (event,) = events
+    (event, transaction) = events
     assert event["level"] == "error"
+    assert transaction["contexts"]["trace"]["status"] == "error"
+
+
+def test_span_status_error(sentry_init, capture_events):
+    sentry_init(integrations=[AnthropicIntegration()], traces_sample_rate=1.0)
+    events = capture_events()
+
+    with start_transaction(name="anthropic"):
+        client = Anthropic(api_key="z")
+        client.messages._post = mock.Mock(
+            side_effect=AnthropicError("API rate limit reached")
+        )
+        with pytest.raises(AnthropicError):
+            client.messages.create(
+                model="some-model",
+                messages=[{"role": "system", "content": "I'm throwing an exception"}],
+                max_tokens=1024,
+            )
+
+    (error, transaction) = events
+    assert error["level"] == "error"
+    assert transaction["spans"][0]["tags"]["status"] == "error"
+    assert transaction["contexts"]["trace"]["status"] == "error"
+
+
+@pytest.mark.asyncio
+async def test_span_status_error_async(sentry_init, capture_events):
+    sentry_init(integrations=[AnthropicIntegration()], traces_sample_rate=1.0)
+    events = capture_events()
+
+    with start_transaction(name="anthropic"):
+        client = AsyncAnthropic(api_key="z")
+        client.messages._post = AsyncMock(
+            side_effect=AnthropicError("API rate limit reached")
+        )
+        with pytest.raises(AnthropicError):
+            await client.messages.create(
+                model="some-model",
+                messages=[{"role": "system", "content": "I'm throwing an exception"}],
+                max_tokens=1024,
+            )
+
+    (error, transaction) = events
+    assert error["level"] == "error"
+    assert transaction["spans"][0]["tags"]["status"] == "error"
+    assert transaction["contexts"]["trace"]["status"] == "error"
 
 
 @pytest.mark.asyncio
@@ -718,8 +764,9 @@ async def test_exception_message_create_async(sentry_init, capture_events):
             max_tokens=1024,
         )
 
-    (event,) = events
+    (event, transaction) = events
     assert event["level"] == "error"
+    assert transaction["contexts"]["trace"]["status"] == "error"
 
 
 def test_span_origin(sentry_init, capture_events):

tests/integrations/cohere/test_cohere.py

@@ -167,6 +167,24 @@ def test_bad_chat(sentry_init, capture_events):
     assert event["level"] == "error"
 
 
+def test_span_status_error(sentry_init, capture_events):
+    sentry_init(integrations=[CohereIntegration()], traces_sample_rate=1.0)
+    events = capture_events()
+
+    with start_transaction(name="test"):
+        client = Client(api_key="z")
+        HTTPXClient.request = mock.Mock(
+            side_effect=httpx.HTTPError("API rate limit reached")
+        )
+        with pytest.raises(httpx.HTTPError):
+            client.chat(model="some-model", message="hello")
+
+    (error, transaction) = events
+    assert error["level"] == "error"
+    assert transaction["spans"][0]["tags"]["status"] == "error"
+    assert transaction["contexts"]["trace"]["status"] == "error"
+
+
 @pytest.mark.parametrize(
     "send_default_pii, include_prompts",
     [(True, True), (True, False), (False, True), (False, False)],

tests/integrations/huggingface_hub/test_huggingface_hub.py

@@ -654,6 +654,25 @@ def test_chat_completion_api_error(
     assert span["data"] == expected_data
 
 
+def test_span_status_error(sentry_init, capture_events, mock_hf_api_with_errors):
+    # type: (Any, Any, Any) -> None
+    sentry_init(traces_sample_rate=1.0)
+    events = capture_events()
+
+    client = InferenceClient(model="test-model")
+
+    with sentry_sdk.start_transaction(name="test"):
+        with pytest.raises(HfHubHTTPError):
+            client.chat_completion(
+                messages=[{"role": "user", "content": "Hello!"}],
+            )
+
+    (error, transaction) = events
+    assert error["level"] == "error"
+    assert transaction["spans"][0]["tags"]["status"] == "error"
+    assert transaction["contexts"]["trace"]["status"] == "error"
+
+
 @pytest.mark.parametrize("send_default_pii", [True, False])
 @pytest.mark.parametrize("include_prompts", [True, False])
 def test_chat_completion_with_tools(

tests/integrations/langchain/test_langchain.py

@@ -249,7 +249,7 @@ def test_langchain_error(sentry_init, capture_events):
         ]
     )
     global stream_result_mock
-    stream_result_mock = Mock(side_effect=Exception("API rate limit error"))
+    stream_result_mock = Mock(side_effect=ValueError("API rate limit error"))
     llm = MockOpenAI(
         model_name="gpt-3.5-turbo",
         temperature=0,
@@ -259,13 +259,53 @@ def test_langchain_error(sentry_init, capture_events):
 
     agent_executor = AgentExecutor(agent=agent, tools=[get_word_length], verbose=True)
 
-    with start_transaction(), pytest.raises(Exception):
+    with start_transaction(), pytest.raises(ValueError):
         list(agent_executor.stream({"input": "How many letters in the word eudca"}))
 
     error = events[0]
     assert error["level"] == "error"
 
 
+def test_span_status_error(sentry_init, capture_events):
+    sentry_init(
+        integrations=[LangchainIntegration(include_prompts=True)],
+        traces_sample_rate=1.0,
+    )
+    events = capture_events()
+
+    with start_transaction(name="test"):
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                (
+                    "system",
+                    "You are very powerful assistant, but don't know current events",
+                ),
+                ("user", "{input}"),
+                MessagesPlaceholder(variable_name="agent_scratchpad"),
+            ]
+        )
+        global stream_result_mock
+        stream_result_mock = Mock(side_effect=ValueError("API rate limit error"))
+        llm = MockOpenAI(
+            model_name="gpt-3.5-turbo",
+            temperature=0,
+            openai_api_key="badkey",
+        )
+        agent = create_openai_tools_agent(llm, [get_word_length], prompt)
+
+        agent_executor = AgentExecutor(
+            agent=agent, tools=[get_word_length], verbose=True
+        )
+
+        with pytest.raises(ValueError):
+            list(agent_executor.stream({"input": "How many letters in the word eudca"}))
+
+    (error, transaction) = events
+    assert error["level"] == "error"
+    assert transaction["spans"][0]["tags"]["status"] == "error"
+    assert transaction["contexts"]["trace"]["status"] == "error"
+
+
 def test_span_origin(sentry_init, capture_events):
     sentry_init(
         integrations=[LangchainIntegration()],