Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 24 additions & 18 deletions sentry_sdk/integrations/litellm.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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
Expand Down Expand Up @@ -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}"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No data is better than wrong data, so please exit early if the key is not in the dictionary.


# 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,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)
Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit b57b7f9. Configure here.

origin=LiteLLMIntegration.origin,
)
span.__enter__()
Comment on lines 125 to 130

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: The _input_callback fails to capture input for text_completion and responses operations because it incorrectly looks for a messages key instead of prompt or input respectively.
Severity: MEDIUM

Suggested Fix

Update _input_callback to handle text_completion and responses operations. Add branches to check the operation type and read the input data from the correct keyword argument: prompt for text_completion and input for responses.

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 in `_input_callback` does not correctly handle
the `text_completion` and `responses` operation types. The code currently branches only
for the `embeddings` operation and defaults to reading from a `messages` key for all
other operations. However, `text_completion` operations use a `prompt` key and
`responses` operations use an `input` key. This discrepancy causes the function to fail
to capture input data for these two new operation types, even when `include_prompts` is
enabled, leading to silent data loss for users expecting this information to be
recorded.

Expand All @@ -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:
Expand Down
80 changes: 80 additions & 0 deletions tests/integrations/litellm/test_litellm.py
Original file line number Diff line number Diff line change
Expand Up @@ -2477,6 +2477,7 @@ def test_response_without_usage(
kwargs = {
"model": "gpt-3.5-turbo",
"messages": messages,
"call_type": "completion",
}

_input_callback(kwargs)
Expand All @@ -2500,6 +2501,7 @@ def test_response_without_usage(
kwargs = {
"model": "gpt-3.5-turbo",
"messages": messages,
"call_type": "completion",
}

_input_callback(kwargs)
Expand Down Expand Up @@ -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(),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

)

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