fixed tests

Author: antonpirker

sentry_sdk/integrations/openai.py

@@ -155,6 +155,7 @@ def _calculate_token_usage(
     )
 
 
+# TODO: rename to _set_input_data and _set_output_data
 def _set_request_data(span, kwargs, operation, integration):
     # type: (Span, dict[str, Any], str, Integration) -> None
     messages = kwargs.get("messages")
@@ -172,6 +173,7 @@ def _set_request_data(span, kwargs, operation, integration):
     ):
         set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages)
 
+    # TODO: make mapping and loop over kwargs to set attributes
     # Common attributes
     model = kwargs.get("model")
     set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "openai")
@@ -215,8 +217,8 @@ def _set_request_data(span, kwargs, operation, integration):
         )
 
 
-def _set_response_data(span, response, kwargs, integration):
-    # type: (Span, Any, dict[str, Any], Integration) -> None
+def _set_response_data(span, response, kwargs, integration, finish_span=True):
+    # type: (Span, Any, dict[str, Any], Integration, bool) -> None
     if hasattr(response, "model"):
         set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, response.model)
 
@@ -236,6 +238,7 @@ def _set_response_data(span, response, kwargs, integration):
                     SPANDATA.GEN_AI_RESPONSE_TEXT,
                     safe_serialize(response_text),
                 )
+        span.__exit__(None, None, None)
 
     elif hasattr(response, "output"):
         if should_send_default_pii() and integration.include_prompts:
@@ -318,6 +321,10 @@ async def new_iterator_async():
             response._iterator = new_iterator_async()
         else:
             response._iterator = new_iterator()
+    else:
+        set_data_normalized(span, "unknown_response", True)
+        if finish_span:
+            span.__exit__(None, None, None)
 
     _calculate_token_usage(messages, response, span, None, integration.count_tokens)
 
@@ -341,16 +348,18 @@ def _new_chat_completion_common(f, *args, **kwargs):
     model = kwargs.get("model")
     operation = "chat"
 
-    with sentry_sdk.start_span(
+    span = sentry_sdk.start_span(
         op=consts.OP.GEN_AI_CHAT,
         name=f"{operation} {model}",
         origin=OpenAIIntegration.origin,
-    ) as span:
-        _set_request_data(span, kwargs, operation, integration)
+    )
+    span.__enter__()
 
-        response = yield f, args, kwargs
+    _set_request_data(span, kwargs, operation, integration)
 
-        _set_response_data(span, response, kwargs, integration)
+    response = yield f, args, kwargs
+
+    _set_response_data(span, response, kwargs, integration)
 
     return response
 
@@ -443,7 +452,7 @@ def _new_embeddings_create_common(f, *args, **kwargs):
 
         response = yield f, args, kwargs
 
-        _set_response_data(span, response, kwargs, integration)
+        _set_response_data(span, response, kwargs, integration, finish_span=False)
 
         return response
 
@@ -525,18 +534,21 @@ def _new_responses_create_common(f, *args, **kwargs):
     model = kwargs.get("model")
     operation = "responses"
 
-    with sentry_sdk.start_span(
+    span = sentry_sdk.start_span(
         op=consts.OP.GEN_AI_RESPONSES,
         name=f"{operation} {model}",
         origin=OpenAIIntegration.origin,
-    ) as span:
-        _set_request_data(span, kwargs, operation, integration)
+    )
+    span.__enter__()
+    _set_request_data(span, kwargs, operation, integration)
 
-        response = yield f, args, kwargs
+    response = yield f, args, kwargs
 
-        _set_response_data(span, response, kwargs, integration)
+    _set_response_data(span, response, kwargs, integration)
 
-        return response
+    span.__exit__(None, None, None)
+
+    return response
 
 
 def _wrap_responses_create(f):

tests/integrations/openai/test_openai.py

@@ -1126,6 +1126,53 @@ async def test_ai_client_span_responses_async_api(sentry_init, capture_events):
     }
 
 
+@pytest.mark.asyncio
+@pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available")
+async def test_ai_client_span_streaming_responses_async_api(
+    sentry_init, capture_events
+):
+    sentry_init(
+        integrations=[OpenAIIntegration(include_prompts=True)],
+        traces_sample_rate=1.0,
+        send_default_pii=True,
+    )
+    events = capture_events()
+
+    client = AsyncOpenAI(api_key="z")
+    client.responses._post = AsyncMock(return_value=EXAMPLE_RESPONSE)
+
+    with start_transaction(name="openai tx"):
+        await client.responses.create(
+            model="gpt-4o",
+            instructions="You are a coding assistant that talks like a pirate.",
+            input="How do I check if a Python object is an instance of a class?",
+            stream=True,
+        )
+
+    (transaction,) = events
+    spans = transaction["spans"]
+
+    assert len(spans) == 1
+    assert spans[0]["op"] == "gen_ai.responses"
+    assert spans[0]["origin"] == "auto.ai.openai"
+    assert spans[0]["data"] == {
+        "ai.streaming": True,
+        "gen_ai.operation.name": "responses",
+        "gen_ai.request.messages": "How do I check if a Python object is an instance of a class?",
+        "gen_ai.request.model": "gpt-4o",
+        "gen_ai.response.model": "response-model-id",
+        "gen_ai.system": "openai",
+        "gen_ai.usage.input_tokens": 20,
+        "gen_ai.usage.input_tokens.cached": 5,
+        "gen_ai.usage.output_tokens": 10,
+        "gen_ai.usage.output_tokens.reasoning": 8,
+        "gen_ai.usage.total_tokens": 30,
+        "gen_ai.response.text": '[{"id": "message-id", "content": [{"annotations": [], "text": "the model response", "type": "output_text"}], "role": "assistant", "status": "completed", "type": "message"}]',
+        "thread.id": mock.ANY,
+        "thread.name": mock.ANY,
+    }
+
+
 @pytest.mark.asyncio
 @pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available")
 async def test_error_in_responses_async_api(sentry_init, capture_events):