diff --git a/sentry_sdk/integrations/huggingface_hub.py b/sentry_sdk/integrations/huggingface_hub.py index 2348f2b155..bb5b1df186 100644 --- a/sentry_sdk/integrations/huggingface_hub.py +++ b/sentry_sdk/integrations/huggingface_hub.py @@ -8,11 +8,9 @@ from sentry_sdk.ai.utils import set_data_normalized from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version -from sentry_sdk.scope import should_send_default_pii from sentry_sdk.utils import ( capture_internal_exceptions, event_from_exception, - has_data_collection_enabled, parse_version, reraise, ) @@ -113,20 +111,8 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": "stream": SPANDATA.GEN_AI_RESPONSE_STREAMING, } - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["inputs"]: - attribute_mapping["tools"] = SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS - else: - # Legacy behaviour where we unconditionally set this. Remove when data collection is fully rolled out + if client.options["data_collection"]["gen_ai"]["inputs"]: attribute_mapping["tools"] = SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS - - # Input attributes - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["inputs"]: - set_data_normalized( - span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompt, unpack=False - ) - elif should_send_default_pii(): set_data_normalized( span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompt, unpack=False ) @@ -206,15 +192,7 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": ) if tool_calls is not None and len(tool_calls) > 0: - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["outputs"]: - set_data_normalized( - span, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - tool_calls, - unpack=False, - ) - elif should_send_default_pii(): + if client.options["data_collection"]["gen_ai"]["outputs"]: set_data_normalized( span, SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, @@ -225,14 +203,7 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": if len(response_text_buffer) > 0: text_response = "".join(response_text_buffer) if text_response: - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["outputs"]: - set_data_normalized( - span, - SPANDATA.GEN_AI_RESPONSE_TEXT, - text_response, - ) - elif should_send_default_pii(): + if client.options["data_collection"]["gen_ai"]["outputs"]: set_data_normalized( span, SPANDATA.GEN_AI_RESPONSE_TEXT, @@ -295,14 +266,7 @@ def new_details_iterator() -> "Iterable[Any]": finish_reason, ) - should_set_response_text = False - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["outputs"]: - should_set_response_text = True - elif should_send_default_pii(): - should_set_response_text = True - - if should_set_response_text: + if client.options["data_collection"]["gen_ai"]["outputs"]: if len(response_text_buffer) > 0: text_response = "".join(response_text_buffer) if text_response: @@ -381,43 +345,29 @@ def new_iterator() -> "Iterable[ChatCompletionStreamOutput]": finish_reason, ) - if tool_calls is not None and len(tool_calls) > 0: - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"][ - "outputs" - ]: - set_data_normalized( - span, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - tool_calls, - unpack=False, - ) - elif should_send_default_pii(): - set_data_normalized( - span, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - tool_calls, - unpack=False, - ) + if ( + tool_calls is not None + and len(tool_calls) > 0 + and client.options["data_collection"]["gen_ai"]["outputs"] + ): + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + tool_calls, + unpack=False, + ) - if len(response_text_buffer) > 0: + if ( + len(response_text_buffer) > 0 + and client.options["data_collection"]["gen_ai"]["outputs"] + ): text_response = "".join(response_text_buffer) if text_response: - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"][ - "outputs" - ]: - set_data_normalized( - span, - SPANDATA.GEN_AI_RESPONSE_TEXT, - text_response, - ) - elif should_send_default_pii(): - set_data_normalized( - span, - SPANDATA.GEN_AI_RESPONSE_TEXT, - text_response, - ) + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TEXT, + text_response, + ) if usage is not None: record_token_usage( diff --git a/tests/integrations/huggingface_hub/test_huggingface_hub.py b/tests/integrations/huggingface_hub/test_huggingface_hub.py index fc12ab3770..734bd2af84 100644 --- a/tests/integrations/huggingface_hub/test_huggingface_hub.py +++ b/tests/integrations/huggingface_hub/test_huggingface_hub.py @@ -468,16 +468,19 @@ def mock_hf_chat_completion_api_streaming_tools(httpx_mock): @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) -@pytest.mark.parametrize("send_default_pii", [True, False]) def test_text_generation( sentry_init: "Any", capture_items: "Any", - send_default_pii: "Any", mock_hf_text_generation_api: "Any", ) -> None: sentry_init( traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, integrations=[HuggingfaceHubIntegration()], ) @@ -520,13 +523,73 @@ def test_text_generation( "sentry.origin": "auto.ai.huggingface_hub", } - if send_default_pii: - expected_data["gen_ai.request.messages"] = "Hello" - expected_data["gen_ai.response.text"] = "[mocked] Hello! How can i help you?" + expected_data["gen_ai.request.messages"] = "Hello" + expected_data["gen_ai.response.text"] = "[mocked] Hello! How can i help you?" + + assert span["attributes"] == ApproxDict(expected_data) + + # text generation does not set the response model + assert "gen_ai.response.model" not in span["attributes"] + + +@pytest.mark.httpx_mock(assert_all_requests_were_expected=False) +def test_text_generation_no_sensitive_data( + sentry_init: "Any", + capture_items: "Any", + mock_hf_text_generation_api: "Any", +) -> None: + sentry_init( + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, + integrations=[HuggingfaceHubIntegration()], + ) + + client = InferenceClient(model="test-model") + items = capture_items("span") + + client.text_generation( + "Hello", + stream=False, + details=True, + ) + + sentry_sdk.flush() + spans = [item.payload for item in items] + 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 not send_default_pii: - assert "gen_ai.request.messages" not in expected_data - assert "gen_ai.response.text" not in expected_data + assert "gen_ai.request.messages" not in expected_data + assert "gen_ai.response.text" not in expected_data assert span["attributes"] == ApproxDict(expected_data) @@ -535,16 +598,19 @@ def test_text_generation( @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) -@pytest.mark.parametrize("send_default_pii", [True, False]) def test_text_generation_streaming( sentry_init: "Any", capture_items: "Any", - send_default_pii: "Any", mock_hf_text_generation_api_streaming: "Any", ) -> None: sentry_init( traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, integrations=[HuggingfaceHubIntegration()], ) @@ -589,13 +655,75 @@ def test_text_generation_streaming( "sentry.origin": "auto.ai.huggingface_hub", } - if send_default_pii: - expected_data["gen_ai.request.messages"] = "Hello" - expected_data["gen_ai.response.text"] = "the mocked model response" + expected_data["gen_ai.request.messages"] = "Hello" + expected_data["gen_ai.response.text"] = "the mocked model response" + + assert span["attributes"] == ApproxDict(expected_data) + + # text generation does not set the response model + assert "gen_ai.response.model" not in span["attributes"] + + +@pytest.mark.httpx_mock(assert_all_requests_were_expected=False) +def test_text_generation_streaming_no_sensitive_data( + sentry_init: "Any", + capture_items: "Any", + mock_hf_text_generation_api_streaming: "Any", +) -> None: + sentry_init( + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, + integrations=[HuggingfaceHubIntegration()], + ) + + client = InferenceClient(model="test-model") + items = capture_items("span") + + for _ in client.text_generation( + prompt="Hello", + stream=True, + details=True, + ): + pass + + sentry_sdk.flush() + spans = [item.payload for item in items] + 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 - if not send_default_pii: - assert "gen_ai.request.messages" not in expected_data - assert "gen_ai.response.text" not in expected_data + 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", + } + + assert "gen_ai.request.messages" not in expected_data + assert "gen_ai.response.text" not in expected_data assert span["attributes"] == ApproxDict(expected_data) @@ -604,16 +732,19 @@ def test_text_generation_streaming( @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) -@pytest.mark.parametrize("send_default_pii", [True, False]) def test_chat_completion( sentry_init: "Any", capture_items: "Any", - send_default_pii: "Any", mock_hf_chat_completion_api: "Any", ) -> None: sentry_init( traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, integrations=[HuggingfaceHubIntegration()], ) @@ -665,38 +796,105 @@ def test_chat_completion( "sentry.origin": "auto.ai.huggingface_hub", } - if send_default_pii: - 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?" - ) + 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?" + + assert span["attributes"] == ApproxDict(expected_data) + + +@pytest.mark.httpx_mock(assert_all_requests_were_expected=False) +def test_chat_completion_no_sensitive_data( + sentry_init: "Any", + capture_items: "Any", + mock_hf_chat_completion_api: "Any", +) -> None: + sentry_init( + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, + integrations=[HuggingfaceHubIntegration()], + ) + + client = get_hf_provider_inference_client() + items = capture_items("span") + + 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] + 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 not send_default_pii: - assert "gen_ai.request.messages" not in expected_data - assert "gen_ai.response.text" not in expected_data + assert "gen_ai.request.messages" not in expected_data + assert "gen_ai.response.text" not in expected_data assert span["attributes"] == ApproxDict(expected_data) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) -@pytest.mark.parametrize("send_default_pii", [True, False]) def test_chat_completion_streaming( sentry_init: "Any", capture_items: "Any", - send_default_pii: "Any", mock_hf_chat_completion_api_streaming: "Any", ) -> None: sentry_init( traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, integrations=[HuggingfaceHubIntegration()], ) @@ -752,21 +950,91 @@ def test_chat_completion_streaming( expected_data["gen_ai.usage.output_tokens"] = 14 expected_data["gen_ai.usage.total_tokens"] = 197 - if send_default_pii: - expected_data["gen_ai.request.messages"] = safe_serialize( + 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" + + assert span["attributes"] == ApproxDict(expected_data) + + +@pytest.mark.httpx_mock(assert_all_requests_were_expected=False) +def test_chat_completion_streaming_no_sensitive_data( + sentry_init: "Any", + capture_items: "Any", + mock_hf_chat_completion_api_streaming: "Any", +) -> None: + sentry_init( + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, + integrations=[HuggingfaceHubIntegration()], + ) + + client = get_hf_provider_inference_client() + items = capture_items("span") + + _ = list( + client.chat_completion( [ { "role": "user", "content": "Message demonstrating the absence of truncation.", }, {"role": "user", "content": "Hello!"}, - ] + ], + stream=True, ) - expected_data["gen_ai.response.text"] = "the mocked model response" + ) + + sentry_sdk.flush() + spans = [item.payload for item in items] + 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" - if not send_default_pii: - assert "gen_ai.request.messages" not in expected_data - assert "gen_ai.response.text" not in expected_data + 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 + + assert "gen_ai.request.messages" not in expected_data + assert "gen_ai.response.text" not in expected_data assert span["attributes"] == ApproxDict(expected_data) @@ -779,6 +1047,7 @@ def test_chat_completion_api_error( ) -> None: sentry_init( traces_sample_rate=1.0, + data_collection={}, ) client = get_hf_provider_inference_client() @@ -797,7 +1066,120 @@ def test_chat_completion_api_error( assert not error["exception"]["values"][-1]["mechanism"]["handled"] sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] + 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" + + assert error["contexts"]["trace"]["trace_id"] == 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) + + +@pytest.mark.httpx_mock(assert_all_requests_were_expected=False) +def test_span_status_error( + sentry_init: "Any", + capture_items: "Any", + mock_hf_api_with_errors: "Any", +) -> None: + client = get_hf_provider_inference_client() + + sentry_init( + traces_sample_rate=1.0, + data_collection={}, + ) + items = capture_items("event", "span") + + with 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" + + +@pytest.mark.httpx_mock(assert_all_requests_were_expected=False) +def test_chat_completion_with_tools( + sentry_init: "Any", + capture_items: "Any", + mock_hf_chat_completion_api_tools: "Any", +): + sentry_init( + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, + integrations=[HuggingfaceHubIntegration()], + ) + + client = get_hf_provider_inference_client() + + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + ] + items = capture_items("span") + + 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] span = None for sp in spans: if "sentry.op" in sp["attributes"] and sp["attributes"]["sentry.op"].startswith( @@ -815,42 +1197,74 @@ def test_chat_completion_api_error( 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" - assert error["contexts"]["trace"]["trace_id"] == span["trace_id"] 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", } + + 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"}]' + ) + assert span["attributes"] == ApproxDict(expected_data) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) -def test_span_status_error( +def test_chat_completion_with_tools_no_sensitive_data( sentry_init: "Any", capture_items: "Any", - mock_hf_api_with_errors: "Any", -) -> None: - client = get_hf_provider_inference_client() - + mock_hf_chat_completion_api_tools: "Any", +): sentry_init( traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, + integrations=[HuggingfaceHubIntegration()], ) - items = capture_items("event", "span") - with pytest.raises(HfHubHTTPError): - client.chat_completion( - messages=[{"role": "user", "content": "Hello!"}], - ) + client = get_hf_provider_inference_client() - (error,) = [item.payload for item in items if item.type == "event"] - assert error["level"] == "error" + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + ] + items = capture_items("span") + + 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"] + spans = [item.payload for item in items] span = None for sp in spans: if "sentry.op" in sp["attributes"] and sp["attributes"]["sentry.op"].startswith( @@ -864,20 +1278,46 @@ def test_span_status_error( assert sp["attributes"]["sentry.op"] == "http.client" assert span is not None - assert span["status"] == "error" + + 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": "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", + } + + 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 "gen_ai.request.available_tools" not in span["attributes"] + + assert span["attributes"] == ApproxDict(expected_data) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) -@pytest.mark.parametrize("send_default_pii", [True, False]) -def test_chat_completion_with_tools( +def test_chat_completion_streaming_with_tools( sentry_init: "Any", capture_items: "Any", - send_default_pii: "Any", - mock_hf_chat_completion_api_tools: "Any", -): + mock_hf_chat_completion_api_streaming_tools: "Any", +) -> None: sentry_init( traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, integrations=[HuggingfaceHubIntegration()], ) @@ -899,10 +1339,13 @@ def test_chat_completion_with_tools( ] items = capture_items("span") - client.chat_completion( - messages=[{"role": "user", "content": "What is the weather in Paris?"}], - tools=tools, - tool_choice="auto", + _ = list( + client.chat_completion( + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + stream=True, + tools=tools, + tool_choice="auto", + ) ) sentry_sdk.flush() @@ -931,41 +1374,42 @@ def test_chat_completion_with_tools( "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, + "gen_ai.response.streaming": True, "sentry.environment": "production", "sentry.op": "gen_ai.chat", "sentry.origin": "auto.ai.huggingface_hub", } - if send_default_pii: - 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 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 not send_default_pii: - 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 + 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"}]' + ) assert span["attributes"] == ApproxDict(expected_data) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) -@pytest.mark.parametrize("send_default_pii", [True, False]) -def test_chat_completion_streaming_with_tools( +def test_chat_completion_streaming_with_tools_no_sensitive_data( sentry_init: "Any", capture_items: "Any", - send_default_pii: "Any", mock_hf_chat_completion_api_streaming_tools: "Any", ) -> None: sentry_init( traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, integrations=[HuggingfaceHubIntegration()], ) @@ -1018,7 +1462,6 @@ def test_chat_completion_streaming_with_tools( 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", @@ -1033,19 +1476,10 @@ def test_chat_completion_streaming_with_tools( expected_data["gen_ai.usage.output_tokens"] = 14 expected_data["gen_ai.usage.total_tokens"] = 197 - if send_default_pii: - 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: - 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 "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 "gen_ai.request.available_tools" not in span["attributes"] assert span["attributes"] == ApproxDict(expected_data) @@ -1068,31 +1502,10 @@ def test_chat_completion_streaming_with_tools( @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", + "data_collection,expected_present,expected_absent", [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": True}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": False}}, - True, - [], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", - ), pytest.param( {"gen_ai": {"inputs": True, "outputs": False}}, - False, [ SPANDATA.GEN_AI_REQUEST_MESSAGES, ], @@ -1103,7 +1516,6 @@ def test_chat_completion_streaming_with_tools( ), pytest.param( {"gen_ai": {"inputs": False, "outputs": True}}, - False, [ SPANDATA.GEN_AI_RESPONSE_TEXT, ], @@ -1112,36 +1524,6 @@ def test_chat_completion_streaming_with_tools( ], id="gen-ai-outputs-enabled-inputs-disabled", ), - pytest.param( - {"gen_ai": {}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", - ), - pytest.param( - None, - True, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="no-gen-ai-config-legacy-pii", - ), - pytest.param( - None, - False, - [], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="no-gen-ai-config-legacy-pii-disabled", - ), ], ) def test_text_generation_data_collection( @@ -1149,13 +1531,11 @@ def test_text_generation_data_collection( capture_items: "Any", mock_hf_text_generation_api: "Any", data_collection: "Any", - send_default_pii: "Any", expected_present: "Any", expected_absent: "Any", ) -> None: sentry_init_kwargs = dict( traces_sample_rate=1.0, - send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration()], ) if data_collection is not None: @@ -1185,7 +1565,6 @@ def test_text_generation_data_collection( for key in expected_absent: assert key not in span_data, f"{key} should not have been collected" - # Data collection never gates non-PII attributes assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "text_completion" assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "test-model" assert span_data[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "length" @@ -1194,31 +1573,10 @@ def test_text_generation_data_collection( @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", + "data_collection,expected_present,expected_absent", [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": True}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": False}}, - True, - [], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", - ), pytest.param( {"gen_ai": {"inputs": True, "outputs": False}}, - False, [ SPANDATA.GEN_AI_REQUEST_MESSAGES, ], @@ -1229,7 +1587,6 @@ def test_text_generation_data_collection( ), pytest.param( {"gen_ai": {"inputs": False, "outputs": True}}, - False, [ SPANDATA.GEN_AI_RESPONSE_TEXT, ], @@ -1238,36 +1595,6 @@ def test_text_generation_data_collection( ], id="gen-ai-outputs-enabled-inputs-disabled", ), - pytest.param( - {"gen_ai": {}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", - ), - pytest.param( - None, - True, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="no-gen-ai-config-legacy-pii", - ), - pytest.param( - None, - False, - [], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="no-gen-ai-config-legacy-pii-disabled", - ), ], ) def test_text_generation_streaming_data_collection( @@ -1275,13 +1602,11 @@ def test_text_generation_streaming_data_collection( capture_items: "Any", mock_hf_text_generation_api_streaming: "Any", data_collection: "Any", - send_default_pii: "Any", expected_present: "Any", expected_absent: "Any", ) -> None: sentry_init_kwargs = dict( traces_sample_rate=1.0, - send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration()], ) if data_collection is not None: @@ -1312,7 +1637,6 @@ def test_text_generation_streaming_data_collection( for key in expected_absent: assert key not in span_data, f"{key} should not have been collected" - # Data collection never gates non-PII attributes assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "text_completion" assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "test-model" assert span_data[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "length" @@ -1322,33 +1646,10 @@ def test_text_generation_streaming_data_collection( @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", + "data_collection,expected_present,expected_absent", [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": True}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - ], - [], - id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": False}}, - True, - [], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - ], - id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", - ), pytest.param( {"gen_ai": {"inputs": True, "outputs": False}}, - False, [ SPANDATA.GEN_AI_REQUEST_MESSAGES, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, @@ -1360,7 +1661,6 @@ def test_text_generation_streaming_data_collection( ), pytest.param( {"gen_ai": {"inputs": False, "outputs": True}}, - False, [ SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, ], @@ -1370,40 +1670,6 @@ def test_text_generation_streaming_data_collection( ], id="gen-ai-outputs-enabled-inputs-disabled", ), - pytest.param( - {"gen_ai": {}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - ], - [], - id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", - ), - pytest.param( - None, - True, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - ], - [], - id="no-gen-ai-config-legacy-pii", - ), - pytest.param( - None, - False, - [ - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - ], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - ], - id="no-gen-ai-config-legacy-pii-disabled", - ), ], ) def test_chat_completion_data_collection_tools( @@ -1411,13 +1677,11 @@ def test_chat_completion_data_collection_tools( capture_items: "Any", mock_hf_chat_completion_api_tools: "Any", data_collection: "Any", - send_default_pii: "Any", expected_present: "Any", expected_absent: "Any", ) -> None: sentry_init_kwargs = dict( traces_sample_rate=1.0, - send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration()], ) if data_collection is not None: @@ -1455,7 +1719,6 @@ def test_chat_completion_data_collection_tools( # This response carries only tool calls, so there is never any response text assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span_data - # Data collection never gates non-PII attributes assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "chat" assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "test-model" assert span_data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "test-model-123" @@ -1467,35 +1730,10 @@ def test_chat_completion_data_collection_tools( @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", + "data_collection,expected_present,expected_absent", [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": True}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - ], - [], - id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": False}}, - True, - [], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - ], - id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", - ), pytest.param( {"gen_ai": {"inputs": True, "outputs": False}}, - False, [ SPANDATA.GEN_AI_REQUEST_MESSAGES, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, @@ -1508,7 +1746,6 @@ def test_chat_completion_data_collection_tools( ), pytest.param( {"gen_ai": {"inputs": False, "outputs": True}}, - False, [ SPANDATA.GEN_AI_RESPONSE_TEXT, SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, @@ -1519,43 +1756,6 @@ def test_chat_completion_data_collection_tools( ], id="gen-ai-outputs-enabled-inputs-disabled", ), - pytest.param( - {"gen_ai": {}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - ], - [], - id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", - ), - pytest.param( - None, - True, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - ], - [], - id="no-gen-ai-config-legacy-pii", - ), - pytest.param( - None, - False, - [ - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - ], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="no-gen-ai-config-legacy-pii-disabled", - ), ], ) def test_chat_completion_streaming_data_collection_tools( @@ -1563,13 +1763,11 @@ def test_chat_completion_streaming_data_collection_tools( capture_items: "Any", mock_hf_chat_completion_api_streaming_tools: "Any", data_collection: "Any", - send_default_pii: "Any", expected_present: "Any", expected_absent: "Any", ) -> None: sentry_init_kwargs = dict( traces_sample_rate=1.0, - send_default_pii=send_default_pii, integrations=[HuggingfaceHubIntegration()], ) if data_collection is not None: @@ -1607,7 +1805,6 @@ def test_chat_completion_streaming_data_collection_tools( for key in expected_absent: assert key not in span_data, f"{key} should not have been collected" - # Data collection never gates non-PII attributes assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "chat" assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "test-model" assert span_data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "test-model-123"