From 2ec0d48b8156e92a789267aa40a21435700d4a7a Mon Sep 17 00:00:00 2001 From: Erica Pisani Date: Tue, 11 Aug 2026 13:09:31 -0400 Subject: [PATCH 1/5] feat(huggingface_hub): Gate prompt/response collection on data_collection option Prompts, response text, tool calls, and available tools are now gated on the `data_collection.gen_ai.inputs`/`outputs` experiment when configured, falling back to the existing `send_default_pii`/`include_prompts` check otherwise. Refs PY-2588 Refs #6748 --- sentry_sdk/integrations/huggingface_hub.py | 106 ++++- .../huggingface_hub/test_huggingface_hub.py | 418 ++++++++++++++++++ 2 files changed, 502 insertions(+), 22 deletions(-) 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..83b567dd0b 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,420 @@ def test_chat_completion_streaming_with_tools( assert "gen_ai.response.tool_calls" not in expected_data assert span["data"] == expected_data + + +DATA_COLLECTION_TEXT_GENERATION_EXPECTED_VALUES = { + SPANDATA.GEN_AI_REQUEST_MESSAGES: "Hello", + SPANDATA.GEN_AI_RESPONSE_TEXT: "[mocked] Hello! How can i help you?", +} + + +def _get_gen_ai_span_data(captured: "Any", stream_gen_ai_spans: "Any") -> "Any": + 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") + ] + return span["attributes"] + + (transaction,) = captured + (span,) = [sp for sp in transaction["spans"] if sp["op"].startswith("gen_ai")] + return span["data"] + + +DATA_COLLECTION_PARAMS = [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + False, + True, + True, + id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": False}}, + True, + True, + False, + False, + id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", + ), + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + False, + False, + True, + False, + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + False, + False, + False, + True, + id="gen-ai-outputs-enabled-inputs-disabled", + ), + pytest.param( + {"gen_ai": {}}, + False, + False, + True, + True, + id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", + ), + pytest.param( + None, + True, + True, + True, + True, + id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + False, + False, + id="no-gen-ai-config-legacy-pii-disabled", + ), +] + +DATA_COLLECTION_PARAM_NAMES = ( + "data_collection,send_default_pii,include_prompts,collect_inputs,collect_outputs" +) + + +def _expected_keys( + collect_inputs: "Any", collect_outputs: "Any", input_keys: "Any", output_keys: "Any" +) -> "Any": + present = (input_keys if collect_inputs else []) + ( + output_keys if collect_outputs else [] + ) + absent = ([] if collect_inputs else input_keys) + ( + [] if collect_outputs else output_keys + ) + return present, absent + + +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.httpx_mock(assert_all_requests_were_expected=False) +@pytest.mark.parametrize(DATA_COLLECTION_PARAM_NAMES, DATA_COLLECTION_PARAMS) +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", + collect_inputs: "Any", + collect_outputs: "Any", + stream_gen_ai_spans: "Any", +) -> None: + expected_present, expected_absent = _expected_keys( + collect_inputs, + collect_outputs, + [SPANDATA.GEN_AI_REQUEST_MESSAGES], + [SPANDATA.GEN_AI_RESPONSE_TEXT], + ) + + 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) + + span_data = _get_gen_ai_span_data(captured, stream_gen_ai_spans) + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + assert span_data[key] == DATA_COLLECTION_TEXT_GENERATION_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 + + +DATA_COLLECTION_TEXT_GENERATION_STREAMING_EXPECTED_VALUES = { + SPANDATA.GEN_AI_REQUEST_MESSAGES: "Hello", + SPANDATA.GEN_AI_RESPONSE_TEXT: "the mocked model response", +} + + +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.httpx_mock(assert_all_requests_were_expected=False) +@pytest.mark.parametrize(DATA_COLLECTION_PARAM_NAMES, DATA_COLLECTION_PARAMS) +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", + collect_inputs: "Any", + collect_outputs: "Any", + stream_gen_ai_spans: "Any", +) -> None: + expected_present, expected_absent = _expected_keys( + collect_inputs, + collect_outputs, + [SPANDATA.GEN_AI_REQUEST_MESSAGES], + [SPANDATA.GEN_AI_RESPONSE_TEXT], + ) + + 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 + + span_data = _get_gen_ai_span_data(captured, stream_gen_ai_spans) + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + assert ( + span_data[key] + == DATA_COLLECTION_TEXT_GENERATION_STREAMING_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 + + +DATA_COLLECTION_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } +] + +DATA_COLLECTION_CHAT_COMPLETION_TOOLS_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"}]', +} + + +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.httpx_mock(assert_all_requests_were_expected=False) +@pytest.mark.parametrize(DATA_COLLECTION_PARAM_NAMES, DATA_COLLECTION_PARAMS) +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", + collect_inputs: "Any", + collect_outputs: "Any", + stream_gen_ai_spans: "Any", +) -> None: + expected_present, expected_absent = _expected_keys( + collect_inputs, + collect_outputs, + [SPANDATA.GEN_AI_REQUEST_MESSAGES, SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS], + [], + ) + + # Legacy behaviour sets the available tools unconditionally, data collection + # gates them on inputs + if data_collection is None or collect_inputs: + expected_present.append(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS) + else: + expected_absent.append(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS) + + 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", + ) + + span_data = _get_gen_ai_span_data(captured, stream_gen_ai_spans) + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + assert ( + span_data[key] == DATA_COLLECTION_CHAT_COMPLETION_TOOLS_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 + + +DATA_COLLECTION_CHAT_COMPLETION_STREAMING_TOOLS_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", +} + + +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.httpx_mock(assert_all_requests_were_expected=False) +@pytest.mark.parametrize(DATA_COLLECTION_PARAM_NAMES, DATA_COLLECTION_PARAMS) +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", + collect_inputs: "Any", + collect_outputs: "Any", + stream_gen_ai_spans: "Any", +) -> None: + expected_present, expected_absent = _expected_keys( + collect_inputs, + collect_outputs, + [SPANDATA.GEN_AI_REQUEST_MESSAGES, SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS], + [SPANDATA.GEN_AI_RESPONSE_TEXT], + ) + + # Legacy behaviour sets the available tools unconditionally, data collection + # gates them on inputs + if data_collection is None or collect_inputs: + expected_present.append(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS) + else: + expected_absent.append(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS) + + 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 + + span_data = _get_gen_ai_span_data(captured, stream_gen_ai_spans) + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + assert ( + span_data[key] + == DATA_COLLECTION_CHAT_COMPLETION_STREAMING_TOOLS_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 From 4962948a3101b98419b28ad2c96b3061f2799f32 Mon Sep 17 00:00:00 2001 From: Erica Pisani Date: Tue, 11 Aug 2026 13:15:49 -0400 Subject: [PATCH 2/5] test(huggingface_hub): Inline data collection parametrize cases Drop the shared DATA_COLLECTION_PARAMS/DATA_COLLECTION_PARAM_NAMES constants and write the parameter names and cases directly in each test's parametrize decorator. --- .../huggingface_hub/test_huggingface_hub.py | 312 ++++++++++++++---- 1 file changed, 244 insertions(+), 68 deletions(-) diff --git a/tests/integrations/huggingface_hub/test_huggingface_hub.py b/tests/integrations/huggingface_hub/test_huggingface_hub.py index 83b567dd0b..58eb890ae4 100644 --- a/tests/integrations/huggingface_hub/test_huggingface_hub.py +++ b/tests/integrations/huggingface_hub/test_huggingface_hub.py @@ -1587,70 +1587,6 @@ def _get_gen_ai_span_data(captured: "Any", stream_gen_ai_spans: "Any") -> "Any": return span["data"] -DATA_COLLECTION_PARAMS = [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": True}}, - False, - False, - True, - True, - id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": False}}, - True, - True, - False, - False, - id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", - ), - pytest.param( - {"gen_ai": {"inputs": True, "outputs": False}}, - False, - False, - True, - False, - id="gen-ai-inputs-enabled-outputs-disabled", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": True}}, - False, - False, - False, - True, - id="gen-ai-outputs-enabled-inputs-disabled", - ), - pytest.param( - {"gen_ai": {}}, - False, - False, - True, - True, - id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", - ), - pytest.param( - None, - True, - True, - True, - True, - id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", - ), - pytest.param( - None, - False, - True, - False, - False, - id="no-gen-ai-config-legacy-pii-disabled", - ), -] - -DATA_COLLECTION_PARAM_NAMES = ( - "data_collection,send_default_pii,include_prompts,collect_inputs,collect_outputs" -) - - def _expected_keys( collect_inputs: "Any", collect_outputs: "Any", input_keys: "Any", output_keys: "Any" ) -> "Any": @@ -1665,7 +1601,67 @@ def _expected_keys( @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) -@pytest.mark.parametrize(DATA_COLLECTION_PARAM_NAMES, DATA_COLLECTION_PARAMS) +@pytest.mark.parametrize( + "data_collection,send_default_pii,include_prompts,collect_inputs,collect_outputs", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + False, + True, + True, + id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": False}}, + True, + True, + False, + False, + id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", + ), + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + False, + False, + True, + False, + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + False, + False, + False, + True, + id="gen-ai-outputs-enabled-inputs-disabled", + ), + pytest.param( + {"gen_ai": {}}, + False, + False, + True, + True, + id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", + ), + pytest.param( + None, + True, + True, + True, + True, + id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + False, + False, + id="no-gen-ai-config-legacy-pii-disabled", + ), + ], +) def test_text_generation_data_collection( sentry_init: "Any", capture_events: "Any", @@ -1731,7 +1727,67 @@ def test_text_generation_data_collection( @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) -@pytest.mark.parametrize(DATA_COLLECTION_PARAM_NAMES, DATA_COLLECTION_PARAMS) +@pytest.mark.parametrize( + "data_collection,send_default_pii,include_prompts,collect_inputs,collect_outputs", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + False, + True, + True, + id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": False}}, + True, + True, + False, + False, + id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", + ), + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + False, + False, + True, + False, + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + False, + False, + False, + True, + id="gen-ai-outputs-enabled-inputs-disabled", + ), + pytest.param( + {"gen_ai": {}}, + False, + False, + True, + True, + id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", + ), + pytest.param( + None, + True, + True, + True, + True, + id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + False, + False, + id="no-gen-ai-config-legacy-pii-disabled", + ), + ], +) def test_text_generation_streaming_data_collection( sentry_init: "Any", capture_events: "Any", @@ -1818,7 +1874,67 @@ def test_text_generation_streaming_data_collection( @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) -@pytest.mark.parametrize(DATA_COLLECTION_PARAM_NAMES, DATA_COLLECTION_PARAMS) +@pytest.mark.parametrize( + "data_collection,send_default_pii,include_prompts,collect_inputs,collect_outputs", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + False, + True, + True, + id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": False}}, + True, + True, + False, + False, + id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", + ), + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + False, + False, + True, + False, + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + False, + False, + False, + True, + id="gen-ai-outputs-enabled-inputs-disabled", + ), + pytest.param( + {"gen_ai": {}}, + False, + False, + True, + True, + id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", + ), + pytest.param( + None, + True, + True, + True, + True, + id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + False, + False, + id="no-gen-ai-config-legacy-pii-disabled", + ), + ], +) def test_chat_completion_data_collection_tools( sentry_init: "Any", capture_events: "Any", @@ -1905,7 +2021,67 @@ def test_chat_completion_data_collection_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_PARAM_NAMES, DATA_COLLECTION_PARAMS) +@pytest.mark.parametrize( + "data_collection,send_default_pii,include_prompts,collect_inputs,collect_outputs", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + False, + True, + True, + id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": False}}, + True, + True, + False, + False, + id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", + ), + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + False, + False, + True, + False, + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + False, + False, + False, + True, + id="gen-ai-outputs-enabled-inputs-disabled", + ), + pytest.param( + {"gen_ai": {}}, + False, + False, + True, + True, + id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", + ), + pytest.param( + None, + True, + True, + True, + True, + id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + False, + False, + id="no-gen-ai-config-legacy-pii-disabled", + ), + ], +) def test_chat_completion_streaming_data_collection_tools( sentry_init: "Any", capture_events: "Any", From 225aa79b50b55ae903ec6414ea6070912b146d46 Mon Sep 17 00:00:00 2001 From: Erica Pisani Date: Tue, 11 Aug 2026 13:19:23 -0400 Subject: [PATCH 3/5] test(huggingface_hub): Inline data collection test helpers Drop _expected_keys and _get_gen_ai_span_data. Each parametrized case now lists the expected present and absent span attributes explicitly, and the streaming/non-streaming span lookup lives in each test body. --- .../huggingface_hub/test_huggingface_hub.py | 360 +++++++++++------- 1 file changed, 221 insertions(+), 139 deletions(-) diff --git a/tests/integrations/huggingface_hub/test_huggingface_hub.py b/tests/integrations/huggingface_hub/test_huggingface_hub.py index 58eb890ae4..df04dfd0e5 100644 --- a/tests/integrations/huggingface_hub/test_huggingface_hub.py +++ b/tests/integrations/huggingface_hub/test_huggingface_hub.py @@ -1574,90 +1574,88 @@ def test_chat_completion_streaming_with_tools( } -def _get_gen_ai_span_data(captured: "Any", stream_gen_ai_spans: "Any") -> "Any": - 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") - ] - return span["attributes"] - - (transaction,) = captured - (span,) = [sp for sp in transaction["spans"] if sp["op"].startswith("gen_ai")] - return span["data"] - - -def _expected_keys( - collect_inputs: "Any", collect_outputs: "Any", input_keys: "Any", output_keys: "Any" -) -> "Any": - present = (input_keys if collect_inputs else []) + ( - output_keys if collect_outputs else [] - ) - absent = ([] if collect_inputs else input_keys) + ( - [] if collect_outputs else output_keys - ) - return present, absent - - @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,collect_inputs,collect_outputs", + "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", [ pytest.param( {"gen_ai": {"inputs": True, "outputs": True}}, False, False, - True, - True, + [ + 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, - False, - False, + [], + [ + 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, - True, - 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, - False, - True, + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + ], id="gen-ai-outputs-enabled-inputs-disabled", ), pytest.param( {"gen_ai": {}}, False, False, - True, - True, + [ + 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, - 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, - False, - False, + [], + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], id="no-gen-ai-config-legacy-pii-disabled", ), ], @@ -1670,17 +1668,10 @@ def test_text_generation_data_collection( data_collection: "Any", send_default_pii: "Any", include_prompts: "Any", - collect_inputs: "Any", - collect_outputs: "Any", + expected_present: "Any", + expected_absent: "Any", stream_gen_ai_spans: "Any", ) -> None: - expected_present, expected_absent = _expected_keys( - collect_inputs, - collect_outputs, - [SPANDATA.GEN_AI_REQUEST_MESSAGES], - [SPANDATA.GEN_AI_RESPONSE_TEXT], - ) - sentry_init_kwargs = dict( traces_sample_rate=1.0, send_default_pii=send_default_pii, @@ -1703,7 +1694,16 @@ def test_text_generation_data_collection( with sentry_sdk.start_transaction(name="test"): client.text_generation("Hello", stream=False, details=True) - span_data = _get_gen_ai_span_data(captured, stream_gen_ai_spans) + 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"] for key in expected_present: assert key in span_data, f"{key} should have been collected" @@ -1728,62 +1728,85 @@ def test_text_generation_data_collection( @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,collect_inputs,collect_outputs", + "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", [ pytest.param( {"gen_ai": {"inputs": True, "outputs": True}}, False, False, - True, - True, + [ + 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, - False, - False, + [], + [ + 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, - True, - 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, - False, - True, + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + ], id="gen-ai-outputs-enabled-inputs-disabled", ), pytest.param( {"gen_ai": {}}, False, False, - True, - True, + [ + 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, - 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, - False, - False, + [], + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], id="no-gen-ai-config-legacy-pii-disabled", ), ], @@ -1796,17 +1819,10 @@ def test_text_generation_streaming_data_collection( data_collection: "Any", send_default_pii: "Any", include_prompts: "Any", - collect_inputs: "Any", - collect_outputs: "Any", + expected_present: "Any", + expected_absent: "Any", stream_gen_ai_spans: "Any", ) -> None: - expected_present, expected_absent = _expected_keys( - collect_inputs, - collect_outputs, - [SPANDATA.GEN_AI_REQUEST_MESSAGES], - [SPANDATA.GEN_AI_RESPONSE_TEXT], - ) - sentry_init_kwargs = dict( traces_sample_rate=1.0, send_default_pii=send_default_pii, @@ -1830,7 +1846,16 @@ def test_text_generation_streaming_data_collection( for _ in client.text_generation(prompt="Hello", stream=True, details=True): pass - span_data = _get_gen_ai_span_data(captured, stream_gen_ai_spans) + 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"] for key in expected_present: assert key in span_data, f"{key} should have been collected" @@ -1875,62 +1900,91 @@ def test_text_generation_streaming_data_collection( @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,collect_inputs,collect_outputs", + "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", [ pytest.param( {"gen_ai": {"inputs": True, "outputs": True}}, False, 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-enabled-override-legacy-off", ), pytest.param( {"gen_ai": {"inputs": False, "outputs": False}}, True, 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-disabled-override-legacy-on", ), pytest.param( {"gen_ai": {"inputs": True, "outputs": False}}, False, False, - True, - 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, - False, - True, + [], + [ + 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, - 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-omitted-default-to-enabled", ), pytest.param( None, True, True, - 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, - False, - 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", ), ], @@ -1943,24 +1997,10 @@ def test_chat_completion_data_collection_tools( data_collection: "Any", send_default_pii: "Any", include_prompts: "Any", - collect_inputs: "Any", - collect_outputs: "Any", + expected_present: "Any", + expected_absent: "Any", stream_gen_ai_spans: "Any", ) -> None: - expected_present, expected_absent = _expected_keys( - collect_inputs, - collect_outputs, - [SPANDATA.GEN_AI_REQUEST_MESSAGES, SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS], - [], - ) - - # Legacy behaviour sets the available tools unconditionally, data collection - # gates them on inputs - if data_collection is None or collect_inputs: - expected_present.append(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS) - else: - expected_absent.append(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS) - sentry_init_kwargs = dict( traces_sample_rate=1.0, send_default_pii=send_default_pii, @@ -1987,7 +2027,16 @@ def test_chat_completion_data_collection_tools( tool_choice="auto", ) - span_data = _get_gen_ai_span_data(captured, stream_gen_ai_spans) + 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"] for key in expected_present: assert key in span_data, f"{key} should have been collected" @@ -2022,62 +2071,100 @@ def test_chat_completion_data_collection_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,collect_inputs,collect_outputs", + "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", [ pytest.param( {"gen_ai": {"inputs": True, "outputs": True}}, False, 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-enabled-override-legacy-off", ), pytest.param( {"gen_ai": {"inputs": False, "outputs": False}}, True, 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-disabled-override-legacy-on", ), pytest.param( {"gen_ai": {"inputs": True, "outputs": False}}, False, False, - True, - 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, - False, - True, + [ + 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, - 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-omitted-default-to-enabled", ), pytest.param( None, True, True, - 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, - False, - 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", ), ], @@ -2090,24 +2177,10 @@ def test_chat_completion_streaming_data_collection_tools( data_collection: "Any", send_default_pii: "Any", include_prompts: "Any", - collect_inputs: "Any", - collect_outputs: "Any", + expected_present: "Any", + expected_absent: "Any", stream_gen_ai_spans: "Any", ) -> None: - expected_present, expected_absent = _expected_keys( - collect_inputs, - collect_outputs, - [SPANDATA.GEN_AI_REQUEST_MESSAGES, SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS], - [SPANDATA.GEN_AI_RESPONSE_TEXT], - ) - - # Legacy behaviour sets the available tools unconditionally, data collection - # gates them on inputs - if data_collection is None or collect_inputs: - expected_present.append(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS) - else: - expected_absent.append(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS) - sentry_init_kwargs = dict( traces_sample_rate=1.0, send_default_pii=send_default_pii, @@ -2136,7 +2209,16 @@ def test_chat_completion_streaming_data_collection_tools( ): pass - span_data = _get_gen_ai_span_data(captured, stream_gen_ai_spans) + 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"] for key in expected_present: assert key in span_data, f"{key} should have been collected" From 85f993b9d24ed9b6432e5c83b051611d554c2ccf Mon Sep 17 00:00:00 2001 From: Erica Pisani Date: Tue, 11 Aug 2026 13:22:09 -0400 Subject: [PATCH 4/5] test(huggingface_hub): Group data collection test constants together The DATA_COLLECTION_* constants were interleaved with the tests that used them. Move them into a single block ahead of the data collection tests. --- .../huggingface_hub/test_huggingface_hub.py | 69 +++++++++---------- 1 file changed, 33 insertions(+), 36 deletions(-) diff --git a/tests/integrations/huggingface_hub/test_huggingface_hub.py b/tests/integrations/huggingface_hub/test_huggingface_hub.py index df04dfd0e5..c2b40f6ebb 100644 --- a/tests/integrations/huggingface_hub/test_huggingface_hub.py +++ b/tests/integrations/huggingface_hub/test_huggingface_hub.py @@ -1568,11 +1568,44 @@ def test_chat_completion_streaming_with_tools( 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"], + }, + }, + } +] + DATA_COLLECTION_TEXT_GENERATION_EXPECTED_VALUES = { SPANDATA.GEN_AI_REQUEST_MESSAGES: "Hello", SPANDATA.GEN_AI_RESPONSE_TEXT: "[mocked] Hello! How can i help you?", } +DATA_COLLECTION_TEXT_GENERATION_STREAMING_EXPECTED_VALUES = { + SPANDATA.GEN_AI_REQUEST_MESSAGES: "Hello", + SPANDATA.GEN_AI_RESPONSE_TEXT: "the mocked model response", +} + +DATA_COLLECTION_CHAT_COMPLETION_TOOLS_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"}]', +} + +DATA_COLLECTION_CHAT_COMPLETION_STREAMING_TOOLS_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", +} + @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @@ -1719,12 +1752,6 @@ def test_text_generation_data_collection( assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 10 -DATA_COLLECTION_TEXT_GENERATION_STREAMING_EXPECTED_VALUES = { - SPANDATA.GEN_AI_REQUEST_MESSAGES: "Hello", - SPANDATA.GEN_AI_RESPONSE_TEXT: "the mocked model response", -} - - @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize( @@ -1875,28 +1902,6 @@ def test_text_generation_streaming_data_collection( assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 10 -DATA_COLLECTION_TOOLS = [ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get current weather", - "parameters": { - "type": "object", - "properties": {"location": {"type": "string"}}, - "required": ["location"], - }, - }, - } -] - -DATA_COLLECTION_CHAT_COMPLETION_TOOLS_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"}]', -} - - @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize( @@ -2060,14 +2065,6 @@ def test_chat_completion_data_collection_tools( assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 18 -DATA_COLLECTION_CHAT_COMPLETION_STREAMING_TOOLS_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", -} - - @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @pytest.mark.parametrize( From 375584527ea2b76bc6b34c83dc88cbe22d883a1e Mon Sep 17 00:00:00 2001 From: Erica Pisani Date: Tue, 11 Aug 2026 13:24:19 -0400 Subject: [PATCH 5/5] test(huggingface_hub): Inline single-use expected value dicts The four DATA_COLLECTION_*_EXPECTED_VALUES dicts were each used by one test. Move them into the test bodies as expected_values; DATA_COLLECTION_TOOLS stays module-level since two tests share it. --- .../huggingface_hub/test_huggingface_hub.py | 62 ++++++++----------- 1 file changed, 27 insertions(+), 35 deletions(-) diff --git a/tests/integrations/huggingface_hub/test_huggingface_hub.py b/tests/integrations/huggingface_hub/test_huggingface_hub.py index c2b40f6ebb..9e7a77f66e 100644 --- a/tests/integrations/huggingface_hub/test_huggingface_hub.py +++ b/tests/integrations/huggingface_hub/test_huggingface_hub.py @@ -1583,29 +1583,6 @@ def test_chat_completion_streaming_with_tools( } ] -DATA_COLLECTION_TEXT_GENERATION_EXPECTED_VALUES = { - SPANDATA.GEN_AI_REQUEST_MESSAGES: "Hello", - SPANDATA.GEN_AI_RESPONSE_TEXT: "[mocked] Hello! How can i help you?", -} - -DATA_COLLECTION_TEXT_GENERATION_STREAMING_EXPECTED_VALUES = { - SPANDATA.GEN_AI_REQUEST_MESSAGES: "Hello", - SPANDATA.GEN_AI_RESPONSE_TEXT: "the mocked model response", -} - -DATA_COLLECTION_CHAT_COMPLETION_TOOLS_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"}]', -} - -DATA_COLLECTION_CHAT_COMPLETION_STREAMING_TOOLS_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", -} - @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.httpx_mock(assert_all_requests_were_expected=False) @@ -1738,9 +1715,14 @@ def test_text_generation_data_collection( (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] == DATA_COLLECTION_TEXT_GENERATION_EXPECTED_VALUES[key] + 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" @@ -1884,12 +1866,14 @@ def test_text_generation_streaming_data_collection( (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] - == DATA_COLLECTION_TEXT_GENERATION_STREAMING_EXPECTED_VALUES[key] - ) + 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" @@ -2043,11 +2027,15 @@ def test_chat_completion_data_collection_tools( (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] == DATA_COLLECTION_CHAT_COMPLETION_TOOLS_EXPECTED_VALUES[key] - ) + 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" @@ -2217,12 +2205,16 @@ def test_chat_completion_streaming_data_collection_tools( (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] - == DATA_COLLECTION_CHAT_COMPLETION_STREAMING_TOOLS_EXPECTED_VALUES[key] - ) + 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"