From 83c08152adb0cb189d017d7cebcd39e2ebe59e13 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Tue, 8 Sep 2026 16:55:22 +0200 Subject: [PATCH] test(huggingface_hub): Parametrize tests on the streaming trace lifecycle --- .../huggingface_hub/test_huggingface_hub.py | 515 +++++++++++++++++- 1 file changed, 499 insertions(+), 16 deletions(-) diff --git a/tests/integrations/huggingface_hub/test_huggingface_hub.py b/tests/integrations/huggingface_hub/test_huggingface_hub.py index c52c38fa6d..c1b38d4377 100644 --- a/tests/integrations/huggingface_hub/test_huggingface_hub.py +++ b/tests/integrations/huggingface_hub/test_huggingface_hub.py @@ -467,6 +467,7 @@ def mock_hf_chat_completion_api_streaming_tools(httpx_mock): yield rsps +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize("send_default_pii", [True, False]) @@ -479,17 +480,73 @@ def test_text_generation( include_prompts: "Any", mock_hf_text_generation_api: "Any", stream_gen_ai_spans: "Any", + span_streaming: "Any", ) -> None: sentry_init( traces_sample_rate=1.0, send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration(include_prompts=include_prompts)], stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", ) client = InferenceClient(model="test-model") - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("transaction", "span") + + with sentry_sdk.traces.start_span(name="test"): + client.text_generation( + "Hello", + stream=False, + details=True, + ) + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + span = None + for sp in spans: + if "sentry.op" in sp["attributes"] and sp["attributes"][ + "sentry.op" + ].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + elif "sentry.op" in sp["attributes"]: + # there should be no other spans, just the gen_ai span + # and optionally some http.client spans from talking to the hf api + assert sp["attributes"]["sentry.op"] == "http.client" + + assert span is not None + + assert span["attributes"]["sentry.op"] == "gen_ai.text_completion" + assert span["name"] == "text_completion test-model" + assert span["attributes"]["sentry.origin"] == "auto.ai.huggingface_hub" + + expected_data = { + "gen_ai.operation.name": "text_completion", + "gen_ai.request.model": "test-model", + "gen_ai.response.finish_reasons": "length", + "gen_ai.response.streaming": False, + "gen_ai.usage.total_tokens": 10, + "sentry.op": "gen_ai.text_completion", + "sentry.origin": "auto.ai.huggingface_hub", + } + + if send_default_pii and include_prompts: + expected_data["gen_ai.request.messages"] = "Hello" + expected_data["gen_ai.response.text"] = ( + "[mocked] Hello! How can i help you?" + ) + + if not send_default_pii or not include_prompts: + assert "gen_ai.request.messages" not in expected_data + assert "gen_ai.response.text" not in expected_data + + assert span["attributes"] == ApproxDict(expected_data) + + # text generation does not set the response model + assert "gen_ai.response.model" not in span["attributes"] + elif stream_gen_ai_spans: items = capture_items("transaction", "span") with sentry_sdk.start_transaction(name="test"): @@ -592,6 +649,7 @@ def test_text_generation( assert "gen_ai.response.model" not in span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize("send_default_pii", [True, False]) @@ -604,17 +662,73 @@ def test_text_generation_streaming( include_prompts: "Any", mock_hf_text_generation_api_streaming: "Any", stream_gen_ai_spans: "Any", + span_streaming: "Any", ) -> None: sentry_init( traces_sample_rate=1.0, send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration(include_prompts=include_prompts)], stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", ) client = InferenceClient(model="test-model") - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + with sentry_sdk.traces.start_span(name="test"): + for _ in client.text_generation( + prompt="Hello", + stream=True, + details=True, + ): + pass + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + span = None + for sp in spans: + if "sentry.op" in sp["attributes"] and sp["attributes"][ + "sentry.op" + ].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + elif "sentry.op" in sp["attributes"]: + # there should be no other spans, just the gen_ai span + # and optionally some http.client spans from talking to the hf api + assert sp["attributes"]["sentry.op"] == "http.client" + + assert span is not None + + assert span["attributes"]["sentry.op"] == "gen_ai.text_completion" + assert span["name"] == "text_completion test-model" + assert span["attributes"]["sentry.origin"] == "auto.ai.huggingface_hub" + + expected_data = { + "gen_ai.operation.name": "text_completion", + "gen_ai.request.model": "test-model", + "gen_ai.response.finish_reasons": "length", + "gen_ai.response.streaming": True, + "gen_ai.usage.total_tokens": 10, + "sentry.environment": "production", + "sentry.op": "gen_ai.text_completion", + "sentry.origin": "auto.ai.huggingface_hub", + } + + if send_default_pii and include_prompts: + expected_data["gen_ai.request.messages"] = "Hello" + expected_data["gen_ai.response.text"] = "the mocked model response" + + if not send_default_pii or not include_prompts: + assert "gen_ai.request.messages" not in expected_data + assert "gen_ai.response.text" not in expected_data + + assert span["attributes"] == ApproxDict(expected_data) + + # text generation does not set the response model + assert "gen_ai.response.model" not in span["attributes"] + elif stream_gen_ai_spans: items = capture_items("transaction", "span") with sentry_sdk.start_transaction(name="test"): @@ -716,6 +830,7 @@ def test_text_generation_streaming( assert "gen_ai.response.model" not in span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize("send_default_pii", [True, False]) @@ -728,17 +843,87 @@ def test_chat_completion( include_prompts: "Any", mock_hf_chat_completion_api: "Any", stream_gen_ai_spans: "Any", + span_streaming: "Any", ) -> None: sentry_init( traces_sample_rate=1.0, send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration(include_prompts=include_prompts)], stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", ) client = get_hf_provider_inference_client() - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + with sentry_sdk.traces.start_span(name="test"): + client.chat_completion( + messages=[ + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "Hello!"}, + ], + stream=False, + ) + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + span = None + for sp in spans: + if "sentry.op" in sp["attributes"] and sp["attributes"][ + "sentry.op" + ].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + elif "sentry.op" in sp["attributes"]: + # there should be no other spans, just the gen_ai span + # and optionally some http.client spans from talking to the hf api + assert sp["attributes"]["sentry.op"] == "http.client" + + assert span is not None + + assert span["attributes"]["sentry.op"] == "gen_ai.chat" + assert span["name"] == "chat test-model" + assert span["attributes"]["sentry.origin"] == "auto.ai.huggingface_hub" + + expected_data = { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "test-model", + "gen_ai.response.finish_reasons": "stop", + "gen_ai.response.model": "test-model-123", + "gen_ai.response.streaming": False, + "gen_ai.usage.input_tokens": 10, + "gen_ai.usage.output_tokens": 8, + "gen_ai.usage.total_tokens": 18, + "sentry.environment": "production", + "sentry.op": "gen_ai.chat", + "sentry.origin": "auto.ai.huggingface_hub", + } + + if send_default_pii and include_prompts: + expected_data["gen_ai.request.messages"] = safe_serialize( + [ + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "Hello!"}, + ] + ) + expected_data["gen_ai.response.text"] = ( + "[mocked] Hello! How can I help you today?" + ) + + if not send_default_pii or not include_prompts: + assert "gen_ai.request.messages" not in expected_data + assert "gen_ai.response.text" not in expected_data + + assert span["attributes"] == ApproxDict(expected_data) + elif stream_gen_ai_spans: items = capture_items("transaction", "span") with sentry_sdk.start_transaction(name="test"): @@ -856,6 +1041,7 @@ def test_chat_completion( assert span["data"] == ApproxDict(expected_data) +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize("send_default_pii", [True, False]) @@ -868,17 +1054,89 @@ def test_chat_completion_streaming( include_prompts: "Any", mock_hf_chat_completion_api_streaming: "Any", stream_gen_ai_spans: "Any", + span_streaming: "Any", ) -> None: sentry_init( traces_sample_rate=1.0, send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration(include_prompts=include_prompts)], stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", ) client = get_hf_provider_inference_client() - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + with sentry_sdk.traces.start_span(name="test"): + _ = list( + client.chat_completion( + [ + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "Hello!"}, + ], + stream=True, + ) + ) + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + span = None + for sp in spans: + if "sentry.op" in sp["attributes"] and sp["attributes"][ + "sentry.op" + ].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + elif "sentry.op" in sp["attributes"]: + # there should be no other spans, just the gen_ai span + # and optionally some http.client spans from talking to the hf api + assert sp["attributes"]["sentry.op"] == "http.client" + + assert span is not None + + assert span["attributes"]["sentry.op"] == "gen_ai.chat" + assert span["name"] == "chat test-model" + assert span["attributes"]["sentry.origin"] == "auto.ai.huggingface_hub" + + expected_data = { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "test-model", + "gen_ai.response.finish_reasons": "stop", + "gen_ai.response.model": "test-model-123", + "gen_ai.response.streaming": True, + "sentry.environment": "production", + "sentry.op": "gen_ai.chat", + "sentry.origin": "auto.ai.huggingface_hub", + } + # usage is not available in older versions of the library + if HF_VERSION and HF_VERSION >= (0, 26, 0): + expected_data["gen_ai.usage.input_tokens"] = 183 + expected_data["gen_ai.usage.output_tokens"] = 14 + expected_data["gen_ai.usage.total_tokens"] = 197 + + if send_default_pii and include_prompts: + expected_data["gen_ai.request.messages"] = safe_serialize( + [ + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "Hello!"}, + ] + ) + expected_data["gen_ai.response.text"] = "the mocked model response" + + if not send_default_pii or not include_prompts: + assert "gen_ai.request.messages" not in expected_data + assert "gen_ai.response.text" not in expected_data + + assert span["attributes"] == ApproxDict(expected_data) + elif stream_gen_ai_spans: items = capture_items("transaction", "span") with sentry_sdk.start_transaction(name="test"): @@ -1000,6 +1258,7 @@ def test_chat_completion_streaming( assert span["data"] == ApproxDict(expected_data) +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) def test_chat_completion_api_error( @@ -1008,14 +1267,61 @@ def test_chat_completion_api_error( capture_items: "Any", mock_hf_api_with_errors: "Any", stream_gen_ai_spans: "Any", + span_streaming: "Any", ) -> None: sentry_init( traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", ) client = get_hf_provider_inference_client() - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("event", "span") + + with sentry_sdk.traces.start_span(name="test"), pytest.raises(HfHubHTTPError): + client.chat_completion( + messages=[{"role": "user", "content": "Hello!"}], + ) + + (error,) = (item.payload for item in items if item.type == "event") + + assert error["exception"]["values"][0]["mechanism"]["type"] == "huggingface_hub" + assert not error["exception"]["values"][0]["mechanism"]["handled"] + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + span = None + for sp in spans: + if "sentry.op" in sp["attributes"] and sp["attributes"][ + "sentry.op" + ].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + elif "sentry.op" in sp["attributes"]: + # there should be no other spans, just the gen_ai span + # and optionally some http.client spans from talking to the hf api + assert sp["attributes"]["sentry.op"] == "http.client" + + assert span is not None + + assert span["attributes"]["sentry.op"] == "gen_ai.chat" + assert span["name"] == "chat test-model" + assert span["attributes"]["sentry.origin"] == "auto.ai.huggingface_hub" + assert span["status"] == "error" + + (manual_span,) = (span for span in spans if span["name"] == "test") + + assert error["contexts"]["trace"]["trace_id"] == manual_span["trace_id"] + expected_data = { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "test-model", + "sentry.environment": "production", + "sentry.op": "gen_ai.chat", + "sentry.origin": "auto.ai.huggingface_hub", + } + assert span["attributes"] == ApproxDict(expected_data) + elif stream_gen_ai_spans: items = capture_items("event", "transaction", "span") with sentry_sdk.start_transaction(name="test"), pytest.raises(HfHubHTTPError): @@ -1105,6 +1411,7 @@ def test_chat_completion_api_error( assert span["data"] == ApproxDict(expected_data) +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) def test_span_status_error( @@ -1113,15 +1420,44 @@ def test_span_status_error( capture_items: "Any", mock_hf_api_with_errors: "Any", stream_gen_ai_spans: "Any", + span_streaming: "Any", ) -> None: client = get_hf_provider_inference_client() sentry_init( traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", ) - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("event", "span") + + with sentry_sdk.traces.start_span(name="test"), pytest.raises(HfHubHTTPError): + client.chat_completion( + messages=[{"role": "user", "content": "Hello!"}], + ) + + (error,) = [item.payload for item in items if item.type == "event"] + assert error["level"] == "error" + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + span = None + for sp in spans: + if "sentry.op" in sp["attributes"] and sp["attributes"][ + "sentry.op" + ].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + elif "sentry.op" in sp["attributes"]: + # there should be no other spans, just the gen_ai span + # and optionally some http.client spans from talking to the hf api + assert sp["attributes"]["sentry.op"] == "http.client" + + assert span is not None + assert span["status"] == "error" + elif stream_gen_ai_spans: items = capture_items("event", "transaction", "span") with sentry_sdk.start_transaction(name="test"), pytest.raises(HfHubHTTPError): @@ -1171,6 +1507,7 @@ def test_span_status_error( assert span["tags"]["status"] == "internal_error" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize("send_default_pii", [True, False]) @@ -1183,12 +1520,14 @@ def test_chat_completion_with_tools( include_prompts: "Any", mock_hf_chat_completion_api_tools: "Any", stream_gen_ai_spans: "Any", + span_streaming: "Any", ): sentry_init( traces_sample_rate=1.0, send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration(include_prompts=include_prompts)], stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", ) client = get_hf_provider_inference_client() @@ -1208,7 +1547,65 @@ def test_chat_completion_with_tools( } ] - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + with sentry_sdk.traces.start_span(name="test"): + client.chat_completion( + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=tools, + tool_choice="auto", + ) + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + span = None + for sp in spans: + if "sentry.op" in sp["attributes"] and sp["attributes"][ + "sentry.op" + ].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + elif "sentry.op" in sp["attributes"]: + # there should be no other spans, just the gen_ai span + # and optionally some http.client spans from talking to the hf api + assert sp["attributes"]["sentry.op"] == "http.client" + + assert span is not None + + assert span["attributes"]["sentry.op"] == "gen_ai.chat" + assert span["name"] == "chat test-model" + assert span["attributes"]["sentry.origin"] == "auto.ai.huggingface_hub" + + expected_data = { + "gen_ai.operation.name": "chat", + "gen_ai.request.available_tools": '[{"type": "function", "function": {"name": "get_weather", "description": "Get current weather", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}}}]', + "gen_ai.request.model": "test-model", + "gen_ai.response.finish_reasons": "tool_calls", + "gen_ai.response.model": "test-model-123", + "gen_ai.usage.input_tokens": 10, + "gen_ai.usage.output_tokens": 8, + "gen_ai.usage.total_tokens": 18, + "sentry.environment": "production", + "sentry.op": "gen_ai.chat", + "sentry.origin": "auto.ai.huggingface_hub", + } + + if send_default_pii and include_prompts: + expected_data["gen_ai.request.messages"] = ( + '[{"role": "user", "content": "What is the weather in Paris?"}]' + ) + expected_data["gen_ai.response.tool_calls"] = ( + '[{"function": {"arguments": {"location": "Paris"}, "name": "get_weather", "description": "None"}, "id": "call_123", "type": "function"}]' + ) + + if not send_default_pii or not include_prompts: + assert "gen_ai.request.messages" not in expected_data + assert "gen_ai.response.text" not in expected_data + assert "gen_ai.response.tool_calls" not in expected_data + + assert span["attributes"] == ApproxDict(expected_data) + elif stream_gen_ai_spans: items = capture_items("transaction", "span") with sentry_sdk.start_transaction(name="test"): @@ -1318,6 +1715,7 @@ def test_chat_completion_with_tools( assert span["data"] == ApproxDict(expected_data) +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize("send_default_pii", [True, False]) @@ -1330,12 +1728,14 @@ def test_chat_completion_streaming_with_tools( include_prompts: "Any", mock_hf_chat_completion_api_streaming_tools: "Any", stream_gen_ai_spans: "Any", + span_streaming: "Any", ) -> None: sentry_init( traces_sample_rate=1.0, send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration(include_prompts=include_prompts)], stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", ) client = get_hf_provider_inference_client() @@ -1355,7 +1755,74 @@ def test_chat_completion_streaming_with_tools( } ] - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + with sentry_sdk.traces.start_span(name="test"): + _ = list( + client.chat_completion( + messages=[ + {"role": "user", "content": "What is the weather in Paris?"} + ], + stream=True, + tools=tools, + tool_choice="auto", + ) + ) + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + span = None + for sp in spans: + if "sentry.op" in sp["attributes"] and sp["attributes"][ + "sentry.op" + ].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + elif "sentry.op" in sp["attributes"]: + # there should be no other spans, just the gen_ai span + # and optionally some http.client spans from talking to the hf api + assert sp["attributes"]["sentry.op"] == "http.client" + + assert span is not None + + assert span["attributes"]["sentry.op"] == "gen_ai.chat" + assert span["name"] == "chat test-model" + assert span["attributes"]["sentry.origin"] == "auto.ai.huggingface_hub" + + expected_data = { + "gen_ai.operation.name": "chat", + "gen_ai.request.available_tools": '[{"type": "function", "function": {"name": "get_weather", "description": "Get current weather", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}}}]', + "gen_ai.request.model": "test-model", + "gen_ai.response.finish_reasons": "tool_calls", + "gen_ai.response.model": "test-model-123", + "gen_ai.response.streaming": True, + "sentry.environment": "production", + "sentry.op": "gen_ai.chat", + "sentry.origin": "auto.ai.huggingface_hub", + } + + if HF_VERSION and HF_VERSION >= (0, 26, 0): + expected_data["gen_ai.usage.input_tokens"] = 183 + expected_data["gen_ai.usage.output_tokens"] = 14 + expected_data["gen_ai.usage.total_tokens"] = 197 + + if send_default_pii and include_prompts: + expected_data["gen_ai.request.messages"] = ( + '[{"role": "user", "content": "What is the weather in Paris?"}]' + ) + expected_data["gen_ai.response.text"] = "response with tool calls follows" + expected_data["gen_ai.response.tool_calls"] = ( + '[{"function": {"arguments": {"location": "Paris"}, "name": "get_weather"}, "id": "call_123", "type": "function", "index": "None"}]' + ) + + if not send_default_pii or not include_prompts: + assert "gen_ai.request.messages" not in expected_data + assert "gen_ai.response.text" not in expected_data + assert "gen_ai.response.tool_calls" not in expected_data + + assert span["attributes"] == ApproxDict(expected_data) + elif stream_gen_ai_spans: items = capture_items("transaction", "span") with sentry_sdk.start_transaction(name="test"): @@ -1585,6 +2052,7 @@ def test_chat_completion_streaming_with_tools( ), ], ) +@pytest.mark.parametrize("span_streaming", [True, False]) def test_text_generation_data_collection( sentry_init: "Any", capture_events: "Any", @@ -1596,12 +2064,14 @@ def test_text_generation_data_collection( expected_present: "Any", expected_absent: "Any", stream_gen_ai_spans: "Any", + span_streaming: "Any", ) -> None: sentry_init_kwargs = dict( traces_sample_rate=1.0, send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration(include_prompts=include_prompts)], stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -1612,14 +2082,15 @@ def test_text_generation_data_collection( captured = ( capture_items("transaction", "span") - if stream_gen_ai_spans + if span_streaming or stream_gen_ai_spans else capture_events() ) with sentry_sdk.start_transaction(name="test"): client.text_generation("Hello", stream=False, details=True) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: + sentry_sdk.flush() spans = [item.payload for item in captured if item.type == "span"] (span,) = [ sp for sp in spans if sp["attributes"]["sentry.op"].startswith("gen_ai") @@ -1735,6 +2206,7 @@ def test_text_generation_data_collection( ), ], ) +@pytest.mark.parametrize("span_streaming", [True, False]) def test_text_generation_streaming_data_collection( sentry_init: "Any", capture_events: "Any", @@ -1746,12 +2218,14 @@ def test_text_generation_streaming_data_collection( expected_present: "Any", expected_absent: "Any", stream_gen_ai_spans: "Any", + span_streaming: "Any", ) -> None: sentry_init_kwargs = dict( traces_sample_rate=1.0, send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration(include_prompts=include_prompts)], stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -1762,7 +2236,7 @@ def test_text_generation_streaming_data_collection( captured = ( capture_items("transaction", "span") - if stream_gen_ai_spans + if span_streaming or stream_gen_ai_spans else capture_events() ) @@ -1770,7 +2244,8 @@ def test_text_generation_streaming_data_collection( for _ in client.text_generation(prompt="Hello", stream=True, details=True): pass - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: + sentry_sdk.flush() spans = [item.payload for item in captured if item.type == "span"] (span,) = [ sp for sp in spans if sp["attributes"]["sentry.op"].startswith("gen_ai") @@ -1895,6 +2370,7 @@ def test_text_generation_streaming_data_collection( ), ], ) +@pytest.mark.parametrize("span_streaming", [True, False]) def test_chat_completion_data_collection_tools( sentry_init: "Any", capture_events: "Any", @@ -1906,12 +2382,14 @@ def test_chat_completion_data_collection_tools( expected_present: "Any", expected_absent: "Any", stream_gen_ai_spans: "Any", + span_streaming: "Any", ) -> None: sentry_init_kwargs = dict( traces_sample_rate=1.0, send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration(include_prompts=include_prompts)], stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -1922,7 +2400,7 @@ def test_chat_completion_data_collection_tools( captured = ( capture_items("transaction", "span") - if stream_gen_ai_spans + if span_streaming or stream_gen_ai_spans else capture_events() ) @@ -1933,7 +2411,8 @@ def test_chat_completion_data_collection_tools( tool_choice="auto", ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: + sentry_sdk.flush() spans = [item.payload for item in captured if item.type == "span"] (span,) = [ sp for sp in spans if sp["attributes"]["sentry.op"].startswith("gen_ai") @@ -2071,6 +2550,7 @@ def test_chat_completion_data_collection_tools( ), ], ) +@pytest.mark.parametrize("span_streaming", [True, False]) def test_chat_completion_streaming_data_collection_tools( sentry_init: "Any", capture_events: "Any", @@ -2082,12 +2562,14 @@ def test_chat_completion_streaming_data_collection_tools( expected_present: "Any", expected_absent: "Any", stream_gen_ai_spans: "Any", + span_streaming: "Any", ) -> None: sentry_init_kwargs = dict( traces_sample_rate=1.0, send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration(include_prompts=include_prompts)], stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -2098,7 +2580,7 @@ def test_chat_completion_streaming_data_collection_tools( captured = ( capture_items("transaction", "span") - if stream_gen_ai_spans + if span_streaming or stream_gen_ai_spans else capture_events() ) @@ -2111,7 +2593,8 @@ def test_chat_completion_streaming_data_collection_tools( ): pass - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: + sentry_sdk.flush() spans = [item.payload for item in captured if item.type == "span"] (span,) = [ sp for sp in spans if sp["attributes"]["sentry.op"].startswith("gen_ai")