From 8f02fad7e7d5488a373e90a95f769bb94cd293fb Mon Sep 17 00:00:00 2001 From: VihaanAgarwal Date: Tue, 7 Jul 2026 10:23:03 -0700 Subject: [PATCH 1/2] fix(litellm): Set operation name from call type instead of chat fallback --- sentry_sdk/integrations/litellm.py | 42 +++++++----- tests/integrations/litellm/test_litellm.py | 80 ++++++++++++++++++++++ 2 files changed, 104 insertions(+), 18 deletions(-) diff --git a/sentry_sdk/integrations/litellm.py b/sentry_sdk/integrations/litellm.py index 49ead6b068..65e7a57657 100644 --- a/sentry_sdk/integrations/litellm.py +++ b/sentry_sdk/integrations/litellm.py @@ -22,7 +22,7 @@ if TYPE_CHECKING: from datetime import datetime - from typing import Any, Dict, List + from typing import Any, Dict, List, Optional, Tuple try: import litellm # type: ignore[import-not-found] @@ -35,6 +35,19 @@ # to every callback, so it lives and dies with the request. _SPAN_KEY = "_sentry_span" +# Call types whose gen_ai operation name we can determine accurately. Everything +# else gets no gen_ai.operation.name attribute, since guessing records wrong data. +_CALL_TYPE_OPERATIONS: "Dict[Any, Tuple[Optional[str], str]]" = { + "completion": ("chat", consts.OP.GEN_AI_CHAT), + "acompletion": ("chat", consts.OP.GEN_AI_CHAT), + "text_completion": ("text_completion", consts.OP.GEN_AI_TEXT_COMPLETION), + "atext_completion": ("text_completion", consts.OP.GEN_AI_TEXT_COMPLETION), + "embedding": ("embeddings", consts.OP.GEN_AI_EMBEDDINGS), + "aembedding": ("embeddings", consts.OP.GEN_AI_EMBEDDINGS), + "responses": ("responses", consts.OP.GEN_AI_RESPONSES), + "aresponses": ("responses", consts.OP.GEN_AI_RESPONSES), +} + def _store_span(kwargs: "Dict[str, Any]", span: "Any") -> None: kwargs[_SPAN_KEY] = span @@ -92,32 +105,24 @@ def _input_callback(kwargs: "Dict[str, Any]") -> None: provider = "unknown" call_type = kwargs.get("call_type", None) - if call_type == "embedding" or call_type == "aembedding": - operation = "embeddings" - else: - operation = "chat" + operation, span_op = _CALL_TYPE_OPERATIONS.get( + call_type, (None, consts.OP.GEN_AI_CHAT) + ) + span_name = f"{operation or call_type or 'unknown'} {model}" # Start a new span/transaction if has_span_streaming_enabled(client.options): span = sentry_sdk.traces.start_span( - name=f"{operation} {model}", + name=span_name, attributes={ - "sentry.op": ( - consts.OP.GEN_AI_CHAT - if operation == "chat" - else consts.OP.GEN_AI_EMBEDDINGS - ), + "sentry.op": span_op, "sentry.origin": LiteLLMIntegration.origin, }, ) else: span = get_start_span_function()( - op=( - consts.OP.GEN_AI_CHAT - if operation == "chat" - else consts.OP.GEN_AI_EMBEDDINGS - ), - name=f"{operation} {model}", + op=span_op, + name=span_name, origin=LiteLLMIntegration.origin, ) span.__enter__() @@ -126,7 +131,8 @@ def _input_callback(kwargs: "Dict[str, Any]") -> None: # Set basic data set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, provider) - set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, operation) + if operation is not None: + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, operation) # Record input/messages if allowed if should_send_default_pii() and integration.include_prompts: diff --git a/tests/integrations/litellm/test_litellm.py b/tests/integrations/litellm/test_litellm.py index 8d0a75d834..42de87844d 100644 --- a/tests/integrations/litellm/test_litellm.py +++ b/tests/integrations/litellm/test_litellm.py @@ -2477,6 +2477,7 @@ def test_response_without_usage( kwargs = { "model": "gpt-3.5-turbo", "messages": messages, + "call_type": "completion", } _input_callback(kwargs) @@ -2500,6 +2501,7 @@ def test_response_without_usage( kwargs = { "model": "gpt-3.5-turbo", "messages": messages, + "call_type": "completion", } _input_callback(kwargs) @@ -3415,3 +3417,81 @@ def test_convert_message_parts_image_url_missing_url(): converted = _convert_message_parts(messages) # Should return item unchanged assert converted[0]["content"][0]["type"] == "image_url" + + +@pytest.mark.parametrize( + "call_type,expected_operation,expected_op", + [ + ("completion", "chat", OP.GEN_AI_CHAT), + ("acompletion", "chat", OP.GEN_AI_CHAT), + ("text_completion", "text_completion", OP.GEN_AI_TEXT_COMPLETION), + ("atext_completion", "text_completion", OP.GEN_AI_TEXT_COMPLETION), + ("embedding", "embeddings", OP.GEN_AI_EMBEDDINGS), + ("aembedding", "embeddings", OP.GEN_AI_EMBEDDINGS), + ("responses", "responses", OP.GEN_AI_RESPONSES), + ("aresponses", "responses", OP.GEN_AI_RESPONSES), + ], +) +def test_operation_name_mapped_from_call_type( + sentry_init, capture_events, call_type, expected_operation, expected_op +): + """Known call types map to their actual operation, not the chat fallback.""" + sentry_init( + integrations=[LiteLLMIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=False, + ) + events = capture_events() + + with start_transaction(name="litellm test"): + kwargs = { + "model": "gpt-3.5-turbo", + "call_type": call_type, + } + + _input_callback(kwargs) + _success_callback( + kwargs, + MockCompletionResponse(), + datetime.now(), + datetime.now(), + ) + + (tx,) = events + (span,) = [s for s in tx["spans"] if s["origin"] == "auto.ai.litellm"] + + assert span["op"] == expected_op + assert span["description"] == f"{expected_operation} gpt-3.5-turbo" + assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == expected_operation + + +def test_operation_name_not_set_for_unknown_call_type(sentry_init, capture_events): + """Call types with no accurate operation name get none, instead of "chat".""" + sentry_init( + integrations=[LiteLLMIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=False, + ) + events = capture_events() + + with start_transaction(name="litellm test"): + kwargs = { + "model": "dall-e-3", + "call_type": "image_generation", + } + + _input_callback(kwargs) + _success_callback( + kwargs, + MockCompletionResponse(model="dall-e-3"), + datetime.now(), + datetime.now(), + ) + + (tx,) = events + (span,) = [s for s in tx["spans"] if s["origin"] == "auto.ai.litellm"] + + assert span["description"] == "image_generation dall-e-3" + assert SPANDATA.GEN_AI_OPERATION_NAME not in span["data"] From b57b7f9a86734bd595f5884c08af4afcd70a02a5 Mon Sep 17 00:00:00 2001 From: VihaanAgarwal Date: Fri, 10 Jul 2026 13:15:42 -0700 Subject: [PATCH 2/2] Skip instrumentation for unknown call types and test through the litellm API --- sentry_sdk/integrations/litellm.py | 21 +-- tests/integrations/litellm/test_litellm.py | 168 +++++++++++++++------ 2 files changed, 129 insertions(+), 60 deletions(-) diff --git a/sentry_sdk/integrations/litellm.py b/sentry_sdk/integrations/litellm.py index 65e7a57657..e97c1817d8 100644 --- a/sentry_sdk/integrations/litellm.py +++ b/sentry_sdk/integrations/litellm.py @@ -22,7 +22,7 @@ if TYPE_CHECKING: from datetime import datetime - from typing import Any, Dict, List, Optional, Tuple + from typing import Any, Dict, List, Tuple try: import litellm # type: ignore[import-not-found] @@ -36,8 +36,8 @@ _SPAN_KEY = "_sentry_span" # Call types whose gen_ai operation name we can determine accurately. Everything -# else gets no gen_ai.operation.name attribute, since guessing records wrong data. -_CALL_TYPE_OPERATIONS: "Dict[Any, Tuple[Optional[str], str]]" = { +# else is not instrumented, since guessing records wrong data. +_CALL_TYPE_OPERATIONS: "Dict[Any, Tuple[str, str]]" = { "completion": ("chat", consts.OP.GEN_AI_CHAT), "acompletion": ("chat", consts.OP.GEN_AI_CHAT), "text_completion": ("text_completion", consts.OP.GEN_AI_TEXT_COMPLETION), @@ -96,6 +96,12 @@ def _input_callback(kwargs: "Dict[str, Any]") -> None: if integration is None: return + call_type = kwargs.get("call_type", None) + if call_type not in _CALL_TYPE_OPERATIONS: + return + + operation, span_op = _CALL_TYPE_OPERATIONS[call_type] + # Get key parameters full_model = kwargs.get("model", "") try: @@ -104,11 +110,7 @@ def _input_callback(kwargs: "Dict[str, Any]") -> None: model = full_model provider = "unknown" - call_type = kwargs.get("call_type", None) - operation, span_op = _CALL_TYPE_OPERATIONS.get( - call_type, (None, consts.OP.GEN_AI_CHAT) - ) - span_name = f"{operation or call_type or 'unknown'} {model}" + span_name = f"{operation} {model}" # Start a new span/transaction if has_span_streaming_enabled(client.options): @@ -131,8 +133,7 @@ def _input_callback(kwargs: "Dict[str, Any]") -> None: # Set basic data set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, provider) - if operation is not None: - set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, operation) + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, operation) # Record input/messages if allowed if should_send_default_pii() and integration.include_prompts: diff --git a/tests/integrations/litellm/test_litellm.py b/tests/integrations/litellm/test_litellm.py index 42de87844d..8faee4599e 100644 --- a/tests/integrations/litellm/test_litellm.py +++ b/tests/integrations/litellm/test_litellm.py @@ -34,7 +34,8 @@ async def __call__(self, *args, **kwargs): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from openai import AsyncOpenAI, OpenAI -from openai.types import CompletionUsage +from openai.types import Completion, CompletionUsage, Image, ImagesResponse +from openai.types.completion_choice import CompletionChoice from sentry_sdk import start_transaction from sentry_sdk._types import BLOB_DATA_SUBSTITUTE @@ -2561,6 +2562,7 @@ def test_litellm_message_truncation(sentry_init, capture_events): kwargs = { "model": "gpt-3.5-turbo", "messages": messages, + "call_type": "completion", } _input_callback(kwargs) @@ -3419,23 +3421,13 @@ def test_convert_message_parts_image_url_missing_url(): assert converted[0]["content"][0]["type"] == "image_url" -@pytest.mark.parametrize( - "call_type,expected_operation,expected_op", - [ - ("completion", "chat", OP.GEN_AI_CHAT), - ("acompletion", "chat", OP.GEN_AI_CHAT), - ("text_completion", "text_completion", OP.GEN_AI_TEXT_COMPLETION), - ("atext_completion", "text_completion", OP.GEN_AI_TEXT_COMPLETION), - ("embedding", "embeddings", OP.GEN_AI_EMBEDDINGS), - ("aembedding", "embeddings", OP.GEN_AI_EMBEDDINGS), - ("responses", "responses", OP.GEN_AI_RESPONSES), - ("aresponses", "responses", OP.GEN_AI_RESPONSES), - ], -) -def test_operation_name_mapped_from_call_type( - sentry_init, capture_events, call_type, expected_operation, expected_op +def test_text_completion_operation_name( + sentry_init, + capture_events, + get_model_response, + reset_litellm_executor, ): - """Known call types map to their actual operation, not the chat fallback.""" + """text_completion calls get the text_completion op, not the chat fallback.""" sentry_init( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], @@ -3444,30 +3436,56 @@ def test_operation_name_mapped_from_call_type( ) events = capture_events() - with start_transaction(name="litellm test"): - kwargs = { - "model": "gpt-3.5-turbo", - "call_type": call_type, - } + client = OpenAI(api_key="test-key") - _input_callback(kwargs) - _success_callback( - kwargs, - MockCompletionResponse(), - datetime.now(), - datetime.now(), + model_response = get_model_response( + Completion( + id="cmpl-test", + choices=[ + CompletionChoice(finish_reason="stop", index=0, text="Test response") + ], + created=1234567890, + model="gpt-3.5-turbo-instruct", + object="text_completion", + usage=CompletionUsage( + prompt_tokens=10, + completion_tokens=20, + total_tokens=30, + ), + ), + serialize_pydantic=True, + request_headers={"X-Stainless-Raw-Response": "true"}, + ) + + with mock.patch.object( + client.completions._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + litellm.text_completion( + model="gpt-3.5-turbo-instruct", + prompt="Hello!", + client=client, ) - (tx,) = events - (span,) = [s for s in tx["spans"] if s["origin"] == "auto.ai.litellm"] + litellm_utils.executor.shutdown(wait=True) + + (event,) = events + (span,) = [s for s in event["spans"] if s["origin"] == "auto.ai.litellm"] - assert span["op"] == expected_op - assert span["description"] == f"{expected_operation} gpt-3.5-turbo" - assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == expected_operation + assert span["op"] == OP.GEN_AI_TEXT_COMPLETION + assert span["description"] == "text_completion gpt-3.5-turbo-instruct" + assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "text_completion" -def test_operation_name_not_set_for_unknown_call_type(sentry_init, capture_events): - """Call types with no accurate operation name get none, instead of "chat".""" +def test_responses_operation_name( + sentry_init, + capture_events, + get_model_response, + nonstreaming_responses_model_response, + reset_litellm_executor, +): + """Responses API calls get the responses op, not the chat fallback.""" sentry_init( integrations=[LiteLLMIntegration()], disabled_integrations=[StdlibIntegration], @@ -3476,22 +3494,72 @@ def test_operation_name_not_set_for_unknown_call_type(sentry_init, capture_event ) events = capture_events() - with start_transaction(name="litellm test"): - kwargs = { - "model": "dall-e-3", - "call_type": "image_generation", - } + client = HTTPHandler() - _input_callback(kwargs) - _success_callback( - kwargs, - MockCompletionResponse(model="dall-e-3"), - datetime.now(), - datetime.now(), + model_response = get_model_response( + nonstreaming_responses_model_response, + serialize_pydantic=True, + ) + + with mock.patch.object( + client, + "post", + return_value=model_response, + ), start_transaction(name="litellm test"): + litellm.responses( + model="gpt-4", + input="Hello!", + client=client, + api_key="test-key", + ) + + litellm_utils.executor.shutdown(wait=True) + + (event,) = events + (span,) = [s for s in event["spans"] if s["origin"] == "auto.ai.litellm"] + + assert span["op"] == OP.GEN_AI_RESPONSES + assert span["description"] == "responses gpt-4" + assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "responses" + + +def test_unknown_call_type_is_not_instrumented( + sentry_init, + capture_events, + get_model_response, + reset_litellm_executor, +): + """Call types with no accurate operation name emit no span at all.""" + sentry_init( + integrations=[LiteLLMIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=False, + ) + events = capture_events() + + client = OpenAI(api_key="test-key") + + model_response = get_model_response( + ImagesResponse( + created=1234567890, + data=[Image(url="https://example.com/image.png")], + ), + serialize_pydantic=True, + ) + + with mock.patch.object( + client.images._client._client, + "send", + return_value=model_response, + ), start_transaction(name="litellm test"): + litellm.image_generation( + model="dall-e-3", + prompt="A cat", + client=client, ) - (tx,) = events - (span,) = [s for s in tx["spans"] if s["origin"] == "auto.ai.litellm"] + litellm_utils.executor.shutdown(wait=True) - assert span["description"] == "image_generation dall-e-3" - assert SPANDATA.GEN_AI_OPERATION_NAME not in span["data"] + (event,) = events + assert [s for s in event["spans"] if s["origin"] == "auto.ai.litellm"] == []