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 thread
alexander-alderman-webb marked this conversation as resolved.

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__()
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 Down
108 changes: 107 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
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,105 @@ 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"