fix(litellm): Set operation name from call type instead of chat fallback

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:

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"]