diff --git a/tests/integrations/huggingface_hub/test_huggingface_hub.py b/tests/integrations/huggingface_hub/test_huggingface_hub.py index 8d8e6a9b2e..a11d7b9169 100644 --- a/tests/integrations/huggingface_hub/test_huggingface_hub.py +++ b/tests/integrations/huggingface_hub/test_huggingface_hub.py @@ -9,7 +9,7 @@ import sentry_sdk from sentry_sdk.consts import SPANDATA from sentry_sdk.integrations.huggingface_hub import HuggingfaceHubIntegration -from sentry_sdk.utils import package_version, safe_serialize +from sentry_sdk.utils import package_version try: from huggingface_hub.utils._errors import HfHubHTTPError @@ -467,789 +467,401 @@ def mock_hf_chat_completion_api_streaming_tools(httpx_mock): yield rsps -@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]) @pytest.mark.parametrize("include_prompts", [True, False]) def test_text_generation( sentry_init: "Any", capture_events: "Any", - capture_items: "Any", send_default_pii: "Any", include_prompts: "Any", mock_hf_text_generation_api: "Any", - stream_gen_ai_spans: "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, + stream_gen_ai_spans=False, ) client = InferenceClient(model="test-model") + events = capture_events() - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with sentry_sdk.start_transaction(name="test"): - client.text_generation( - "Hello", - stream=False, - details=True, - ) - - spans = [item.payload for item in items if item.type == "span"] - span = None - for sp in spans: - if sp["attributes"]["sentry.op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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, - "process.runtime.name": mock.ANY, - "process.runtime.version": mock.ANY, - "sentry.environment": "production", - "sentry.op": "gen_ai.text_completion", - "sentry.origin": "auto.ai.huggingface_hub", - "sentry.release": mock.ANY, - "sentry.sdk.name": "sentry.python", - "sentry.sdk.version": mock.ANY, - "sentry.segment.id": mock.ANY, - "sentry.segment.name": "test", - "server.address": mock.ANY, - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } - - 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"] == expected_data - - # text generation does not set the response model - assert "gen_ai.response.model" not in span["attributes"] - else: - events = capture_events() + with sentry_sdk.start_transaction(name="test"): + client.text_generation( + "Hello", + stream=False, + details=True, + ) - with sentry_sdk.start_transaction(name="test"): - client.text_generation( - "Hello", - stream=False, - details=True, - ) + (transaction,) = events - (transaction,) = events - - span = None - for sp in transaction["spans"]: - if sp["op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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["op"] == "http.client" - - assert span is not None - - assert span["op"] == "gen_ai.text_completion" - assert span["description"] == "text_completion test-model" - assert span["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, - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } + span = None + for sp in transaction["spans"]: + if sp["op"].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + else: + # 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["op"] == "http.client" + + assert span is not None + + assert span["op"] == "gen_ai.text_completion" + assert span["description"] == "text_completion test-model" + assert span["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, + "thread.id": mock.ANY, + "thread.name": mock.ANY, + } - 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 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 + 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["data"] == expected_data + assert span["data"] == expected_data - # text generation does not set the response model - assert "gen_ai.response.model" not in span["data"] + # text generation does not set the response model + assert "gen_ai.response.model" not in span["data"] -@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]) @pytest.mark.parametrize("include_prompts", [True, False]) def test_text_generation_streaming( sentry_init: "Any", capture_events: "Any", - capture_items: "Any", send_default_pii: "Any", include_prompts: "Any", mock_hf_text_generation_api_streaming: "Any", - stream_gen_ai_spans: "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, + stream_gen_ai_spans=False, ) client = InferenceClient(model="test-model") + events = capture_events() - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with sentry_sdk.start_transaction(name="test"): - for _ in client.text_generation( - prompt="Hello", - stream=True, - details=True, - ): - pass - - spans = [item.payload for item in items if item.type == "span"] - span = None - for sp in spans: - if sp["attributes"]["sentry.op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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, - "process.runtime.name": mock.ANY, - "process.runtime.version": mock.ANY, - "sentry.environment": "production", - "sentry.op": "gen_ai.text_completion", - "sentry.origin": "auto.ai.huggingface_hub", - "sentry.release": mock.ANY, - "sentry.sdk.name": "sentry.python", - "sentry.sdk.version": mock.ANY, - "sentry.segment.id": mock.ANY, - "sentry.segment.name": "test", - "server.address": mock.ANY, - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } - - 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"] == expected_data + with sentry_sdk.start_transaction(name="test"): + for _ in client.text_generation( + prompt="Hello", + stream=True, + details=True, + ): + pass - # text generation does not set the response model - assert "gen_ai.response.model" not in span["attributes"] - else: - events = capture_events() + (transaction,) = events - with sentry_sdk.start_transaction(name="test"): - for _ in client.text_generation( - prompt="Hello", - stream=True, - details=True, - ): - pass - - (transaction,) = events - - span = None - for sp in transaction["spans"]: - if sp["op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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["op"] == "http.client" - - assert span is not None - - assert span["op"] == "gen_ai.text_completion" - assert span["description"] == "text_completion test-model" - assert span["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, - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } + span = None + for sp in transaction["spans"]: + if sp["op"].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + else: + # 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["op"] == "http.client" + + assert span is not None + + assert span["op"] == "gen_ai.text_completion" + assert span["description"] == "text_completion test-model" + assert span["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, + "thread.id": mock.ANY, + "thread.name": mock.ANY, + } - 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 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 + 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["data"] == expected_data + assert span["data"] == expected_data - # text generation does not set the response model - assert "gen_ai.response.model" not in span["data"] + # text generation does not set the response model + assert "gen_ai.response.model" not in span["data"] -@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]) @pytest.mark.parametrize("include_prompts", [True, False]) def test_chat_completion( sentry_init: "Any", capture_events: "Any", - capture_items: "Any", send_default_pii: "Any", include_prompts: "Any", mock_hf_chat_completion_api: "Any", - stream_gen_ai_spans: "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, + stream_gen_ai_spans=False, ) client = get_hf_provider_inference_client() + events = capture_events() - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with sentry_sdk.start_transaction(name="test"): - client.chat_completion( - messages=[ - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "Hello!"}, - ], - stream=False, - ) - - spans = [item.payload for item in items if item.type == "span"] - span = None - for sp in spans: - if sp["attributes"]["sentry.op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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, - "process.runtime.name": mock.ANY, - "process.runtime.version": mock.ANY, - "sentry.environment": "production", - "sentry.op": "gen_ai.chat", - "sentry.origin": "auto.ai.huggingface_hub", - "sentry.release": mock.ANY, - "sentry.sdk.name": "sentry.python", - "sentry.sdk.version": mock.ANY, - "sentry.segment.id": mock.ANY, - "sentry.segment.name": "test", - "server.address": mock.ANY, - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } - - 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 + with sentry_sdk.start_transaction(name="test"): + client.chat_completion( + messages=[{"role": "user", "content": "Hello!"}], + stream=False, + ) - assert span["attributes"] == expected_data - else: - events = capture_events() + (transaction,) = events - with sentry_sdk.start_transaction(name="test"): - client.chat_completion( - messages=[{"role": "user", "content": "Hello!"}], - stream=False, - ) - - (transaction,) = events - - span = None - for sp in transaction["spans"]: - if sp["op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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["op"] == "http.client" - - assert span is not None - - assert span["op"] == "gen_ai.chat" - assert span["description"] == "chat test-model" - assert span["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, - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } + span = None + for sp in transaction["spans"]: + if sp["op"].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + else: + # 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["op"] == "http.client" + + assert span is not None + + assert span["op"] == "gen_ai.chat" + assert span["description"] == "chat test-model" + assert span["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, + "thread.id": mock.ANY, + "thread.name": mock.ANY, + } - if send_default_pii and include_prompts: - expected_data["gen_ai.request.messages"] = ( - '[{"role": "user", "content": "Hello!"}]' - ) - expected_data["gen_ai.response.text"] = ( - "[mocked] Hello! How can I help you today?" - ) + if send_default_pii and include_prompts: + expected_data["gen_ai.request.messages"] = ( + '[{"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 + 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["data"] == expected_data + assert span["data"] == expected_data -@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]) @pytest.mark.parametrize("include_prompts", [True, False]) def test_chat_completion_streaming( sentry_init: "Any", capture_events: "Any", - capture_items: "Any", send_default_pii: "Any", include_prompts: "Any", mock_hf_chat_completion_api_streaming: "Any", - stream_gen_ai_spans: "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, + stream_gen_ai_spans=False, ) client = get_hf_provider_inference_client() + events = capture_events() - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with sentry_sdk.start_transaction(name="test"): - _ = list( - client.chat_completion( - [ - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "Hello!"}, - ], - stream=True, - ) - ) - - spans = [item.payload for item in items if item.type == "span"] - span = None - for sp in spans: - if sp["attributes"]["sentry.op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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, - "process.runtime.name": mock.ANY, - "process.runtime.version": mock.ANY, - "sentry.environment": "production", - "sentry.op": "gen_ai.chat", - "sentry.origin": "auto.ai.huggingface_hub", - "sentry.release": mock.ANY, - "sentry.sdk.name": "sentry.python", - "sentry.sdk.version": mock.ANY, - "sentry.segment.id": mock.ANY, - "sentry.segment.name": "test", - "server.address": mock.ANY, - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } - # 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!"}, - ] + with sentry_sdk.start_transaction(name="test"): + _ = list( + client.chat_completion( + [{"role": "user", "content": "Hello!"}], + stream=True, ) - 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 + (transaction,) = events - assert span["attributes"] == expected_data - else: - events = capture_events() - - with sentry_sdk.start_transaction(name="test"): - _ = list( - client.chat_completion( - [{"role": "user", "content": "Hello!"}], - stream=True, - ) - ) + span = None + for sp in transaction["spans"]: + if sp["op"].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + else: + # 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["op"] == "http.client" + + assert span is not None + + assert span["op"] == "gen_ai.chat" + assert span["description"] == "chat test-model" + assert span["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, + "thread.id": mock.ANY, + "thread.name": mock.ANY, + } + # 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 - (transaction,) = events - - span = None - for sp in transaction["spans"]: - if sp["op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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["op"] == "http.client" - - assert span is not None - - assert span["op"] == "gen_ai.chat" - assert span["description"] == "chat test-model" - assert span["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, - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } - # 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"] = ( - '[{"role": "user", "content": "Hello!"}]' - ) - expected_data["gen_ai.response.text"] = "the mocked model response" + if send_default_pii and include_prompts: + expected_data["gen_ai.request.messages"] = ( + '[{"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 + 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["data"] == expected_data + assert span["data"] == expected_data -@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( sentry_init: "Any", capture_events: "Any", - capture_items: "Any", mock_hf_api_with_errors: "Any", - stream_gen_ai_spans: "Any", ) -> None: sentry_init( traces_sample_rate=1.0, - stream_gen_ai_spans=stream_gen_ai_spans, + stream_gen_ai_spans=False, ) client = get_hf_provider_inference_client() - if stream_gen_ai_spans: - items = capture_items("event", "transaction", "span") + events = capture_events() - with sentry_sdk.start_transaction(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"] + with sentry_sdk.start_transaction(name="test"), pytest.raises(HfHubHTTPError): + client.chat_completion( + messages=[{"role": "user", "content": "Hello!"}], + ) - spans = [item.payload for item in items if item.type == "span"] - span = None - for sp in spans: - if sp["attributes"]["sentry.op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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" + ( + error, + transaction, + ) = events - assert span is not None + assert error["exception"]["values"][0]["mechanism"]["type"] == "huggingface_hub" + assert not error["exception"]["values"][0]["mechanism"]["handled"] - 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" + span = None + for sp in transaction["spans"]: + if sp["op"].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + else: + # 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["op"] == "http.client" - (transaction,) = (item.payload for item in items if item.type == "transaction") + assert span is not None - assert ( - error["contexts"]["trace"]["trace_id"] - == transaction["contexts"]["trace"]["trace_id"] - ) - expected_data = { - "gen_ai.operation.name": "chat", - "gen_ai.request.model": "test-model", - "process.runtime.name": mock.ANY, - "process.runtime.version": mock.ANY, - "sentry.environment": "production", - "sentry.op": "gen_ai.chat", - "sentry.origin": "auto.ai.huggingface_hub", - "sentry.release": mock.ANY, - "sentry.sdk.name": "sentry.python", - "sentry.sdk.version": mock.ANY, - "sentry.segment.id": mock.ANY, - "sentry.segment.name": "test", - "server.address": mock.ANY, - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } - assert span["attributes"] == expected_data - else: - events = capture_events() + assert span["op"] == "gen_ai.chat" + assert span["description"] == "chat test-model" + assert span["origin"] == "auto.ai.huggingface_hub" + assert span["status"] == "internal_error" + assert span.get("tags", {}).get("status") == "internal_error" - with sentry_sdk.start_transaction(name="test"), pytest.raises(HfHubHTTPError): - client.chat_completion( - messages=[{"role": "user", "content": "Hello!"}], - ) - - ( - error, - transaction, - ) = events - - assert error["exception"]["values"][0]["mechanism"]["type"] == "huggingface_hub" - assert not error["exception"]["values"][0]["mechanism"]["handled"] - - span = None - for sp in transaction["spans"]: - if sp["op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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["op"] == "http.client" - - assert span is not None - - assert span["op"] == "gen_ai.chat" - assert span["description"] == "chat test-model" - assert span["origin"] == "auto.ai.huggingface_hub" - assert span["status"] == "internal_error" - assert span.get("tags", {}).get("status") == "internal_error" - - assert ( - error["contexts"]["trace"]["trace_id"] - == transaction["contexts"]["trace"]["trace_id"] - ) - expected_data = { - "gen_ai.operation.name": "chat", - "gen_ai.request.model": "test-model", - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } - assert span["data"] == expected_data + assert ( + error["contexts"]["trace"]["trace_id"] + == transaction["contexts"]["trace"]["trace_id"] + ) + expected_data = { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "test-model", + "thread.id": mock.ANY, + "thread.name": mock.ANY, + } + assert span["data"] == expected_data -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) def test_span_status_error( sentry_init: "Any", capture_events: "Any", - capture_items: "Any", mock_hf_api_with_errors: "Any", - stream_gen_ai_spans: "Any", ) -> None: client = get_hf_provider_inference_client() sentry_init( traces_sample_rate=1.0, - stream_gen_ai_spans=stream_gen_ai_spans, + stream_gen_ai_spans=False, ) + events = capture_events() - if stream_gen_ai_spans: - items = capture_items("event", "transaction", "span") - - with sentry_sdk.start_transaction(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" - - spans = [item.payload for item in items if item.type == "span"] - span = None - for sp in spans: - if sp["attributes"]["sentry.op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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" - else: - events = capture_events() - - with sentry_sdk.start_transaction(name="test"), pytest.raises(HfHubHTTPError): - client.chat_completion( - messages=[{"role": "user", "content": "Hello!"}], - ) + with sentry_sdk.start_transaction(name="test"), pytest.raises(HfHubHTTPError): + client.chat_completion( + messages=[{"role": "user", "content": "Hello!"}], + ) - (error, transaction) = events - assert error["level"] == "error" + (error, transaction) = events + assert error["level"] == "error" - span = None - for sp in transaction["spans"]: - if sp["op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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["op"] == "http.client" + span = None + for sp in transaction["spans"]: + if sp["op"].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + else: + # 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["op"] == "http.client" - assert span is not None - assert span["status"] == "internal_error" - assert span["tags"]["status"] == "internal_error" + assert span is not None + assert span["status"] == "internal_error" + assert span["tags"]["status"] == "internal_error" -@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]) @pytest.mark.parametrize("include_prompts", [True, False]) def test_chat_completion_with_tools( sentry_init: "Any", capture_events: "Any", - capture_items: "Any", send_default_pii: "Any", include_prompts: "Any", mock_hf_chat_completion_api_tools: "Any", - stream_gen_ai_spans: "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, + stream_gen_ai_spans=False, ) client = get_hf_provider_inference_client() @@ -1268,147 +880,77 @@ def test_chat_completion_with_tools( }, } ] + events = capture_events() - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with sentry_sdk.start_transaction(name="test"): - client.chat_completion( - messages=[{"role": "user", "content": "What is the weather in Paris?"}], - tools=tools, - tool_choice="auto", - ) - - spans = [item.payload for item in items if item.type == "span"] - span = None - for sp in spans: - if sp["attributes"]["sentry.op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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, - "process.runtime.name": mock.ANY, - "process.runtime.version": mock.ANY, - "sentry.environment": "production", - "sentry.op": "gen_ai.chat", - "sentry.origin": "auto.ai.huggingface_hub", - "sentry.release": mock.ANY, - "sentry.sdk.name": "sentry.python", - "sentry.sdk.version": mock.ANY, - "sentry.segment.id": mock.ANY, - "sentry.segment.name": "test", - "server.address": mock.ANY, - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } - - 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"] == expected_data - else: - events = capture_events() + with sentry_sdk.start_transaction(name="test"): + client.chat_completion( + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=tools, + tool_choice="auto", + ) - with sentry_sdk.start_transaction(name="test"): - client.chat_completion( - messages=[{"role": "user", "content": "What is the weather in Paris?"}], - tools=tools, - tool_choice="auto", - ) + (transaction,) = events - (transaction,) = events - - span = None - for sp in transaction["spans"]: - if sp["op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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["op"] == "http.client" - - assert span is not None - - assert span["op"] == "gen_ai.chat" - assert span["description"] == "chat test-model" - assert span["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, - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } + span = None + for sp in transaction["spans"]: + if sp["op"].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + else: + # 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["op"] == "http.client" + + assert span is not None + + assert span["op"] == "gen_ai.chat" + assert span["description"] == "chat test-model" + assert span["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, + "thread.id": mock.ANY, + "thread.name": mock.ANY, + } - 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 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 + 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["data"] == expected_data + assert span["data"] == expected_data -@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]) @pytest.mark.parametrize("include_prompts", [True, False]) def test_chat_completion_streaming_with_tools( sentry_init: "Any", capture_events: "Any", - capture_items: "Any", send_default_pii: "Any", include_prompts: "Any", mock_hf_chat_completion_api_streaming_tools: "Any", - stream_gen_ai_spans: "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, + stream_gen_ai_spans=False, ) client = get_hf_provider_inference_client() @@ -1427,145 +969,67 @@ def test_chat_completion_streaming_with_tools( }, } ] + events = capture_events() - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with sentry_sdk.start_transaction(name="test"): - _ = list( - client.chat_completion( - messages=[ - {"role": "user", "content": "What is the weather in Paris?"} - ], - stream=True, - tools=tools, - tool_choice="auto", - ) - ) - - spans = [item.payload for item in items if item.type == "span"] - span = None - for sp in spans: - if sp["attributes"]["sentry.op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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, - "process.runtime.name": mock.ANY, - "process.runtime.version": mock.ANY, - "sentry.environment": "production", - "sentry.op": "gen_ai.chat", - "sentry.origin": "auto.ai.huggingface_hub", - "sentry.release": mock.ANY, - "sentry.sdk.name": "sentry.python", - "sentry.sdk.version": mock.ANY, - "sentry.segment.id": mock.ANY, - "sentry.segment.name": "test", - "server.address": mock.ANY, - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } - - 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"}]' + with sentry_sdk.start_transaction(name="test"): + _ = list( + client.chat_completion( + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + stream=True, + tools=tools, + tool_choice="auto", ) + ) - 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"] == expected_data - else: - events = capture_events() - - with sentry_sdk.start_transaction(name="test"): - _ = list( - client.chat_completion( - messages=[ - {"role": "user", "content": "What is the weather in Paris?"} - ], - stream=True, - tools=tools, - tool_choice="auto", - ) - ) + (transaction,) = events - (transaction,) = events - - span = None - for sp in transaction["spans"]: - if sp["op"].startswith("gen_ai"): - assert span is None, "there is exactly one gen_ai span" - span = sp - else: - # 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["op"] == "http.client" - - assert span is not None - - assert span["op"] == "gen_ai.chat" - assert span["description"] == "chat test-model" - assert span["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, - "thread.id": mock.ANY, - "thread.name": mock.ANY, - } + span = None + for sp in transaction["spans"]: + if sp["op"].startswith("gen_ai"): + assert span is None, "there is exactly one gen_ai span" + span = sp + else: + # 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["op"] == "http.client" + + assert span is not None + + assert span["op"] == "gen_ai.chat" + assert span["description"] == "chat test-model" + assert span["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, + "thread.id": mock.ANY, + "thread.name": mock.ANY, + } - 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 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 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 + 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["data"] == expected_data + assert span["data"] == expected_data DATA_COLLECTION_TOOLS = [ @@ -1584,7 +1048,6 @@ def test_chat_completion_streaming_with_tools( ] -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize( "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", @@ -1673,20 +1136,18 @@ def test_chat_completion_streaming_with_tools( def test_text_generation_data_collection( sentry_init: "Any", capture_events: "Any", - capture_items: "Any", mock_hf_text_generation_api: "Any", data_collection: "Any", send_default_pii: "Any", include_prompts: "Any", expected_present: "Any", expected_absent: "Any", - stream_gen_ai_spans: "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, + stream_gen_ai_spans=False, ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -1695,25 +1156,13 @@ def test_text_generation_data_collection( client = InferenceClient(model="test-model") - captured = ( - capture_items("transaction", "span") - if stream_gen_ai_spans - else capture_events() - ) + captured = capture_events() with sentry_sdk.start_transaction(name="test"): client.text_generation("Hello", stream=False, details=True) - - if stream_gen_ai_spans: - 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") - ] - span_data = span["attributes"] - else: - (transaction,) = captured - (span,) = [sp for sp in transaction["spans"] if sp["op"].startswith("gen_ai")] - span_data = span["data"] + (transaction,) = captured + (span,) = [sp for sp in transaction["spans"] if sp["op"].startswith("gen_ai")] + span_data = span["data"] expected_values = { SPANDATA.GEN_AI_REQUEST_MESSAGES: "Hello", @@ -1734,7 +1183,6 @@ def test_text_generation_data_collection( assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 10 -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize( "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", @@ -1823,20 +1271,18 @@ def test_text_generation_data_collection( def test_text_generation_streaming_data_collection( sentry_init: "Any", capture_events: "Any", - capture_items: "Any", mock_hf_text_generation_api_streaming: "Any", data_collection: "Any", send_default_pii: "Any", include_prompts: "Any", expected_present: "Any", expected_absent: "Any", - stream_gen_ai_spans: "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, + stream_gen_ai_spans=False, ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -1845,26 +1291,14 @@ def test_text_generation_streaming_data_collection( client = InferenceClient(model="test-model") - captured = ( - capture_items("transaction", "span") - if stream_gen_ai_spans - else capture_events() - ) + captured = capture_events() with sentry_sdk.start_transaction(name="test"): for _ in client.text_generation(prompt="Hello", stream=True, details=True): pass - - if stream_gen_ai_spans: - 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") - ] - span_data = span["attributes"] - else: - (transaction,) = captured - (span,) = [sp for sp in transaction["spans"] if sp["op"].startswith("gen_ai")] - span_data = span["data"] + (transaction,) = captured + (span,) = [sp for sp in transaction["spans"] if sp["op"].startswith("gen_ai")] + span_data = span["data"] expected_values = { SPANDATA.GEN_AI_REQUEST_MESSAGES: "Hello", @@ -1886,7 +1320,6 @@ def test_text_generation_streaming_data_collection( assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 10 -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize( "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", @@ -1983,20 +1416,18 @@ def test_text_generation_streaming_data_collection( def test_chat_completion_data_collection_tools( sentry_init: "Any", capture_events: "Any", - capture_items: "Any", mock_hf_chat_completion_api_tools: "Any", data_collection: "Any", send_default_pii: "Any", include_prompts: "Any", expected_present: "Any", expected_absent: "Any", - stream_gen_ai_spans: "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, + stream_gen_ai_spans=False, ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -2005,11 +1436,7 @@ def test_chat_completion_data_collection_tools( client = get_hf_provider_inference_client() - captured = ( - capture_items("transaction", "span") - if stream_gen_ai_spans - else capture_events() - ) + captured = capture_events() with sentry_sdk.start_transaction(name="test"): client.chat_completion( @@ -2017,17 +1444,9 @@ def test_chat_completion_data_collection_tools( tools=DATA_COLLECTION_TOOLS, tool_choice="auto", ) - - if stream_gen_ai_spans: - 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") - ] - span_data = span["attributes"] - else: - (transaction,) = captured - (span,) = [sp for sp in transaction["spans"] if sp["op"].startswith("gen_ai")] - span_data = span["data"] + (transaction,) = captured + (span,) = [sp for sp in transaction["spans"] if sp["op"].startswith("gen_ai")] + span_data = span["data"] expected_values = { SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS: '[{"type": "function", "function": {"name": "get_weather", "description": "Get current weather", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}}}]', @@ -2055,7 +1474,6 @@ def test_chat_completion_data_collection_tools( assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 18 -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize( "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", @@ -2159,20 +1577,18 @@ def test_chat_completion_data_collection_tools( def test_chat_completion_streaming_data_collection_tools( sentry_init: "Any", capture_events: "Any", - capture_items: "Any", mock_hf_chat_completion_api_streaming_tools: "Any", data_collection: "Any", send_default_pii: "Any", include_prompts: "Any", expected_present: "Any", expected_absent: "Any", - stream_gen_ai_spans: "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, + stream_gen_ai_spans=False, ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -2181,11 +1597,7 @@ def test_chat_completion_streaming_data_collection_tools( client = get_hf_provider_inference_client() - captured = ( - capture_items("transaction", "span") - if stream_gen_ai_spans - else capture_events() - ) + captured = capture_events() with sentry_sdk.start_transaction(name="test"): for _ in client.chat_completion( @@ -2195,17 +1607,9 @@ def test_chat_completion_streaming_data_collection_tools( stream=True, ): pass - - if stream_gen_ai_spans: - 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") - ] - span_data = span["attributes"] - else: - (transaction,) = captured - (span,) = [sp for sp in transaction["spans"] if sp["op"].startswith("gen_ai")] - span_data = span["data"] + (transaction,) = captured + (span,) = [sp for sp in transaction["spans"] if sp["op"].startswith("gen_ai")] + span_data = span["data"] expected_values = { SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS: '[{"type": "function", "function": {"name": "get_weather", "description": "Get current weather", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}}}]',