Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
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
43 changes: 25 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, 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 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),
"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 @@ -83,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
Comment on lines +99 to +101

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 returns early for unmapped call_types, preventing span creation and disabling exception capture for litellm calls like image_generation.
Severity: MEDIUM

Suggested Fix

To restore exception capturing for all litellm call types, ensure a span is created even for unmapped types. Consider creating a fallback span with a generic operation, similar to the previous behavior, instead of returning early from _input_callback when a call_type is not found.

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#L99-L101

Potential issue: The `_input_callback` returns early for any `litellm` `call_type` not
in `_CALL_TYPE_OPERATIONS`, such as `image_generation` or `speech`. This prevents a
Sentry span from being created for these calls. As a result, if an exception occurs
during one of these unmapped calls, the `_failure_callback` is unable to find a span to
process. It returns immediately, and the exception is not captured by Sentry. This is a
regression that silently disables error reporting for previously-instrumented call
types.

Did we get this right? 👍 / 👎 to inform future reviews.


operation, span_op = _CALL_TYPE_OPERATIONS[call_type]

# Get key parameters
full_model = kwargs.get("model", "")
try:
Expand All @@ -91,33 +110,21 @@ def _input_callback(kwargs: "Dict[str, Any]") -> None:
model = full_model
provider = "unknown"

call_type = kwargs.get("call_type", None)
if call_type == "embedding" or call_type == "aembedding":
operation = "embeddings"
else:
operation = "chat"
span_name = f"{operation} {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,

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

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

_input_callback(kwargs)
Expand Down Expand Up @@ -2559,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)
Expand Down Expand Up @@ -3415,3 +3419,147 @@ 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"


def test_text_completion_operation_name(
sentry_init,
capture_events,
get_model_response,
reset_litellm_executor,
):
"""text_completion calls get the text_completion op, not the chat fallback."""
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(
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,
)

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_TEXT_COMPLETION
assert span["description"] == "text_completion gpt-3.5-turbo-instruct"
assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "text_completion"


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],
traces_sample_rate=1.0,
stream_gen_ai_spans=False,
)
events = capture_events()

client = HTTPHandler()

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

litellm_utils.executor.shutdown(wait=True)

(event,) = events
assert [s for s in event["spans"] if s["origin"] == "auto.ai.litellm"] == []