diff --git a/sentry_sdk/integrations/litellm.py b/sentry_sdk/integrations/litellm.py index 49ead6b068..e97c1817d8 100644 --- a/sentry_sdk/integrations/litellm.py +++ b/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, 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 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 @@ -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 + + operation, span_op = _CALL_TYPE_OPERATIONS[call_type] + # Get key parameters full_model = kwargs.get("model", "") try: @@ -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, origin=LiteLLMIntegration.origin, ) span.__enter__() diff --git a/tests/integrations/litellm/test_litellm.py b/tests/integrations/litellm/test_litellm.py index 8d0a75d834..8faee4599e 100644 --- a/tests/integrations/litellm/test_litellm.py +++ b/tests/integrations/litellm/test_litellm.py @@ -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 @@ -2477,6 +2478,7 @@ def test_response_without_usage( kwargs = { "model": "gpt-3.5-turbo", "messages": messages, + "call_type": "completion", } _input_callback(kwargs) @@ -2500,6 +2502,7 @@ def test_response_without_usage( kwargs = { "model": "gpt-3.5-turbo", "messages": messages, + "call_type": "completion", } _input_callback(kwargs) @@ -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) @@ -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"] == []