fix(litellm): Set operation name from call type instead of chat fallback - #6792
fix(litellm): Set operation name from call type instead of chat fallback#6792VihaanAgarwal wants to merge 4 commits into
Conversation
| 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}" |
There was a problem hiding this comment.
No data is better than wrong data, so please exit early if the key is not in the dictionary.
| _input_callback(kwargs) | ||
| _success_callback( | ||
| kwargs, | ||
| MockCompletionResponse(), | ||
| datetime.now(), | ||
| datetime.now(), |
There was a problem hiding this comment.
For tests, the expectation is that you use the library like a user would and assert that the correct telemetry was emitted.
The existing tests have prior art for how to achieve this.
|
Fair points, both addressed. Unknown call types now bail before any span is created, so nothing gets recorded for them. Also replaced the callback-level tests with ones that go through litellm.text_completion / litellm.responses / litellm.image_generation with mocked transports and assert on the captured spans. |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit b57b7f9. Configure here.
| ), | ||
| name=f"{operation} {model}", | ||
| op=span_op, | ||
| name=span_name, |
There was a problem hiding this comment.
New ops skip prompt capture
Medium Severity
_CALL_TYPE_OPERATIONS now yields text_completion and responses, but prompt recording still only special-cases embeddings and otherwise reads messages. Those call types take prompt and input, so with include_prompts enabled their inputs are never attached to the span.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit b57b7f9. Configure here.
![sentry[bot]](https://avatars.githubusercontent.com/in/12637?s=60&v=4){kind=small}
| 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.""" |
There was a problem hiding this comment.
Please do not test a negative.
We test that telemetry is emitted for certain operations, and not that telemetry doesn't exist sometimes.
There was a problem hiding this comment.
Makes sense, dropped it (and the now-unused image imports). The text_completion and responses tests still cover the mapping.
|
Also extended the two new tests to assert the prompt is recorded, since the review bots flagged that Anything else you'd like changed here? |
| 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__() |
There was a problem hiding this comment.
Bug: Input prompts for litellm.text_completion and litellm.responses are not captured because the code only looks for a messages argument, not prompt or input.
Severity: MEDIUM
Suggested Fix
Update the input capture logic to check for kwargs.get("prompt") and kwargs.get("input") in addition to kwargs.get("messages"). This will ensure the input content is correctly extracted for text_completion and responses calls.
Prompt for AI Agent
Review the code at the location below. A potential bug has been identified by an AI
agent. Verify if this is a real issue. If it is, propose a fix; if not, explain why it's
not valid.
Location: sentry_sdk/integrations/litellm.py#L125-L130
Potential issue: The input capture logic for the LiteLLM integration only checks for a
`messages` keyword argument to record the input prompt. However, the
`litellm.text_completion` function uses a `prompt` argument, and `litellm.responses`
uses an `input` argument. Because the implementation does not check for these
alternative argument names, the input prompts for these specific function calls will not
be captured and attached to the transaction span data. This represents a regression that
will cause existing tests, which assert the presence of this data, to fail.
2 participants
Fixes #6442.
The input callback mapped every call type except
embedding/aembeddingtochat, so alitellm.image_generation()call showed up in Sentry aschat dall-e-3withgen_ai.operation.name: chat. Same story for text completions, speech, transcription, rerank and the rest.This replaces the binary fallback with an explicit table of the call types we can name accurately:
completion/acompletion->chattext_completion/atext_completion->text_completionembedding/aembedding->embeddingsresponses/aresponses->responsesEach maps to its matching
OP.GEN_AI_*span op (the constants fortext_completionandresponsesalready exist and are used by the langchain and openai integrations).For everything else,
gen_ai.operation.nameis simply not set. Recording a guess is what this issue is about, and litellm has 140+ call types (batch, files, assistants, video...), most of which have no semconv equivalent. The raw call type still ends up in the span name (image_generation dall-e-3), so the spans stay identifiable. I keptgen_ai.chatas the span op fallback for unmapped types so those spans keep flowing into the AI views like before; happy to change that if you'd rather they drop out.One test tweak:
test_response_without_usagebuilt callback kwargs without acall_typeand asserted the span was namedchat ..., which only held because of the fallback. Real callbacks always carrycall_type(litellm'sLogging.pre_callsets it), so I added"call_type": "completion"to those kwargs rather than special casing a missing key as chat.Tests cover all eight mapped call types plus the unknown case. The litellm test module passes locally except
test_multiple_providers, which errors at setup on current master too (missing optional provider package in my env).