From 8f02fad7e7d5488a373e90a95f769bb94cd293fb Mon Sep 17 00:00:00 2001 From: VihaanAgarwal Date: Tue, 7 Jul 2026 10:23:03 -0700 Subject: [PATCH] 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"]