From 8d62c84a166649e3d9b00deea4c915d50ffded0b Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Fri, 4 Sep 2026 12:45:57 +0200 Subject: [PATCH 1/4] chore(litellm): Remove transaction-based tracing --- sentry_sdk/integrations/litellm.py | 27 +- tests/integrations/litellm/test_litellm.py | 2723 ++++++-------------- 2 files changed, 812 insertions(+), 1938 deletions(-) diff --git a/sentry_sdk/integrations/litellm.py b/sentry_sdk/integrations/litellm.py index 7c0a84d867..2ef9433ac5 100644 --- a/sentry_sdk/integrations/litellm.py +++ b/sentry_sdk/integrations/litellm.py @@ -5,7 +5,6 @@ from sentry_sdk import consts from sentry_sdk.ai.monitoring import record_token_usage from sentry_sdk.ai.utils import ( - get_start_span_function, set_data_normalized, transform_openai_content_part, truncate_and_annotate_embedding_inputs, @@ -101,29 +100,17 @@ def _input_callback(kwargs: "Dict[str, Any]") -> None: operation = "chat" # Start a new span/transaction - if has_span_streaming_enabled(client.options): - span = sentry_sdk.traces.start_span( - name=f"{operation} {model}", - attributes={ - "sentry.op": ( - consts.OP.GEN_AI_CHAT - if operation == "chat" - else consts.OP.GEN_AI_EMBEDDINGS - ), - "sentry.origin": LiteLLMIntegration.origin, - }, - ) - else: - span = get_start_span_function()( - op=( + span = sentry_sdk.traces.start_span( + name=f"{operation} {model}", + attributes={ + "sentry.op": ( consts.OP.GEN_AI_CHAT if operation == "chat" else consts.OP.GEN_AI_EMBEDDINGS ), - name=f"{operation} {model}", - origin=LiteLLMIntegration.origin, - ) - span.__enter__() + "sentry.origin": LiteLLMIntegration.origin, + }, + ) _store_span(kwargs, span) diff --git a/tests/integrations/litellm/test_litellm.py b/tests/integrations/litellm/test_litellm.py index e4371262d1..50cd4e690a 100644 --- a/tests/integrations/litellm/test_litellm.py +++ b/tests/integrations/litellm/test_litellm.py @@ -135,7 +135,6 @@ def __init__( self.created = 1234567890 -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "send_default_pii, include_prompts", [ @@ -148,21 +147,18 @@ def __init__( def test_nonstreaming_chat_completion( reset_litellm_executor, sentry_init, - capture_events, capture_items, send_default_pii, include_prompts, get_model_response, nonstreaming_chat_completions_model_response, - span_streaming, ): sentry_init( integrations=[LiteLLMIntegration(include_prompts=include_prompts)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [ @@ -187,108 +183,58 @@ def test_nonstreaming_chat_completion( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("transaction", "span") - if span_streaming: - items = capture_items("transaction", "span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), sentry_sdk.traces.start_span(name="litellm test"): - 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 if item.type == "span"] - assert spans[2]["name"] == "litellm test" - 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_SYSTEM] == "openai" - assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" - - if send_default_pii and include_prompts: - 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 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 - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - litellm.completion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - ) - - litellm_utils.executor.shutdown(wait=True) - - assert len(events) == 1 - (event,) = events - - assert event["type"] == "transaction" - assert event["transaction"] == "litellm test" - - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), sentry_sdk.traces.start_span(name="litellm test"): + litellm.completion( + model="gpt-3.5-turbo", + messages=messages, + client=client, ) - assert len(chat_spans) == 1 - span = chat_spans[0] - assert span["op"] == OP.GEN_AI_CHAT - assert span["description"] == "chat gpt-3.5-turbo" - 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_SYSTEM] == "openai" - assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + litellm_utils.executor.shutdown(wait=True) - if send_default_pii and include_prompts: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span["data"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT in span["data"] - else: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[2]["name"] == "litellm test" + 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_SYSTEM] == "openai" + assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + + if send_default_pii and include_prompts: + 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 span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 - assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 - assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + 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("span_streaming", [True, False]) @pytest.mark.asyncio(loop_scope="session") @pytest.mark.parametrize( "send_default_pii, include_prompts", @@ -301,21 +247,18 @@ def test_nonstreaming_chat_completion( ) async def test_async_nonstreaming_chat_completion( sentry_init, - capture_events, capture_items, send_default_pii, include_prompts, get_model_response, nonstreaming_chat_completions_model_response, - span_streaming, ): sentry_init( integrations=[LiteLLMIntegration(include_prompts=include_prompts)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [ @@ -340,110 +283,59 @@ async def test_async_nonstreaming_chat_completion( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("transaction", "span") - if span_streaming: - items = capture_items("transaction", "span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), sentry_sdk.traces.start_span(name="litellm test"): - 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 if item.type == "span"] - assert spans[2]["name"] == "litellm test" - 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_SYSTEM] == "openai" - assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" - - if send_default_pii and include_prompts: - 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 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 - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - ) - - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) - - assert len(events) == 1 - (event,) = events - - assert event["type"] == "transaction" - assert event["transaction"] == "litellm test" - - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), sentry_sdk.traces.start_span(name="litellm test"): + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=messages, + client=client, ) - assert len(chat_spans) == 1 - span = chat_spans[0] - assert span["op"] == OP.GEN_AI_CHAT - assert span["description"] == "chat gpt-3.5-turbo" - 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_SYSTEM] == "openai" - assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) - if send_default_pii and include_prompts: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span["data"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT in span["data"] - else: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[2]["name"] == "litellm test" + 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_SYSTEM] == "openai" + assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + + if send_default_pii and include_prompts: + 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 span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 - assert span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 - assert span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + 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("span_streaming", [True, False]) @pytest.mark.parametrize( "send_default_pii, include_prompts", [ @@ -456,22 +348,19 @@ async def test_async_nonstreaming_chat_completion( def test_streaming_chat_completion( reset_litellm_executor, sentry_init, - capture_events, capture_items, send_default_pii, include_prompts, get_model_response, server_side_event_chunks, streaming_chat_completions_model_response, - span_streaming, ): sentry_init( integrations=[LiteLLMIntegration(include_prompts=include_prompts)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [ @@ -488,75 +377,39 @@ def test_streaming_chat_completion( ), request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - response = litellm.completion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - stream=True, - ) - for _ in response: - pass - - streaming_handler.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"]["sentry.op"] == OP.GEN_AI_CHAT - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + response = litellm.completion( + model="gpt-3.5-turbo", + messages=messages, + client=client, + stream=True, ) - assert len(chat_spans) == 1 - span = chat_spans[0] + for _ in response: + pass - assert span["attributes"]["sentry.op"] == OP.GEN_AI_CHAT - assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - response = litellm.completion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - stream=True, - ) - for _ in response: - pass + streaming_handler.executor.shutdown(wait=True) - streaming_handler.executor.shutdown(wait=True) - - assert len(events) == 1 - (event,) = events - - assert event["type"] == "transaction" - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" - ) - assert len(chat_spans) == 1 - span = chat_spans[0] + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(chat_spans) == 1 + span = chat_spans[0] - assert span["op"] == OP.GEN_AI_CHAT - assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + assert span["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio(loop_scope="session") @pytest.mark.parametrize( "send_default_pii, include_prompts", @@ -569,7 +422,6 @@ def test_streaming_chat_completion( ) async def test_async_streaming_chat_completion( sentry_init, - capture_events, capture_items, send_default_pii, include_prompts, @@ -577,15 +429,13 @@ async def test_async_streaming_chat_completion( async_iterator, server_side_event_chunks, streaming_chat_completions_model_response, - span_streaming, ): sentry_init( integrations=[LiteLLMIntegration(include_prompts=include_prompts)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [ @@ -604,85 +454,46 @@ async def test_async_streaming_chat_completion( ), request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - response = await litellm.acompletion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - stream=True, - ) - async for _ in response: - pass - - 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"]["sentry.op"] == OP.GEN_AI_CHAT - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + response = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=messages, + client=client, + stream=True, ) - assert len(chat_spans) == 1 - span = chat_spans[0] - - assert span["attributes"]["sentry.op"] == OP.GEN_AI_CHAT - assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - response = await litellm.acompletion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - stream=True, - ) - async for _ in response: - pass - - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) + async for _ in response: + pass - assert len(events) == 1 - (event,) = events + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) - assert event["type"] == "transaction" - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" - ) - assert len(chat_spans) == 1 - span = chat_spans[0] + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(chat_spans) == 1 + span = chat_spans[0] - assert span["op"] == OP.GEN_AI_CHAT - assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + assert span["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True -@pytest.mark.parametrize("span_streaming", [True, False]) def test_embeddings_create( sentry_init, - capture_events, capture_items, get_model_response, openai_embedding_model_response, clear_litellm_cache, - span_streaming, ): """ Test that litellm.embedding() calls are properly instrumented. @@ -695,8 +506,7 @@ def test_embeddings_create( disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) client = OpenAI(api_key="test-key") @@ -706,102 +516,56 @@ def test_embeddings_create( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ): - with start_transaction(name="litellm test"): - 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) - - # Response is processed by litellm, so just check it exists - assert response is not None - sentry_sdk.flush() - spans = [item.payload for item in items] - spans = list( - x - for x in spans - if x["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.embeddings._client._client, + "send", + return_value=model_response, + ): + with start_transaction(name="litellm test"): + 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) - assert len(spans) == 1 - span = spans[0] + # Response is processed by litellm, so just check it exists + assert response is not None + sentry_sdk.flush() + spans = [item.payload for item in items] + spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) - assert span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - assert span["name"] == "embeddings text-embedding-ada-002" - assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 - assert ( - span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] - == "text-embedding-ada-002" - ) - # Check that embeddings input is captured (it's JSON serialized) - embeddings_input = span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] - assert json.loads(embeddings_input) == ["Hello, world!"] - else: - events = capture_events() - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ): - with start_transaction(name="litellm test"): - 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) - - # Response is processed by litellm, so just check it exists - assert response is not None - assert len(events) == 1 - (event,) = events - - assert event["type"] == "transaction" - spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_EMBEDDINGS and x["origin"] == "auto.ai.litellm" - ) - assert len(spans) == 1 - span = spans[0] - - assert span["op"] == OP.GEN_AI_EMBEDDINGS - assert span["description"] == "embeddings text-embedding-ada-002" - assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 - assert ( - span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-ada-002" - ) - # Check that embeddings input is captured (it's JSON serialized) - embeddings_input = span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] - assert json.loads(embeddings_input) == ["Hello, world!"] + assert len(spans) == 1 + span = spans[0] + + assert span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + assert span["name"] == "embeddings text-embedding-ada-002" + assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" + assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 + assert ( + span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] + == "text-embedding-ada-002" + ) + # Check that embeddings input is captured (it's JSON serialized) + embeddings_input = span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] + assert json.loads(embeddings_input) == ["Hello, world!"] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio(loop_scope="session") async def test_async_embeddings_create( sentry_init, - capture_events, capture_items, get_model_response, openai_embedding_model_response, clear_litellm_cache, - span_streaming, ): """ Test that litellm.embedding() calls are properly instrumented. @@ -814,8 +578,7 @@ async def test_async_embeddings_create( disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) client = AsyncOpenAI(api_key="test-key") @@ -825,103 +588,56 @@ async def test_async_embeddings_create( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ): - with start_transaction(name="litellm test"): - response = await litellm.aembedding( - model="text-embedding-ada-002", - input="Hello, world!", - client=client, - ) - - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) - - # Response is processed by litellm, so just check it exists - assert response is not None - sentry_sdk.flush() - spans = [item.payload for item in items] - spans = list( - x - for x in spans - if x["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" - ) - assert len(spans) == 1 - span = spans[0] - - assert span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - assert span["name"] == "embeddings text-embedding-ada-002" - assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 - assert ( - span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] - == "text-embedding-ada-002" + with mock.patch.object( + client.embeddings._client._client, + "send", + return_value=model_response, + ): + with start_transaction(name="litellm test"): + response = await litellm.aembedding( + model="text-embedding-ada-002", + input="Hello, world!", + client=client, ) - # Check that embeddings input is captured (it's JSON serialized) - embeddings_input = span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] - assert json.loads(embeddings_input) == ["Hello, world!"] - else: - events = capture_events() - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ): - with start_transaction(name="litellm test"): - response = await litellm.aembedding( - model="text-embedding-ada-002", - input="Hello, world!", - client=client, - ) - - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) - - # Response is processed by litellm, so just check it exists - assert response is not None - assert len(events) == 1 - (event,) = events - - assert event["type"] == "transaction" - spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_EMBEDDINGS and x["origin"] == "auto.ai.litellm" - ) - assert len(spans) == 1 - span = spans[0] - - assert span["op"] == OP.GEN_AI_EMBEDDINGS - assert span["description"] == "embeddings text-embedding-ada-002" - assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 - assert ( - span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-ada-002" - ) - # Check that embeddings input is captured (it's JSON serialized) - embeddings_input = span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] - assert json.loads(embeddings_input) == ["Hello, world!"] + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) + + # Response is processed by litellm, so just check it exists + assert response is not None + sentry_sdk.flush() + spans = [item.payload for item in items] + spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(spans) == 1 + span = spans[0] + + assert span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + assert span["name"] == "embeddings text-embedding-ada-002" + assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" + assert span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 + assert ( + span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] + == "text-embedding-ada-002" + ) + # Check that embeddings input is captured (it's JSON serialized) + embeddings_input = span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] + + assert json.loads(embeddings_input) == ["Hello, world!"] -@pytest.mark.parametrize("span_streaming", [True, False]) def test_embeddings_create_with_list_input( sentry_init, - capture_events, capture_items, get_model_response, openai_embedding_model_response, clear_litellm_cache, - span_streaming, ): """Test embedding with list input.""" sentry_init( @@ -929,8 +645,7 @@ def test_embeddings_create_with_list_input( disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) client = OpenAI(api_key="test-key") @@ -940,100 +655,55 @@ def test_embeddings_create_with_list_input( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ): - with start_transaction(name="litellm test"): - response = litellm.embedding( - model="text-embedding-ada-002", - input=["First text", "Second text", "Third text"], - client=client, - ) - # Allow time for callbacks to complete (they may run in separate threads) - time.sleep(0.1) - - # Response is processed by litellm, so just check it exists - assert response is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - spans = list( - x - for x in spans - if x["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" - ) - assert len(spans) == 1 - span = spans[0] - - assert span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - # Check that list of embeddings input is captured (it's JSON serialized) - embeddings_input = span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] - - assert json.loads(embeddings_input) == [ - "First text", - "Second text", - "Third text", - ] - else: - events = capture_events() - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ): - with start_transaction(name="litellm test"): - response = litellm.embedding( - model="text-embedding-ada-002", - input=["First text", "Second text", "Third text"], - client=client, - ) - # Allow time for callbacks to complete (they may run in separate threads) - time.sleep(0.1) - - # Response is processed by litellm, so just check it exists - assert response is not None - assert len(events) == 1 - (event,) = events - - assert event["type"] == "transaction" - spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_EMBEDDINGS and x["origin"] == "auto.ai.litellm" + with mock.patch.object( + client.embeddings._client._client, + "send", + return_value=model_response, + ): + with start_transaction(name="litellm test"): + response = litellm.embedding( + model="text-embedding-ada-002", + input=["First text", "Second text", "Third text"], + client=client, ) - assert len(spans) == 1 - span = spans[0] + # Allow time for callbacks to complete (they may run in separate threads) + time.sleep(0.1) - assert span["op"] == OP.GEN_AI_EMBEDDINGS - assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - # Check that list of embeddings input is captured (it's JSON serialized) - embeddings_input = span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] - assert json.loads(embeddings_input) == [ - "First text", - "Second text", - "Third text", - ] + # Response is processed by litellm, so just check it exists + assert response is not None + + sentry_sdk.flush() + spans = [item.payload for item in items] + spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(spans) == 1 + span = spans[0] + + assert span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" + # Check that list of embeddings input is captured (it's JSON serialized) + embeddings_input = span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] + + assert json.loads(embeddings_input) == [ + "First text", + "Second text", + "Third text", + ] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio(loop_scope="session") async def test_async_embeddings_create_with_list_input( sentry_init, - capture_events, capture_items, get_model_response, openai_embedding_model_response, clear_litellm_cache, - span_streaming, ): """Test embedding with list input.""" sentry_init( @@ -1041,8 +711,7 @@ async def test_async_embeddings_create_with_list_input( disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) client = AsyncOpenAI(api_key="test-key") @@ -1052,100 +721,54 @@ async def test_async_embeddings_create_with_list_input( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ): - with start_transaction(name="litellm test"): - response = await litellm.aembedding( - model="text-embedding-ada-002", - input=["First text", "Second text", "Third text"], - client=client, - ) - - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) - - # Response is processed by litellm, so just check it exists - assert response is not None - sentry_sdk.flush() - spans = [item.payload for item in items] - spans = list( - x - for x in spans - if x["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" - ) - assert len(spans) == 1 - span = spans[0] - - assert span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - # Check that list of embeddings input is captured (it's JSON serialized) - embeddings_input = span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] - - assert json.loads(embeddings_input) == [ - "First text", - "Second text", - "Third text", - ] - else: - events = capture_events() - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ): - with start_transaction(name="litellm test"): - response = await litellm.aembedding( - model="text-embedding-ada-002", - input=["First text", "Second text", "Third text"], - client=client, - ) - - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) - - # Response is processed by litellm, so just check it exists - assert response is not None - assert len(events) == 1 - (event,) = events - - assert event["type"] == "transaction" - spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_EMBEDDINGS and x["origin"] == "auto.ai.litellm" + with mock.patch.object( + client.embeddings._client._client, + "send", + return_value=model_response, + ): + with start_transaction(name="litellm test"): + response = await litellm.aembedding( + model="text-embedding-ada-002", + input=["First text", "Second text", "Third text"], + client=client, ) - assert len(spans) == 1 - span = spans[0] - assert span["op"] == OP.GEN_AI_EMBEDDINGS - assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - # Check that list of embeddings input is captured (it's JSON serialized) - embeddings_input = span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] - assert json.loads(embeddings_input) == [ - "First text", - "Second text", - "Third text", - ] + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) + + # Response is processed by litellm, so just check it exists + assert response is not None + sentry_sdk.flush() + spans = [item.payload for item in items] + spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(spans) == 1 + span = spans[0] + + assert span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" + # Check that list of embeddings input is captured (it's JSON serialized) + embeddings_input = span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] + + assert json.loads(embeddings_input) == [ + "First text", + "Second text", + "Third text", + ] -@pytest.mark.parametrize("span_streaming", [True, False]) def test_embeddings_no_pii( sentry_init, - capture_events, capture_items, get_model_response, openai_embedding_model_response, clear_litellm_cache, - span_streaming, ): """Test that PII is not captured when disabled.""" sentry_init( @@ -1153,8 +776,7 @@ def test_embeddings_no_pii( disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=False, # PII disabled - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) client = OpenAI(api_key="test-key") @@ -1164,86 +786,47 @@ def test_embeddings_no_pii( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ): - with start_transaction(name="litellm test"): - 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) - - # Response is processed by litellm, so just check it exists - assert response is not None - sentry_sdk.flush() - spans = [item.payload for item in items] - spans = list( - x - for x in spans - if x["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.embeddings._client._client, + "send", + return_value=model_response, + ): + with start_transaction(name="litellm test"): + response = litellm.embedding( + model="text-embedding-ada-002", + input="Hello, world!", + client=client, ) - assert len(spans) == 1 - span = spans[0] + # Allow time for callbacks to complete (they may run in separate threads) + time.sleep(0.1) - 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"] - else: - events = capture_events() - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ): - with start_transaction(name="litellm test"): - 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) - - # Response is processed by litellm, so just check it exists - assert response is not None - assert len(events) == 1 - (event,) = events - - assert event["type"] == "transaction" - spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_EMBEDDINGS and x["origin"] == "auto.ai.litellm" - ) - assert len(spans) == 1 - span = spans[0] + # Response is processed by litellm, so just check it exists + assert response is not None + sentry_sdk.flush() + spans = [item.payload for item in items] + spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(spans) == 1 + span = spans[0] - assert span["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["data"] + 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.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio(loop_scope="session") async def test_async_embeddings_no_pii( sentry_init, - capture_events, capture_items, get_model_response, openai_embedding_model_response, clear_litellm_cache, - span_streaming, ): """Test that PII is not captured when disabled.""" sentry_init( @@ -1251,8 +834,7 @@ async def test_async_embeddings_no_pii( disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=False, # PII disabled - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) client = AsyncOpenAI(api_key="test-key") @@ -1262,93 +844,52 @@ async def test_async_embeddings_no_pii( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ): - with start_transaction(name="litellm test"): - response = await litellm.aembedding( - model="text-embedding-ada-002", - input="Hello, world!", - client=client, - ) - - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) - - # Response is processed by litellm, so just check it exists - assert response is not None - sentry_sdk.flush() - spans = [item.payload for item in items] - spans = list( - x - for x in spans - if x["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.embeddings._client._client, + "send", + return_value=model_response, + ): + with start_transaction(name="litellm test"): + response = await litellm.aembedding( + model="text-embedding-ada-002", + input="Hello, world!", + client=client, ) - assert len(spans) == 1 - 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"] - else: - events = capture_events() - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ): - with start_transaction(name="litellm test"): - response = await litellm.aembedding( - model="text-embedding-ada-002", - input="Hello, world!", - client=client, - ) - - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) - - # Response is processed by litellm, so just check it exists - assert response is not None - assert len(events) == 1 - (event,) = events - - assert event["type"] == "transaction" - spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_EMBEDDINGS and x["origin"] == "auto.ai.litellm" - ) - assert len(spans) == 1 - span = spans[0] + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) - assert span["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["data"] + # Response is processed by litellm, so just check it exists + assert response is not None + sentry_sdk.flush() + spans = [item.payload for item in items] + spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(spans) == 1 + 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.parametrize("span_streaming", [True, False]) def test_exception_handling( reset_litellm_executor, sentry_init, - capture_events, capture_items, get_rate_limit_model_response, - span_streaming, ): sentry_init( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [{"role": "user", "content": "Hello!"}] @@ -1356,63 +897,35 @@ def test_exception_handling( client = OpenAI(api_key="test-key") model_response = get_rate_limit_model_response() + items = capture_items("event") - if span_streaming: - items = capture_items("event") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"), pytest.raises( - litellm.RateLimitError - ): - litellm.completion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - ) - - # Find the error event - error_events = [item.payload for item in items] - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"), pytest.raises( - litellm.RateLimitError - ): - litellm.completion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - ) + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"), pytest.raises(litellm.RateLimitError): + litellm.completion( + model="gpt-3.5-turbo", + messages=messages, + client=client, + ) - # Should have error event and transaction - assert len(events) >= 1 - # Find the error event - error_events = [e for e in events if e.get("level") == "error"] + # Find the error event + error_events = [item.payload for item in items] assert len(error_events) == 1 -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio(loop_scope="session") async def test_async_exception_handling( sentry_init, - capture_events, capture_items, get_rate_limit_model_response, - span_streaming, ): sentry_init( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [{"role": "user", "content": "Hello!"}] @@ -1420,64 +933,36 @@ async def test_async_exception_handling( client = AsyncOpenAI(api_key="test-key") model_response = get_rate_limit_model_response() + items = capture_items("event") - if span_streaming: - items = capture_items("event") - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"), pytest.raises( - litellm.RateLimitError - ): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - ) - - # Find the error event - error_events = [item.payload for item in items] - else: - events = capture_events() - - with mock.patch.object( - client.embeddings._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"), pytest.raises( - litellm.RateLimitError - ): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - ) + with mock.patch.object( + client.embeddings._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"), pytest.raises(litellm.RateLimitError): + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=messages, + client=client, + ) - # Should have error event and transaction - assert len(events) >= 1 - # Find the error event - error_events = [e for e in events if e.get("level") == "error"] + # Find the error event + error_events = [item.payload for item in items] assert len(error_events) == 1 -@pytest.mark.parametrize("span_streaming", [True, False]) def test_span_origin( reset_litellm_executor, sentry_init, - capture_events, capture_items, get_model_response, nonstreaming_chat_completions_model_response, - span_streaming, ): sentry_init( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [{"role": "user", "content": "Hello!"}] @@ -1499,68 +984,43 @@ def test_span_origin( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), sentry_sdk.traces.start_span(name="litellm test"): - 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] - assert spans[2]["is_segment"] is True - assert spans[2]["attributes"]["sentry.origin"] == "manual" - assert spans[1]["attributes"]["sentry.origin"] == "auto.ai.litellm" - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - litellm.completion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - ) + items = capture_items("span") - litellm_utils.executor.shutdown(wait=True) + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), sentry_sdk.traces.start_span(name="litellm test"): + litellm.completion( + model="gpt-3.5-turbo", + messages=messages, + client=client, + ) - (event,) = events + litellm_utils.executor.shutdown(wait=True) - assert event["contexts"]["trace"]["origin"] == "manual" - assert event["spans"][0]["origin"] == "auto.ai.litellm" + sentry_sdk.flush() + spans = [item.payload for item in items] + assert spans[2]["is_segment"] is True + assert spans[2]["attributes"]["sentry.origin"] == "manual" + assert spans[1]["attributes"]["sentry.origin"] == "auto.ai.litellm" -@pytest.mark.parametrize("span_streaming", [True, False]) def test_multiple_providers( reset_litellm_executor, sentry_init, - capture_events, capture_items, get_model_response, nonstreaming_chat_completions_model_response, nonstreaming_anthropic_model_response, nonstreaming_google_genai_model_response, - span_streaming, ): """Test that the integration correctly identifies different providers.""" sentry_init( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [{"role": "user", "content": "Hello!"}] @@ -1581,163 +1041,90 @@ def test_multiple_providers( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - openai_client.completions._client._client, - "send", - return_value=openai_model_response, - ), sentry_sdk.traces.start_span(name="test gpt-3.5-turbo"): - litellm.completion( - model="gpt-3.5-turbo", - messages=messages, - client=openai_client, - ) - - litellm_utils.executor.shutdown(wait=True) - - _reset_litellm_executor() - - anthropic_client = HTTPHandler() - anthropic_model_response = get_model_response( - nonstreaming_anthropic_model_response, - serialize_pydantic=True, - request_headers={"X-Stainless-Raw-Response": "true"}, - ) - - with mock.patch.object( - anthropic_client, - "post", - return_value=anthropic_model_response, - ), sentry_sdk.traces.start_span(name="test claude-3-opus-20240229"): - litellm.completion( - model="claude-3-opus-20240229", - messages=messages, - client=anthropic_client, - api_key="test-key", - ) - - litellm_utils.executor.shutdown(wait=True) - - _reset_litellm_executor() - - gemini_client = HTTPHandler() - gemini_model_response = get_model_response( - nonstreaming_google_genai_model_response, - serialize_pydantic=True, + with mock.patch.object( + openai_client.completions._client._client, + "send", + return_value=openai_model_response, + ), sentry_sdk.traces.start_span(name="test gpt-3.5-turbo"): + litellm.completion( + model="gpt-3.5-turbo", + messages=messages, + client=openai_client, ) - with mock.patch.object( - gemini_client, - "post", - return_value=gemini_model_response, - ), sentry_sdk.traces.start_span(name="test gemini/gemini-pro"): - litellm.completion( - model="gemini/gemini-pro", - messages=messages, - client=gemini_client, - api_key="test-key", - ) - - litellm_utils.executor.shutdown(wait=True) - - sentry_sdk.flush() - spans = [item.payload for item in items] - for span in spans: - if span["is_segment"] is True: - continue - # The provider should be detected by litellm.get_llm_provider - assert SPANDATA.GEN_AI_SYSTEM in span["attributes"] - else: - events = capture_events() - - with mock.patch.object( - openai_client.completions._client._client, - "send", - return_value=openai_model_response, - ), start_transaction(name="test gpt-3.5-turbo"): - litellm.completion( - model="gpt-3.5-turbo", - messages=messages, - client=openai_client, - ) + litellm_utils.executor.shutdown(wait=True) - litellm_utils.executor.shutdown(wait=True) + _reset_litellm_executor() - _reset_litellm_executor() + anthropic_client = HTTPHandler() + anthropic_model_response = get_model_response( + nonstreaming_anthropic_model_response, + serialize_pydantic=True, + request_headers={"X-Stainless-Raw-Response": "true"}, + ) - anthropic_client = HTTPHandler() - anthropic_model_response = get_model_response( - nonstreaming_anthropic_model_response, - serialize_pydantic=True, - request_headers={"X-Stainless-Raw-Response": "true"}, + with mock.patch.object( + anthropic_client, + "post", + return_value=anthropic_model_response, + ), sentry_sdk.traces.start_span(name="test claude-3-opus-20240229"): + litellm.completion( + model="claude-3-opus-20240229", + messages=messages, + client=anthropic_client, + api_key="test-key", ) - with mock.patch.object( - anthropic_client, - "post", - return_value=anthropic_model_response, - ), start_transaction(name="test claude-3-opus-20240229"): - litellm.completion( - model="claude-3-opus-20240229", - messages=messages, - client=anthropic_client, - api_key="test-key", - ) + litellm_utils.executor.shutdown(wait=True) - litellm_utils.executor.shutdown(wait=True) + _reset_litellm_executor() - _reset_litellm_executor() + gemini_client = HTTPHandler() + gemini_model_response = get_model_response( + nonstreaming_google_genai_model_response, + serialize_pydantic=True, + ) - gemini_client = HTTPHandler() - gemini_model_response = get_model_response( - nonstreaming_google_genai_model_response, - serialize_pydantic=True, + with mock.patch.object( + gemini_client, + "post", + return_value=gemini_model_response, + ), sentry_sdk.traces.start_span(name="test gemini/gemini-pro"): + litellm.completion( + model="gemini/gemini-pro", + messages=messages, + client=gemini_client, + api_key="test-key", ) - with mock.patch.object( - gemini_client, - "post", - return_value=gemini_model_response, - ), start_transaction(name="test gemini/gemini-pro"): - litellm.completion( - model="gemini/gemini-pro", - messages=messages, - client=gemini_client, - api_key="test-key", - ) - - litellm_utils.executor.shutdown(wait=True) - - assert len(events) == 3 + litellm_utils.executor.shutdown(wait=True) - for i in range(3): - span = events[i]["spans"][0] - # The provider should be detected by litellm.get_llm_provider - assert SPANDATA.GEN_AI_SYSTEM in span["data"] + sentry_sdk.flush() + spans = [item.payload for item in items] + for span in spans: + if span["is_segment"] is True: + continue + # The provider should be detected by litellm.get_llm_provider + assert SPANDATA.GEN_AI_SYSTEM in span["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio(loop_scope="session") async def test_async_multiple_providers( sentry_init, - capture_events, capture_items, get_model_response, nonstreaming_chat_completions_model_response, nonstreaming_anthropic_model_response, nonstreaming_google_genai_model_response, - span_streaming, ): """Test that the integration correctly identifies different providers.""" sentry_init( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [{"role": "user", "content": "Hello!"}] @@ -1758,167 +1145,91 @@ async def test_async_multiple_providers( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("transaction", "span") - if span_streaming: - items = capture_items("transaction", "span") - - with mock.patch.object( - openai_client.completions._client._client, - "send", - return_value=openai_model_response, - ), sentry_sdk.traces.start_span(name="test gpt-3.5-turbo"): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=messages, - client=openai_client, - ) - - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) - - _reset_litellm_executor() - - anthropic_client = AsyncHTTPHandler() - anthropic_model_response = get_model_response( - nonstreaming_anthropic_model_response, - serialize_pydantic=True, - request_headers={"X-Stainless-Raw-Response": "True"}, + with mock.patch.object( + openai_client.completions._client._client, + "send", + return_value=openai_model_response, + ), sentry_sdk.traces.start_span(name="test gpt-3.5-turbo"): + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=messages, + client=openai_client, ) - with mock.patch.object( - anthropic_client, - "post", - return_value=anthropic_model_response, - ), sentry_sdk.traces.start_span(name="test claude-3-opus-20240229"): - await litellm.acompletion( - model="claude-3-opus-20240229", - messages=messages, - client=anthropic_client, - api_key="test-key", - ) + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) - - _reset_litellm_executor() - - gemini_client = AsyncHTTPHandler() - gemini_model_response = get_model_response( - nonstreaming_google_genai_model_response, - serialize_pydantic=True, - ) - - with mock.patch.object( - gemini_client, - "post", - return_value=gemini_model_response, - ), sentry_sdk.traces.start_span(name="test gemini/gemini-pro"): - await litellm.acompletion( - model="gemini/gemini-pro", - messages=messages, - client=gemini_client, - api_key="test-key", - ) - - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) - - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - for span in spans: - if span["is_segment"] is True: - continue - # The provider should be detected by litellm.get_llm_provider - assert SPANDATA.GEN_AI_SYSTEM in span["attributes"] - else: - events = capture_events() - - with mock.patch.object( - openai_client.completions._client._client, - "send", - return_value=openai_model_response, - ), start_transaction(name="test gpt-3.5-turbo"): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=messages, - client=openai_client, - ) - - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) + _reset_litellm_executor() - _reset_litellm_executor() + anthropic_client = AsyncHTTPHandler() + anthropic_model_response = get_model_response( + nonstreaming_anthropic_model_response, + serialize_pydantic=True, + request_headers={"X-Stainless-Raw-Response": "True"}, + ) - anthropic_client = AsyncHTTPHandler() - anthropic_model_response = get_model_response( - nonstreaming_anthropic_model_response, - serialize_pydantic=True, - request_headers={"X-Stainless-Raw-Response": "True"}, + with mock.patch.object( + anthropic_client, + "post", + return_value=anthropic_model_response, + ), sentry_sdk.traces.start_span(name="test claude-3-opus-20240229"): + await litellm.acompletion( + model="claude-3-opus-20240229", + messages=messages, + client=anthropic_client, + api_key="test-key", ) - with mock.patch.object( - anthropic_client, - "post", - return_value=anthropic_model_response, - ), start_transaction(name="test claude-3-opus-20240229"): - await litellm.acompletion( - model="claude-3-opus-20240229", - messages=messages, - client=anthropic_client, - api_key="test-key", - ) + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) + _reset_litellm_executor() - _reset_litellm_executor() + gemini_client = AsyncHTTPHandler() + gemini_model_response = get_model_response( + nonstreaming_google_genai_model_response, + serialize_pydantic=True, + ) - gemini_client = AsyncHTTPHandler() - gemini_model_response = get_model_response( - nonstreaming_google_genai_model_response, - serialize_pydantic=True, + with mock.patch.object( + gemini_client, + "post", + return_value=gemini_model_response, + ), sentry_sdk.traces.start_span(name="test gemini/gemini-pro"): + await litellm.acompletion( + model="gemini/gemini-pro", + messages=messages, + client=gemini_client, + api_key="test-key", ) - with mock.patch.object( - gemini_client, - "post", - return_value=gemini_model_response, - ), start_transaction(name="test gemini/gemini-pro"): - await litellm.acompletion( - model="gemini/gemini-pro", - messages=messages, - client=gemini_client, - api_key="test-key", - ) - - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) - - assert len(events) == 3 + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) - for i in range(3): - span = events[i]["spans"][0] - # The provider should be detected by litellm.get_llm_provider - assert SPANDATA.GEN_AI_SYSTEM in span["data"] + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + for span in spans: + if span["is_segment"] is True: + continue + # The provider should be detected by litellm.get_llm_provider + assert SPANDATA.GEN_AI_SYSTEM in span["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) def test_additional_parameters( reset_litellm_executor, sentry_init, - capture_events, capture_items, get_model_response, nonstreaming_chat_completions_model_response, - span_streaming, ): """Test that additional parameters are captured.""" sentry_init( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [{"role": "user", "content": "Hello!"}] @@ -1939,98 +1250,57 @@ def test_additional_parameters( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - litellm.completion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - temperature=0.7, - max_tokens=100, - top_p=0.9, - frequency_penalty=0.5, - presence_penalty=0.5, - ) - - 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"]["sentry.op"] == OP.GEN_AI_CHAT - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + litellm.completion( + model="gpt-3.5-turbo", + messages=messages, + client=client, + temperature=0.7, + max_tokens=100, + top_p=0.9, + frequency_penalty=0.5, + presence_penalty=0.5, ) - assert len(chat_spans) == 1 - span = chat_spans[0] - - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.5 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.5 - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - litellm.completion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - temperature=0.7, - max_tokens=100, - top_p=0.9, - frequency_penalty=0.5, - presence_penalty=0.5, - ) - litellm_utils.executor.shutdown(wait=True) + litellm_utils.executor.shutdown(wait=True) - (event,) = events - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" - ) - assert len(chat_spans) == 1 - span = chat_spans[0] + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(chat_spans) == 1 + span = chat_spans[0] - assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.5 - assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.5 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.5 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.5 -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio(loop_scope="session") async def test_async_additional_parameters( sentry_init, - capture_events, capture_items, get_model_response, nonstreaming_chat_completions_model_response, - span_streaming, ): """Test that additional parameters are captured.""" sentry_init( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [{"role": "user", "content": "Hello!"}] @@ -2051,99 +1321,57 @@ async def test_async_additional_parameters( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - temperature=0.7, - max_tokens=100, - top_p=0.9, - frequency_penalty=0.5, - presence_penalty=0.5, - ) - - 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"]["sentry.op"] == OP.GEN_AI_CHAT - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=messages, + client=client, + temperature=0.7, + max_tokens=100, + top_p=0.9, + frequency_penalty=0.5, + presence_penalty=0.5, ) - assert len(chat_spans) == 1 - span = chat_spans[0] - - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.5 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.5 - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - temperature=0.7, - max_tokens=100, - top_p=0.9, - frequency_penalty=0.5, - presence_penalty=0.5, - ) - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) - (event,) = events - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" - ) - assert len(chat_spans) == 1 - span = chat_spans[0] + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(chat_spans) == 1 + span = chat_spans[0] - assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.5 - assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.5 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.5 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.5 -@pytest.mark.parametrize("span_streaming", [True, False]) def test_no_integration( reset_litellm_executor, sentry_init, - capture_events, capture_items, get_model_response, nonstreaming_chat_completions_model_response, - span_streaming, ): """Test that when integration is not enabled, callbacks don't break.""" sentry_init( disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [{"role": "user", "content": "Hello!"}] @@ -2164,74 +1392,44 @@ def test_no_integration( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - 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"]["sentry.op"] == OP.GEN_AI_CHAT - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + litellm.completion( + model="gpt-3.5-turbo", + messages=messages, + client=client, ) - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - litellm.completion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - ) - litellm_utils.executor.shutdown(wait=True) + litellm_utils.executor.shutdown(wait=True) - (event,) = events - # Should still have the transaction, but no child spans since integration is off - assert event["type"] == "transaction" - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" - ) + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) assert len(chat_spans) == 0 -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio(loop_scope="session") async def test_async_no_integration( sentry_init, - capture_events, capture_items, get_model_response, nonstreaming_chat_completions_model_response, - span_streaming, ): """Test that when integration is not enabled, callbacks don't break.""" sentry_init( disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [{"role": "user", "content": "Hello!"}] @@ -2252,74 +1450,43 @@ async def test_async_no_integration( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - 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"]["sentry.op"] == OP.GEN_AI_CHAT - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=messages, + client=client, ) - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=messages, - client=client, - ) - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) - (event,) = events - # Should still have the transaction, but no child spans since integration is off - assert event["type"] == "transaction" - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" - ) + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) assert len(chat_spans) == 0 -@pytest.mark.parametrize("span_streaming", [True, False]) def test_response_without_usage( sentry_init, - capture_events, capture_items, - span_streaming, ): """Test handling of responses without usage information.""" sentry_init( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [{"role": "user", "content": "Hello!"}] @@ -2333,92 +1500,7 @@ def test_response_without_usage( "choices": [MockChoice()], }, )() - - if span_streaming: - items = capture_items("span") - - with start_transaction(name="litellm test"): - kwargs = { - "model": "gpt-3.5-turbo", - "messages": messages, - } - - _input_callback(kwargs) - _success_callback( - kwargs, - mock_response, - datetime.now(), - datetime.now(), - ) - - sentry_sdk.flush() - (span,) = (item.payload for item in items) - - # Span should still be created even without usage info - assert span["attributes"]["sentry.op"] == OP.GEN_AI_CHAT - assert span["name"] == "chat gpt-3.5-turbo" - else: - events = capture_events() - - with start_transaction(name="litellm test"): - kwargs = { - "model": "gpt-3.5-turbo", - "messages": messages, - } - - _input_callback(kwargs) - _success_callback( - kwargs, - mock_response, - datetime.now(), - datetime.now(), - ) - - (event,) = events - (span,) = event["spans"] - - # Span should still be created even without usage info - assert span["op"] == OP.GEN_AI_CHAT - assert span["description"] == "chat gpt-3.5-turbo" - - -def test_integration_setup(sentry_init): - """Test that the integration sets up the callbacks correctly.""" - sentry_init( - integrations=[LiteLLMIntegration()], - disabled_integrations=[StdlibIntegration], - traces_sample_rate=1.0, - stream_gen_ai_spans=False, - ) - - # Check that callbacks are registered - assert _input_callback in (litellm.input_callback or []) - assert _success_callback in (litellm.success_callback or []) - assert _failure_callback in (litellm.failure_callback or []) - - -def test_litellm_message_truncation(sentry_init, capture_events): - """Test that large messages are truncated properly in LiteLLM integration.""" - sentry_init( - integrations=[LiteLLMIntegration(include_prompts=True)], - disabled_integrations=[StdlibIntegration], - traces_sample_rate=1.0, - send_default_pii=True, - stream_gen_ai_spans=False, - ) - events = capture_events() - - large_content = ( - "This is a very long message that will exceed our size limits. " * 1000 - ) - messages = [ - {"role": "user", "content": "small message 1"}, - {"role": "assistant", "content": large_content}, - {"role": "user", "content": large_content}, - {"role": "assistant", "content": "small message 4"}, - {"role": "user", "content": "small message 5"}, - ] - mock_response = MockCompletionResponse() + items = capture_items("span") with start_transaction(name="litellm test"): kwargs = { @@ -2434,26 +1516,26 @@ def test_litellm_message_truncation(sentry_init, capture_events): datetime.now(), ) - assert len(events) > 0 - tx = events[0] - assert tx["type"] == "transaction" + sentry_sdk.flush() + (span,) = (item.payload for item in items) - chat_spans = [ - span for span in tx.get("spans", []) if span.get("op") == OP.GEN_AI_CHAT - ] - assert len(chat_spans) > 0 + # Span should still be created even without usage info + assert span["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + assert span["name"] == "chat gpt-3.5-turbo" - chat_span = chat_spans[0] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in chat_span["data"] - messages_data = chat_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert isinstance(messages_data, str) +def test_integration_setup(sentry_init): + """Test that the integration sets up the callbacks correctly.""" + sentry_init( + integrations=[LiteLLMIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + ) - parsed_messages = json.loads(messages_data) - assert isinstance(parsed_messages, list) - assert len(parsed_messages) == 1 - assert "small message 5" in str(parsed_messages[0]) - assert tx["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 5 + # Check that callbacks are registered + assert _input_callback in (litellm.input_callback or []) + assert _success_callback in (litellm.success_callback or []) + assert _failure_callback in (litellm.failure_callback or []) IMAGE_DATA = b"fake_image_data_12345" @@ -2461,23 +1543,19 @@ def test_litellm_message_truncation(sentry_init, capture_events): IMAGE_DATA_URI = f"data:image/png;base64,{IMAGE_B64}" -@pytest.mark.parametrize("span_streaming", [True, False]) def test_binary_content_encoding_image_url( reset_litellm_executor, sentry_init, - capture_events, capture_items, get_model_response, nonstreaming_chat_completions_model_response, - span_streaming, ): sentry_init( integrations=[LiteLLMIntegration(include_prompts=True)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [ @@ -2509,61 +1587,33 @@ def test_binary_content_encoding_image_url( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - litellm.completion( - model="gpt-4-vision-preview", - messages=messages, - client=client, - custom_llm_provider="openai", - ) - - 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"]["sentry.op"] == OP.GEN_AI_CHAT - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + litellm.completion( + model="gpt-4-vision-preview", + messages=messages, + client=client, + custom_llm_provider="openai", ) - assert len(chat_spans) == 1 - span = chat_spans[0] - messages_data = json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - litellm.completion( - model="gpt-4-vision-preview", - messages=messages, - client=client, - custom_llm_provider="openai", - ) - litellm_utils.executor.shutdown(wait=True) + litellm_utils.executor.shutdown(wait=True) - (event,) = events - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" - ) - assert len(chat_spans) == 1 - span = chat_spans[0] - messages_data = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(chat_spans) == 1 + span = chat_spans[0] + messages_data = json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) blob_item = next( ( @@ -2584,23 +1634,19 @@ def test_binary_content_encoding_image_url( ) -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio(loop_scope="session") async def test_async_binary_content_encoding_image_url( sentry_init, - capture_events, capture_items, get_model_response, nonstreaming_chat_completions_model_response, - span_streaming, ): sentry_init( integrations=[LiteLLMIntegration(include_prompts=True)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [ @@ -2632,63 +1678,34 @@ async def test_async_binary_content_encoding_image_url( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - await litellm.acompletion( - model="gpt-4-vision-preview", - messages=messages, - client=client, - custom_llm_provider="openai", - ) - - 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"]["sentry.op"] == OP.GEN_AI_CHAT - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + await litellm.acompletion( + model="gpt-4-vision-preview", + messages=messages, + client=client, + custom_llm_provider="openai", ) - assert len(chat_spans) == 1 - span = chat_spans[0] - messages_data = json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - await litellm.acompletion( - model="gpt-4-vision-preview", - messages=messages, - client=client, - custom_llm_provider="openai", - ) - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) - (event,) = events - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" - ) - assert len(chat_spans) == 1 - span = chat_spans[0] - messages_data = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(chat_spans) == 1 + span = chat_spans[0] + messages_data = json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) blob_item = next( ( @@ -2709,23 +1726,19 @@ async def test_async_binary_content_encoding_image_url( ) -@pytest.mark.parametrize("span_streaming", [True, False]) def test_binary_content_encoding_mixed_content( reset_litellm_executor, sentry_init, - capture_events, capture_items, get_model_response, nonstreaming_chat_completions_model_response, - span_streaming, ): sentry_init( integrations=[LiteLLMIntegration(include_prompts=True)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [ @@ -2758,61 +1771,33 @@ def test_binary_content_encoding_mixed_content( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - litellm.completion( - model="gpt-4-vision-preview", - messages=messages, - client=client, - custom_llm_provider="openai", - ) - - 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"]["sentry.op"] == OP.GEN_AI_CHAT - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + litellm.completion( + model="gpt-4-vision-preview", + messages=messages, + client=client, + custom_llm_provider="openai", ) - assert len(chat_spans) == 1 - span = chat_spans[0] - messages_data = json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - litellm.completion( - model="gpt-4-vision-preview", - messages=messages, - client=client, - custom_llm_provider="openai", - ) - litellm_utils.executor.shutdown(wait=True) + litellm_utils.executor.shutdown(wait=True) - (event,) = events - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" - ) - assert len(chat_spans) == 1 - span = chat_spans[0] - messages_data = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(chat_spans) == 1 + span = chat_spans[0] + messages_data = json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) content_items = [ item for msg in messages_data if "content" in msg for item in msg["content"] @@ -2821,23 +1806,19 @@ def test_binary_content_encoding_mixed_content( assert any(item.get("type") == "blob" for item in content_items) -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio(loop_scope="session") async def test_async_binary_content_encoding_mixed_content( sentry_init, - capture_events, capture_items, get_model_response, nonstreaming_chat_completions_model_response, - span_streaming, ): sentry_init( integrations=[LiteLLMIntegration(include_prompts=True)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [ @@ -2870,63 +1851,34 @@ async def test_async_binary_content_encoding_mixed_content( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - await litellm.acompletion( - model="gpt-4-vision-preview", - messages=messages, - client=client, - custom_llm_provider="openai", - ) - - 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"]["sentry.op"] == OP.GEN_AI_CHAT - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + await litellm.acompletion( + model="gpt-4-vision-preview", + messages=messages, + client=client, + custom_llm_provider="openai", ) - assert len(chat_spans) == 1 - span = chat_spans[0] - messages_data = json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - await litellm.acompletion( - model="gpt-4-vision-preview", - messages=messages, - client=client, - custom_llm_provider="openai", - ) - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) - (event,) = events - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" - ) - assert len(chat_spans) == 1 - span = chat_spans[0] - messages_data = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(chat_spans) == 1 + span = chat_spans[0] + messages_data = json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) content_items = [ item for msg in messages_data if "content" in msg for item in msg["content"] @@ -2935,23 +1887,19 @@ async def test_async_binary_content_encoding_mixed_content( assert any(item.get("type") == "blob" for item in content_items) -@pytest.mark.parametrize("span_streaming", [True, False]) def test_binary_content_encoding_uri_type( reset_litellm_executor, sentry_init, - capture_events, capture_items, get_model_response, nonstreaming_chat_completions_model_response, - span_streaming, ): sentry_init( integrations=[LiteLLMIntegration(include_prompts=True)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [ @@ -2982,60 +1930,32 @@ def test_binary_content_encoding_uri_type( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) - - if span_streaming: - items = capture_items("span") - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - litellm.completion( - model="gpt-4-vision-preview", - messages=messages, - client=client, - custom_llm_provider="openai", - ) - - 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"]["sentry.op"] == OP.GEN_AI_CHAT - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + items = capture_items("span") + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + litellm.completion( + model="gpt-4-vision-preview", + messages=messages, + client=client, + custom_llm_provider="openai", ) - assert len(chat_spans) == 1 - span = chat_spans[0] - messages_data = json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - litellm.completion( - model="gpt-4-vision-preview", - messages=messages, - client=client, - custom_llm_provider="openai", - ) - litellm_utils.executor.shutdown(wait=True) + litellm_utils.executor.shutdown(wait=True) - (event,) = events - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" - ) - assert len(chat_spans) == 1 - span = chat_spans[0] - messages_data = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(chat_spans) == 1 + span = chat_spans[0] + messages_data = json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) uri_item = next( ( @@ -3051,23 +1971,19 @@ def test_binary_content_encoding_uri_type( assert uri_item["uri"] == "https://example.com/image.jpg" -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio(loop_scope="session") async def test_async_binary_content_encoding_uri_type( sentry_init, - capture_events, capture_items, get_model_response, nonstreaming_chat_completions_model_response, - span_streaming, ): sentry_init( integrations=[LiteLLMIntegration(include_prompts=True)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) messages = [ @@ -3098,63 +2014,34 @@ async def test_async_binary_content_encoding_uri_type( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - await litellm.acompletion( - model="gpt-4-vision-preview", - messages=messages, - client=client, - custom_llm_provider="openai", - ) - - 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"]["sentry.op"] == OP.GEN_AI_CHAT - and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + await litellm.acompletion( + model="gpt-4-vision-preview", + messages=messages, + client=client, + custom_llm_provider="openai", ) - assert len(chat_spans) == 1 - span = chat_spans[0] - messages_data = json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) - else: - events = capture_events() - - with mock.patch.object( - client.completions._client._client, - "send", - return_value=model_response, - ), start_transaction(name="litellm test"): - await litellm.acompletion( - model="gpt-4-vision-preview", - messages=messages, - client=client, - custom_llm_provider="openai", - ) - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) - (event,) = events - chat_spans = list( - x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" - ) - assert len(chat_spans) == 1 - span = chat_spans[0] - messages_data = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + sentry_sdk.flush() + spans = [item.payload for item in items] + chat_spans = list( + x + for x in spans + if x["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + and x["attributes"]["sentry.origin"] == "auto.ai.litellm" + ) + assert len(chat_spans) == 1 + span = chat_spans[0] + messages_data = json.loads(span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) uri_item = next( ( @@ -3334,7 +2221,7 @@ def test_convert_message_parts_image_url_missing_url(): def test_chat_completion_data_collection( reset_litellm_executor, sentry_init, - capture_events, + capture_items, get_model_response, nonstreaming_chat_completions_model_response, data_collection, @@ -3348,8 +2235,7 @@ def test_chat_completion_data_collection( disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -3376,7 +2262,7 @@ def test_chat_completion_data_collection( request_headers={"X-Stainless-Raw-Response": "true"}, ) - events = capture_events() + items = capture_items("span") with mock.patch.object( client.completions._client._client, @@ -3391,13 +2277,14 @@ def test_chat_completion_data_collection( litellm_utils.executor.shutdown(wait=True) - (event,) = events + spans = [item.payload for item in items] (span,) = [ x - for x in event["spans"] - if x["op"] == OP.GEN_AI_CHAT and x["origin"] == "auto.ai.litellm" + for x in spans + if x["attributes"].get("sentry.op") == OP.GEN_AI_CHAT + and x["attributes"].get("sentry.origin") == "auto.ai.litellm" ] - span_data = span["data"] + span_data = span["attributes"] for key in expected_present: assert key in span_data, f"{key} should have been collected" @@ -3462,7 +2349,7 @@ def test_chat_completion_data_collection( ) def test_embeddings_data_collection( sentry_init, - capture_events, + capture_items, get_model_response, openai_embedding_model_response, data_collection, @@ -3475,8 +2362,7 @@ def test_embeddings_data_collection( disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -3491,7 +2377,7 @@ def test_embeddings_data_collection( request_headers={"X-Stainless-Raw-Response": "true"}, ) - events = capture_events() + items = capture_items("span") with mock.patch.object( client.embeddings._client._client, @@ -3507,13 +2393,14 @@ def test_embeddings_data_collection( # Allow time for callbacks to complete (they may run in separate threads) time.sleep(0.1) - (event,) = events + spans = [item.payload for item in items] (span,) = [ x - for x in event["spans"] - if x["op"] == OP.GEN_AI_EMBEDDINGS and x["origin"] == "auto.ai.litellm" + 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["data"] + span_data = span["attributes"] if should_collect_input: assert json.loads(span_data[SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == [ From ceea29793ccae9b23646d37b659ab62bca2da07d Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Fri, 4 Sep 2026 12:51:49 +0200 Subject: [PATCH 2/4] remove start_transaction --- tests/integrations/litellm/test_litellm.py | 159 ++++++++++----------- 1 file changed, 75 insertions(+), 84 deletions(-) diff --git a/tests/integrations/litellm/test_litellm.py b/tests/integrations/litellm/test_litellm.py index 50cd4e690a..dd84302c50 100644 --- a/tests/integrations/litellm/test_litellm.py +++ b/tests/integrations/litellm/test_litellm.py @@ -36,7 +36,6 @@ async def __call__(self, *args, **kwargs): from openai import AsyncOpenAI, OpenAI from openai.types import CompletionUsage -from sentry_sdk import start_transaction from sentry_sdk._types import BLOB_DATA_SUBSTITUTE from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations.litellm import ( @@ -383,7 +382,7 @@ def test_streaming_chat_completion( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"): + ): response = litellm.completion( model="gpt-3.5-turbo", messages=messages, @@ -460,7 +459,7 @@ async def test_async_streaming_chat_completion( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"): + ): response = await litellm.acompletion( model="gpt-3.5-turbo", messages=messages, @@ -523,14 +522,13 @@ def test_embeddings_create( "send", return_value=model_response, ): - with start_transaction(name="litellm test"): - 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) + 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) # Response is processed by litellm, so just check it exists assert response is not None @@ -595,15 +593,14 @@ async def test_async_embeddings_create( "send", return_value=model_response, ): - with start_transaction(name="litellm test"): - response = await litellm.aembedding( - model="text-embedding-ada-002", - input="Hello, world!", - client=client, - ) + response = await litellm.aembedding( + model="text-embedding-ada-002", + input="Hello, world!", + client=client, + ) - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) # Response is processed by litellm, so just check it exists assert response is not None @@ -662,14 +659,13 @@ def test_embeddings_create_with_list_input( "send", return_value=model_response, ): - with start_transaction(name="litellm test"): - response = litellm.embedding( - model="text-embedding-ada-002", - input=["First text", "Second text", "Third text"], - client=client, - ) - # Allow time for callbacks to complete (they may run in separate threads) - time.sleep(0.1) + response = litellm.embedding( + model="text-embedding-ada-002", + input=["First text", "Second text", "Third text"], + client=client, + ) + # Allow time for callbacks to complete (they may run in separate threads) + time.sleep(0.1) # Response is processed by litellm, so just check it exists assert response is not None @@ -728,15 +724,14 @@ async def test_async_embeddings_create_with_list_input( "send", return_value=model_response, ): - with start_transaction(name="litellm test"): - response = await litellm.aembedding( - model="text-embedding-ada-002", - input=["First text", "Second text", "Third text"], - client=client, - ) + response = await litellm.aembedding( + model="text-embedding-ada-002", + input=["First text", "Second text", "Third text"], + client=client, + ) - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) # Response is processed by litellm, so just check it exists assert response is not None @@ -793,14 +788,13 @@ def test_embeddings_no_pii( "send", return_value=model_response, ): - with start_transaction(name="litellm test"): - 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) + 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) # Response is processed by litellm, so just check it exists assert response is not None @@ -851,15 +845,14 @@ async def test_async_embeddings_no_pii( "send", return_value=model_response, ): - with start_transaction(name="litellm test"): - response = await litellm.aembedding( - model="text-embedding-ada-002", - input="Hello, world!", - client=client, - ) + response = await litellm.aembedding( + model="text-embedding-ada-002", + input="Hello, world!", + client=client, + ) - await GLOBAL_LOGGING_WORKER.flush() - await asyncio.sleep(0.5) + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.5) # Response is processed by litellm, so just check it exists assert response is not None @@ -903,7 +896,7 @@ def test_exception_handling( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"), pytest.raises(litellm.RateLimitError): + ), pytest.raises(litellm.RateLimitError): litellm.completion( model="gpt-3.5-turbo", messages=messages, @@ -939,7 +932,7 @@ async def test_async_exception_handling( client.embeddings._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"), pytest.raises(litellm.RateLimitError): + ), pytest.raises(litellm.RateLimitError): await litellm.acompletion( model="gpt-3.5-turbo", messages=messages, @@ -1256,7 +1249,7 @@ def test_additional_parameters( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"): + ): litellm.completion( model="gpt-3.5-turbo", messages=messages, @@ -1327,7 +1320,7 @@ async def test_async_additional_parameters( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"): + ): await litellm.acompletion( model="gpt-3.5-turbo", messages=messages, @@ -1398,7 +1391,7 @@ def test_no_integration( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"): + ): litellm.completion( model="gpt-3.5-turbo", messages=messages, @@ -1456,7 +1449,7 @@ async def test_async_no_integration( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"): + ): await litellm.acompletion( model="gpt-3.5-turbo", messages=messages, @@ -1502,19 +1495,18 @@ def test_response_without_usage( )() items = capture_items("span") - with start_transaction(name="litellm test"): - kwargs = { - "model": "gpt-3.5-turbo", - "messages": messages, - } + kwargs = { + "model": "gpt-3.5-turbo", + "messages": messages, + } - _input_callback(kwargs) - _success_callback( - kwargs, - mock_response, - datetime.now(), - datetime.now(), - ) + _input_callback(kwargs) + _success_callback( + kwargs, + mock_response, + datetime.now(), + datetime.now(), + ) sentry_sdk.flush() (span,) = (item.payload for item in items) @@ -1593,7 +1585,7 @@ def test_binary_content_encoding_image_url( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"): + ): litellm.completion( model="gpt-4-vision-preview", messages=messages, @@ -1684,7 +1676,7 @@ async def test_async_binary_content_encoding_image_url( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"): + ): await litellm.acompletion( model="gpt-4-vision-preview", messages=messages, @@ -1777,7 +1769,7 @@ def test_binary_content_encoding_mixed_content( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"): + ): litellm.completion( model="gpt-4-vision-preview", messages=messages, @@ -1857,7 +1849,7 @@ async def test_async_binary_content_encoding_mixed_content( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"): + ): await litellm.acompletion( model="gpt-4-vision-preview", messages=messages, @@ -1935,7 +1927,7 @@ def test_binary_content_encoding_uri_type( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"): + ): litellm.completion( model="gpt-4-vision-preview", messages=messages, @@ -2020,7 +2012,7 @@ async def test_async_binary_content_encoding_uri_type( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"): + ): await litellm.acompletion( model="gpt-4-vision-preview", messages=messages, @@ -2268,7 +2260,7 @@ def test_chat_completion_data_collection( client.completions._client._client, "send", return_value=model_response, - ), start_transaction(name="litellm test"): + ): litellm.completion( model="gpt-3.5-turbo", messages=messages, @@ -2384,14 +2376,13 @@ def test_embeddings_data_collection( "send", return_value=model_response, ): - with start_transaction(name="litellm test"): - 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) + 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,) = [ From 4c4dd890a26a67fda76e6874154aee4d8b2c2715 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Fri, 4 Sep 2026 12:57:27 +0200 Subject: [PATCH 3/4] remove manual spans and stop capturing transactions --- tests/integrations/litellm/test_litellm.py | 34 ++++++++++------------ 1 file changed, 15 insertions(+), 19 deletions(-) diff --git a/tests/integrations/litellm/test_litellm.py b/tests/integrations/litellm/test_litellm.py index dd84302c50..8f34ce0a3d 100644 --- a/tests/integrations/litellm/test_litellm.py +++ b/tests/integrations/litellm/test_litellm.py @@ -182,13 +182,13 @@ def test_nonstreaming_chat_completion( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) - items = capture_items("transaction", "span") + items = capture_items("span") with mock.patch.object( client.completions._client._client, "send", return_value=model_response, - ), sentry_sdk.traces.start_span(name="litellm test"): + ): litellm.completion( model="gpt-3.5-turbo", messages=messages, @@ -198,8 +198,7 @@ def test_nonstreaming_chat_completion( litellm_utils.executor.shutdown(wait=True) sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[2]["name"] == "litellm test" + spans = [item.payload for item in items] chat_spans = list( x for x in spans @@ -282,13 +281,13 @@ async def test_async_nonstreaming_chat_completion( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) - items = capture_items("transaction", "span") + items = capture_items("span") with mock.patch.object( client.completions._client._client, "send", return_value=model_response, - ), sentry_sdk.traces.start_span(name="litellm test"): + ): await litellm.acompletion( model="gpt-3.5-turbo", messages=messages, @@ -299,8 +298,7 @@ async def test_async_nonstreaming_chat_completion( await asyncio.sleep(0.5) sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[2]["name"] == "litellm test" + spans = [item.payload for item in items] chat_spans = list( x for x in spans @@ -983,7 +981,7 @@ def test_span_origin( client.completions._client._client, "send", return_value=model_response, - ), sentry_sdk.traces.start_span(name="litellm test"): + ): litellm.completion( model="gpt-3.5-turbo", messages=messages, @@ -994,8 +992,6 @@ def test_span_origin( sentry_sdk.flush() spans = [item.payload for item in items] - assert spans[2]["is_segment"] is True - assert spans[2]["attributes"]["sentry.origin"] == "manual" assert spans[1]["attributes"]["sentry.origin"] == "auto.ai.litellm" @@ -1040,7 +1036,7 @@ def test_multiple_providers( openai_client.completions._client._client, "send", return_value=openai_model_response, - ), sentry_sdk.traces.start_span(name="test gpt-3.5-turbo"): + ): litellm.completion( model="gpt-3.5-turbo", messages=messages, @@ -1062,7 +1058,7 @@ def test_multiple_providers( anthropic_client, "post", return_value=anthropic_model_response, - ), sentry_sdk.traces.start_span(name="test claude-3-opus-20240229"): + ): litellm.completion( model="claude-3-opus-20240229", messages=messages, @@ -1084,7 +1080,7 @@ def test_multiple_providers( gemini_client, "post", return_value=gemini_model_response, - ), sentry_sdk.traces.start_span(name="test gemini/gemini-pro"): + ): litellm.completion( model="gemini/gemini-pro", messages=messages, @@ -1138,13 +1134,13 @@ async def test_async_multiple_providers( serialize_pydantic=True, request_headers={"X-Stainless-Raw-Response": "true"}, ) - items = capture_items("transaction", "span") + items = capture_items("span") with mock.patch.object( openai_client.completions._client._client, "send", return_value=openai_model_response, - ), sentry_sdk.traces.start_span(name="test gpt-3.5-turbo"): + ): await litellm.acompletion( model="gpt-3.5-turbo", messages=messages, @@ -1167,7 +1163,7 @@ async def test_async_multiple_providers( anthropic_client, "post", return_value=anthropic_model_response, - ), sentry_sdk.traces.start_span(name="test claude-3-opus-20240229"): + ): await litellm.acompletion( model="claude-3-opus-20240229", messages=messages, @@ -1190,7 +1186,7 @@ async def test_async_multiple_providers( gemini_client, "post", return_value=gemini_model_response, - ), sentry_sdk.traces.start_span(name="test gemini/gemini-pro"): + ): await litellm.acompletion( model="gemini/gemini-pro", messages=messages, @@ -1202,7 +1198,7 @@ async def test_async_multiple_providers( await asyncio.sleep(0.5) sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] + spans = [item.payload for item in items] for span in spans: if span["is_segment"] is True: continue From db62e4835d3924afbc3616170118ae00aeaa9651 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Fri, 4 Sep 2026 13:37:22 +0200 Subject: [PATCH 4/4] flush in test and stop skipping segments --- tests/integrations/litellm/test_litellm.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/integrations/litellm/test_litellm.py b/tests/integrations/litellm/test_litellm.py index 8f34ce0a3d..655d3221ca 100644 --- a/tests/integrations/litellm/test_litellm.py +++ b/tests/integrations/litellm/test_litellm.py @@ -1093,8 +1093,6 @@ def test_multiple_providers( sentry_sdk.flush() spans = [item.payload for item in items] for span in spans: - if span["is_segment"] is True: - continue # The provider should be detected by litellm.get_llm_provider assert SPANDATA.GEN_AI_SYSTEM in span["attributes"] @@ -1200,8 +1198,6 @@ async def test_async_multiple_providers( sentry_sdk.flush() spans = [item.payload for item in items] for span in spans: - if span["is_segment"] is True: - continue # The provider should be detected by litellm.get_llm_provider assert SPANDATA.GEN_AI_SYSTEM in span["attributes"] @@ -2265,6 +2261,7 @@ def test_chat_completion_data_collection( litellm_utils.executor.shutdown(wait=True) + sentry_sdk.flush() spans = [item.payload for item in items] (span,) = [ x