diff --git a/sentry_sdk/integrations/litellm.py b/sentry_sdk/integrations/litellm.py index 467f7c07ad..ea0f8897e9 100644 --- a/sentry_sdk/integrations/litellm.py +++ b/sentry_sdk/integrations/litellm.py @@ -10,10 +10,8 @@ ) from sentry_sdk.consts import 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 ( event_from_exception, - has_data_collection_enabled, package_version, ) @@ -113,14 +111,7 @@ def _input_callback(kwargs: "Dict[str, Any]") -> None: set_data_normalized(span, SPANDATA.GEN_AI_PROVIDER_NAME, provider) set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, operation) - # Record input/messages if allowed - record_inputs = False - if has_data_collection_enabled(client.options): - record_inputs = client.options["data_collection"]["gen_ai"]["inputs"] - elif should_send_default_pii(): - record_inputs = True - - if record_inputs: + if client.options["data_collection"]["gen_ai"]["inputs"]: if operation == "embeddings": # For embeddings, look for the 'input' parameter embedding_input = kwargs.get("input") @@ -194,14 +185,7 @@ def _success_callback( span, SPANDATA.GEN_AI_RESPONSE_MODEL, completion_response.model ) - # Record response content if allowed - record_outputs = False - if has_data_collection_enabled(client.options): - record_outputs = client.options["data_collection"]["gen_ai"]["outputs"] - elif should_send_default_pii(): - record_outputs = True - - if record_outputs: + if client.options["data_collection"]["gen_ai"]["outputs"]: if hasattr(completion_response, "choices"): response_messages = [] for choice in completion_response.choices: diff --git a/tests/integrations/litellm/test_litellm.py b/tests/integrations/litellm/test_litellm.py index 145ddaa6e3..7131399eed 100644 --- a/tests/integrations/litellm/test_litellm.py +++ b/tests/integrations/litellm/test_litellm.py @@ -134,15 +134,10 @@ def __init__( self.created = 1234567890 -@pytest.mark.parametrize( - "send_default_pii", - [True, False], -) def test_nonstreaming_chat_completion( reset_litellm_executor, sentry_init, capture_items, - send_default_pii, get_model_response, nonstreaming_chat_completions_model_response, ): @@ -150,7 +145,12 @@ def test_nonstreaming_chat_completion( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) messages = [ @@ -208,18 +208,96 @@ def test_nonstreaming_chat_completion( assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "openai" assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" - if send_default_pii: - assert json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) == [ - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "Hello!"}, - ] - assert SPANDATA.GEN_AI_RESPONSE_TEXT in span["attributes"] - else: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) == [ + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "Hello!"}, + ] + assert SPANDATA.GEN_AI_RESPONSE_TEXT in span["attributes"] + + assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert span["attributes"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + + +def test_nonstreaming_chat_completion_no_sensitive_data( + reset_litellm_executor, + sentry_init, + capture_items, + get_model_response, + nonstreaming_chat_completions_model_response, +): + sentry_init( + integrations=[LiteLLMIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, + ) + + messages = [ + {"role": "user", "content": "Message demonstrating the absence of truncation."}, + {"role": "user", "content": "Hello!"}, + ] + + client = OpenAI(api_key="test-key") + + model_response = get_model_response( + nonstreaming_chat_completions_model_response( + response_id="chatcmpl-test", + response_model="gpt-3.5-turbo", + message_content="Test response", + created=1234567890, + usage=CompletionUsage( + prompt_tokens=10, + completion_tokens=20, + total_tokens=30, + ), + ), + serialize_pydantic=True, + request_headers={"X-Stainless-Raw-Response": "true"}, + ) + items = capture_items("span") + + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ): + litellm.completion( + model="gpt-3.5-turbo", + messages=messages, + client=client, + ) + + litellm_utils.executor.shutdown(wait=True) + + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"].get("sentry.op") == OP.GEN_AI_CHAT + and x["attributes"].get("sentry.origin") == "auto.ai.litellm" + ) + assert len(chat_spans) == 1 + span = chat_spans[0] + + assert span["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + assert span["name"] == "chat gpt-3.5-turbo" + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gpt-3.5-turbo" + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-3.5-turbo" + assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "openai" + assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 @@ -227,14 +305,9 @@ def test_nonstreaming_chat_completion( @pytest.mark.asyncio(loop_scope="session") -@pytest.mark.parametrize( - "send_default_pii", - [True, False], -) async def test_async_nonstreaming_chat_completion( sentry_init, capture_items, - send_default_pii, get_model_response, nonstreaming_chat_completions_model_response, ): @@ -242,7 +315,12 @@ async def test_async_nonstreaming_chat_completion( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) messages = [ @@ -301,33 +379,107 @@ async def test_async_nonstreaming_chat_completion( assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "openai" assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" - if send_default_pii: - assert json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) == [ - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "Hello!"}, - ] - assert SPANDATA.GEN_AI_RESPONSE_TEXT in span["attributes"] - else: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + assert json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) == [ + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "Hello!"}, + ] + assert SPANDATA.GEN_AI_RESPONSE_TEXT in span["attributes"] + + assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + assert span["attributes"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + + +@pytest.mark.asyncio(loop_scope="session") +async def test_async_nonstreaming_chat_completion_no_sensitive_data( + sentry_init, + capture_items, + get_model_response, + nonstreaming_chat_completions_model_response, +): + sentry_init( + integrations=[LiteLLMIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, + ) + + messages = [ + {"role": "user", "content": "Message demonstrating the absence of truncation."}, + {"role": "user", "content": "Hello!"}, + ] + + client = AsyncOpenAI(api_key="test-key") + + model_response = get_model_response( + nonstreaming_chat_completions_model_response( + response_id="chatcmpl-test", + response_model="gpt-3.5-turbo", + message_content="Test response", + created=1234567890, + usage=CompletionUsage( + prompt_tokens=10, + completion_tokens=20, + total_tokens=30, + ), + ), + serialize_pydantic=True, + request_headers={"X-Stainless-Raw-Response": "true"}, + ) + items = capture_items("span") + + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ): + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=messages, + client=client, + ) + + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) + + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"].get("sentry.op") == OP.GEN_AI_CHAT + and x["attributes"].get("sentry.origin") == "auto.ai.litellm" + ) + assert len(chat_spans) == 1 + span = chat_spans[0] + + assert span["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + assert span["name"] == "chat gpt-3.5-turbo" + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gpt-3.5-turbo" + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-3.5-turbo" + assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "openai" + assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 assert span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 assert span["attributes"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 -@pytest.mark.parametrize( - "send_default_pii", - [True, False], -) def test_streaming_chat_completion( reset_litellm_executor, sentry_init, capture_items, - send_default_pii, get_model_response, server_side_event_chunks, streaming_chat_completions_model_response, @@ -336,7 +488,7 @@ def test_streaming_chat_completion( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={}, ) messages = [ @@ -387,14 +539,9 @@ def test_streaming_chat_completion( @pytest.mark.asyncio(loop_scope="session") -@pytest.mark.parametrize( - "send_default_pii", - [True, False], -) async def test_async_streaming_chat_completion( sentry_init, capture_items, - send_default_pii, get_model_response, async_iterator, server_side_event_chunks, @@ -404,7 +551,7 @@ async def test_async_streaming_chat_completion( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={}, ) messages = [ @@ -474,7 +621,12 @@ def test_embeddings_create( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) client = OpenAI(api_key="test-key") @@ -544,7 +696,12 @@ async def test_async_embeddings_create( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) client = AsyncOpenAI(api_key="test-key") @@ -609,7 +766,12 @@ def test_embeddings_create_with_list_input( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) client = OpenAI(api_key="test-key") @@ -673,7 +835,12 @@ async def test_async_embeddings_create_with_list_input( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) client = AsyncOpenAI(api_key="test-key") @@ -724,19 +891,23 @@ async def test_async_embeddings_create_with_list_input( ] -def test_embeddings_no_pii( +def test_embeddings_no_sensitive_data( sentry_init, capture_items, get_model_response, openai_embedding_model_response, clear_litellm_cache, ): - """Test that PII is not captured when disabled.""" sentry_init( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=False, # PII disabled + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, ) client = OpenAI(api_key="test-key") @@ -775,24 +946,27 @@ def test_embeddings_no_pii( span = spans[0] assert span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - # Check that embeddings input is NOT captured when PII is disabled assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["attributes"] @pytest.mark.asyncio(loop_scope="session") -async def test_async_embeddings_no_pii( +async def test_async_embeddings_no_sensitive_data( sentry_init, capture_items, get_model_response, openai_embedding_model_response, clear_litellm_cache, ): - """Test that PII is not captured when disabled.""" sentry_init( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=False, # PII disabled + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, ) client = AsyncOpenAI(api_key="test-key") @@ -832,7 +1006,6 @@ async def test_async_embeddings_no_pii( span = spans[0] assert span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - # Check that embeddings input is NOT captured when PII is disabled assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["attributes"] @@ -846,6 +1019,7 @@ def test_exception_handling( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + data_collection={}, ) messages = [{"role": "user", "content": "Hello!"}] @@ -881,6 +1055,7 @@ async def test_async_exception_handling( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + data_collection={}, ) messages = [{"role": "user", "content": "Hello!"}] @@ -917,6 +1092,7 @@ def test_span_origin( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + data_collection={}, ) messages = [{"role": "user", "content": "Hello!"}] @@ -972,6 +1148,7 @@ def test_multiple_providers( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + data_collection={}, ) messages = [{"role": "user", "content": "Hello!"}] @@ -1073,6 +1250,7 @@ async def test_async_multiple_providers( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + data_collection={}, ) messages = [{"role": "user", "content": "Hello!"}] @@ -1175,6 +1353,7 @@ def test_additional_parameters( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + data_collection={}, ) messages = [{"role": "user", "content": "Hello!"}] @@ -1245,6 +1424,7 @@ async def test_async_additional_parameters( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + data_collection={}, ) messages = [{"role": "user", "content": "Hello!"}] @@ -1315,6 +1495,7 @@ def test_no_integration( sentry_init( disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + data_collection={}, ) messages = [{"role": "user", "content": "Hello!"}] @@ -1372,6 +1553,7 @@ async def test_async_no_integration( sentry_init( disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + data_collection={}, ) messages = [{"role": "user", "content": "Hello!"}] @@ -1428,6 +1610,7 @@ def test_response_without_usage( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + data_collection={}, ) messages = [{"role": "user", "content": "Hello!"}] @@ -1470,6 +1653,7 @@ def test_integration_setup(sentry_init): integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + data_collection={}, ) # Check that callbacks are registered @@ -1494,7 +1678,12 @@ def test_binary_content_encoding_image_url( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) messages = [ @@ -1584,7 +1773,12 @@ async def test_async_binary_content_encoding_image_url( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) messages = [ @@ -1675,7 +1869,12 @@ def test_binary_content_encoding_mixed_content( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) messages = [ @@ -1754,7 +1953,12 @@ async def test_async_binary_content_encoding_mixed_content( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) messages = [ @@ -1834,7 +2038,12 @@ def test_binary_content_encoding_uri_type( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) messages = [ @@ -1917,7 +2126,12 @@ async def test_async_binary_content_encoding_uri_type( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) messages = [ @@ -2092,57 +2306,20 @@ def test_convert_message_parts_image_url_missing_url(): @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], [SPANDATA.GEN_AI_RESPONSE_TEXT], id="gen-ai-inputs-enabled-outputs-disabled", ), pytest.param( {"gen_ai": {"inputs": False, "outputs": True}}, - False, [SPANDATA.GEN_AI_RESPONSE_TEXT], [SPANDATA.GEN_AI_REQUEST_MESSAGES], 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_chat_completion_data_collection( @@ -2152,7 +2329,6 @@ def test_chat_completion_data_collection( get_model_response, nonstreaming_chat_completions_model_response, data_collection, - send_default_pii, expected_present, expected_absent, ): @@ -2160,7 +2336,6 @@ def test_chat_completion_data_collection( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=send_default_pii, ) if data_collection is not None: sentry_init_kwargs["data_collection"] = data_collection @@ -2226,108 +2401,7 @@ def test_chat_completion_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] == "chat" assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "gpt-3.5-turbo" assert span_data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-3.5-turbo" assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 - - -@pytest.mark.parametrize( - "data_collection, send_default_pii, should_collect_input", - [ - pytest.param( - {"gen_ai": {"inputs": True}}, - False, - True, - id="gen-ai-inputs-enabled-override-legacy-off", - ), - pytest.param( - {"gen_ai": {"inputs": False}}, - True, - False, - id="gen-ai-inputs-disabled-override-legacy-on", - ), - pytest.param( - {"gen_ai": {}}, - False, - True, - id="gen-ai-inputs-omitted-default-to-enabled", - ), - pytest.param( - None, - True, - True, - id="no-gen-ai-config-legacy-pii", - ), - pytest.param( - None, - False, - False, - id="no-gen-ai-config-legacy-pii-disabled", - ), - ], -) -def test_embeddings_data_collection( - sentry_init, - capture_items, - get_model_response, - openai_embedding_model_response, - data_collection, - send_default_pii, - should_collect_input, -): - sentry_init_kwargs = dict( - integrations=[LiteLLMIntegration()], - disabled_integrations=[StdlibIntegration], - traces_sample_rate=1.0, - send_default_pii=send_default_pii, - ) - if data_collection is not None: - sentry_init_kwargs["data_collection"] = data_collection - - sentry_init(**sentry_init_kwargs) - - client = OpenAI(api_key="test-key") - - model_response = get_model_response( - openai_embedding_model_response, - serialize_pydantic=True, - request_headers={"X-Stainless-Raw-Response": "true"}, - ) - - items = capture_items("span") - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ): - litellm.embedding( - model="text-embedding-ada-002", - input="Hello, world!", - client=client, - ) - # Allow time for callbacks to complete (they may run in separate threads) - time.sleep(0.1) - - spans = [item.payload for item in items] - (span,) = [ - x - for x in spans - if x["attributes"].get("sentry.op") == OP.GEN_AI_EMBEDDINGS - and x["attributes"].get("sentry.origin") == "auto.ai.litellm" - ] - span_data = span["attributes"] - - if should_collect_input: - assert json.loads(span_data[SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == [ - "Hello, world!" - ] - else: - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span_data - - # Data collection never gates non-PII attributes - assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-ada-002" - assert span_data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5