diff --git a/sentry_sdk/integrations/huggingface_hub.py b/sentry_sdk/integrations/huggingface_hub.py index f93d90e2cb..fa6dae5214 100644 --- a/sentry_sdk/integrations/huggingface_hub.py +++ b/sentry_sdk/integrations/huggingface_hub.py @@ -18,6 +18,7 @@ from sentry_sdk.utils import ( capture_internal_exceptions, event_from_exception, + has_data_collection_enabled, reraise, ) @@ -74,7 +75,8 @@ def _capture_exception(exc: "Any") -> None: def _wrap_huggingface_task(f: "Callable[..., Any]", op: str) -> "Callable[..., Any]": @wraps(f) def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": - integration = sentry_sdk.get_client().get_integration(HuggingfaceHubIntegration) + client = sentry_sdk.get_client() + integration = client.get_integration(HuggingfaceHubIntegration) if integration is None: return f(*args, **kwargs) @@ -91,12 +93,12 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": # invalid call, dont instrument, let it return error return f(*args, **kwargs) - client = args[0] - model = client.model or kwargs.get("model") or "" + hf_client = args[0] + model = hf_client.model or kwargs.get("model") or "" operation_name = op.split(".")[-1] span: "Union[Span, StreamedSpan]" - if has_span_streaming_enabled(sentry_sdk.get_client().options): + if has_span_streaming_enabled(client.options): span = sentry_sdk.traces.start_span( name=f"{operation_name} {model}", attributes={ @@ -117,14 +119,7 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": if model: _set_span_data_attribute(span, SPANDATA.GEN_AI_REQUEST_MODEL, model) - # Input attributes - if should_send_default_pii() and integration.include_prompts: - set_data_normalized( - span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompt, unpack=False - ) - attribute_mapping = { - "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, @@ -134,6 +129,24 @@ 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 + 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() and integration.include_prompts: + set_data_normalized( + span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompt, unpack=False + ) + for attribute, span_attribute in attribute_mapping.items(): value = kwargs.get(attribute, None) if value is not None: @@ -210,8 +223,16 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": finish_reason, ) - if should_send_default_pii() and integration.include_prompts: - if tool_calls is not None and len(tool_calls) > 0: + 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"]["inputs"]: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + tool_calls, + unpack=False, + ) + elif should_send_default_pii() and integration.include_prompts: set_data_normalized( span, SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, @@ -219,9 +240,17 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": unpack=False, ) - if len(response_text_buffer) > 0: - text_response = "".join(response_text_buffer) - if text_response: + 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() and integration.include_prompts: set_data_normalized( span, SPANDATA.GEN_AI_RESPONSE_TEXT, @@ -284,7 +313,14 @@ def new_details_iterator() -> "Iterable[Any]": finish_reason, ) - if should_send_default_pii() and integration.include_prompts: + 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() and integration.include_prompts: + should_set_response_text = True + + if should_set_response_text: if len(response_text_buffer) > 0: text_response = "".join(response_text_buffer) if text_response: @@ -363,8 +399,21 @@ def new_iterator() -> "Iterable[ChatCompletionStreamOutput]": finish_reason, ) - if should_send_default_pii() and integration.include_prompts: - if tool_calls is not None and len(tool_calls) > 0: + 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"][ + "inputs" + ]: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + tool_calls, + unpack=False, + ) + elif ( + should_send_default_pii() + and integration.include_prompts + ): set_data_normalized( span, SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, @@ -372,9 +421,22 @@ def new_iterator() -> "Iterable[ChatCompletionStreamOutput]": unpack=False, ) - if len(response_text_buffer) > 0: - text_response = "".join(response_text_buffer) - if text_response: + 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() + and integration.include_prompts + ): set_data_normalized( span, SPANDATA.GEN_AI_RESPONSE_TEXT, diff --git a/tests/integrations/huggingface_hub/test_huggingface_hub.py b/tests/integrations/huggingface_hub/test_huggingface_hub.py index 91e0909731..9e7a77f66e 100644 --- a/tests/integrations/huggingface_hub/test_huggingface_hub.py +++ b/tests/integrations/huggingface_hub/test_huggingface_hub.py @@ -7,6 +7,7 @@ from huggingface_hub import InferenceClient 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 @@ -1565,3 +1566,667 @@ def test_chat_completion_streaming_with_tools( assert "gen_ai.response.tool_calls" not in expected_data assert span["data"] == expected_data + + +DATA_COLLECTION_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } +] + + +@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", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + 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, + 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, + False, + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + ], + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + False, + False, + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + ], + id="gen-ai-outputs-enabled-inputs-disabled", + ), + pytest.param( + {"gen_ai": {}}, + False, + 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, + True, + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + [], + id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + [], + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + id="no-gen-ai-config-legacy-pii-disabled", + ), + ], +) +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, + ) + if data_collection is not None: + sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**sentry_init_kwargs) + + client = InferenceClient(model="test-model") + + captured = ( + capture_items("transaction", "span") + if stream_gen_ai_spans + else capture_events() + ) + + with sentry_sdk.start_transaction(name="test"): + client.text_generation("Hello", stream=False, details=True) + + if stream_gen_ai_spans: + 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"] + + expected_values = { + SPANDATA.GEN_AI_REQUEST_MESSAGES: "Hello", + SPANDATA.GEN_AI_RESPONSE_TEXT: "[mocked] Hello! How can i help you?", + } + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + assert span_data[key] == expected_values[key] + + 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" + 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", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + 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, + 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, + False, + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + ], + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + False, + False, + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + ], + id="gen-ai-outputs-enabled-inputs-disabled", + ), + pytest.param( + {"gen_ai": {}}, + False, + 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, + True, + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + [], + id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + [], + [ + 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( + 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, + ) + if data_collection is not None: + sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**sentry_init_kwargs) + + client = InferenceClient(model="test-model") + + captured = ( + capture_items("transaction", "span") + if stream_gen_ai_spans + else 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"] + + expected_values = { + SPANDATA.GEN_AI_REQUEST_MESSAGES: "Hello", + SPANDATA.GEN_AI_RESPONSE_TEXT: "the mocked model response", + } + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + assert span_data[key] == expected_values[key] + + 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" + assert span_data[SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + 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", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + 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, + 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, + False, + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, + ], + [], + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + False, + False, + [], + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, + ], + id="gen-ai-outputs-enabled-inputs-disabled", + ), + pytest.param( + {"gen_ai": {}}, + False, + 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, + 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-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + [ + 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( + 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, + ) + if data_collection is not None: + sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**sentry_init_kwargs) + + client = get_hf_provider_inference_client() + + captured = ( + capture_items("transaction", "span") + if stream_gen_ai_spans + else capture_events() + ) + + with sentry_sdk.start_transaction(name="test"): + client.chat_completion( + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + 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"] + + 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"]}}}]', + SPANDATA.GEN_AI_REQUEST_MESSAGES: '[{"role": "user", "content": "What is the weather in Paris?"}]', + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: '[{"function": {"arguments": {"location": "Paris"}, "name": "get_weather", "description": "None"}, "id": "call_123", "type": "function"}]', + } + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + assert span_data[key] == expected_values[key] + + for key in expected_absent: + assert key not in span_data, f"{key} should not have been collected" + + # 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" + assert span_data[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "tool_calls" + assert span_data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert span_data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 8 + 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", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + 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, + 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, + False, + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, + ], + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + False, + False, + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, + ], + id="gen-ai-outputs-enabled-inputs-disabled", + ), + pytest.param( + {"gen_ai": {}}, + False, + 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, + 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-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + [ + 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( + 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, + ) + if data_collection is not None: + sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**sentry_init_kwargs) + + client = get_hf_provider_inference_client() + + captured = ( + capture_items("transaction", "span") + if stream_gen_ai_spans + else capture_events() + ) + + with sentry_sdk.start_transaction(name="test"): + for _ in client.chat_completion( + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=DATA_COLLECTION_TOOLS, + tool_choice="auto", + 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"] + + 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"]}}}]', + SPANDATA.GEN_AI_REQUEST_MESSAGES: '[{"role": "user", "content": "What is the weather in Paris?"}]', + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS: '[{"function": {"arguments": {"location": "Paris"}, "name": "get_weather"}, "id": "call_123", "type": "function", "index": "None"}]', + SPANDATA.GEN_AI_RESPONSE_TEXT: "response with tool calls follows", + } + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + assert span_data[key] == expected_values[key] + + 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" + assert span_data[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "tool_calls" + assert span_data[SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + + if HF_VERSION and HF_VERSION >= (0, 26, 0): + assert span_data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 183 + assert span_data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 14 + assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 197