diff --git a/sentry_sdk/integrations/litellm.py b/sentry_sdk/integrations/litellm.py index 08cb217962..1e3f18ddf6 100644 --- a/sentry_sdk/integrations/litellm.py +++ b/sentry_sdk/integrations/litellm.py @@ -14,7 +14,7 @@ from sentry_sdk.utils import event_from_exception if TYPE_CHECKING: - from typing import Any, Dict + from typing import Any, Dict, List from datetime import datetime try: @@ -35,6 +35,68 @@ def _get_metadata_dict(kwargs: "Dict[str, Any]") -> "Dict[str, Any]": return metadata +def _convert_message_parts(messages: "List[Dict[str, Any]]") -> "List[Dict[str, Any]]": + """ + Convert the message parts from OpenAI format to the `gen_ai.request.messages` format. + e.g: + { + "role": "user", + "content": [ + { + "text": "How many ponies do you see in the image?", + "type": "text" + }, + { + "type": "image_url", + "image_url": { + "url": "data:image/jpeg;base64,...", + "detail": "high" + } + } + ] + } + becomes: + { + "role": "user", + "content": [ + { + "text": "How many ponies do you see in the image?", + "type": "text" + }, + { + "type": "blob", + "modality": "image", + "mime_type": "image/jpeg", + "content": "data:image/jpeg;base64,..." + } + ] + } + """ + + def _map_item(item: "Dict[str, Any]") -> "Dict[str, Any]": + if item.get("type") == "image_url": + image_url = item.get("image_url") or {} + if image_url.get("url", "").startswith("data:"): + return { + "type": "blob", + "modality": "image", + "mime_type": item["image_url"]["url"].split(";base64,")[0], + "content": item["image_url"]["url"].split(";base64,")[1], + } + else: + return { + "type": "uri", + "uri": item["image_url"]["url"], + } + return item + + for message in messages: + content = message.get("content") + if isinstance(content, list): + message["content"] = [_map_item(item) for item in content] + return messages + + def _input_callback(kwargs: "Dict[str, Any]") -> None: """Handle the start of a request.""" integration = sentry_sdk.get_client().get_integration(LiteLLMIntegration) @@ -101,6 +163,7 @@ def _input_callback(kwargs: "Dict[str, Any]") -> None: messages = kwargs.get("messages", []) if messages: scope = sentry_sdk.get_current_scope() + messages = _convert_message_parts(messages) messages_data = truncate_and_annotate_messages(messages, span, scope) if messages_data is not None: set_data_normalized( diff --git a/tests/integrations/litellm/test_litellm.py b/tests/integrations/litellm/test_litellm.py index 1b925fb61f..4f81c54b63 100644 --- a/tests/integrations/litellm/test_litellm.py +++ b/tests/integrations/litellm/test_litellm.py @@ -1,3 +1,4 @@ +import base64 import json import pytest import time @@ -23,6 +24,7 @@ async def __call__(self, *args, **kwargs): from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations.litellm import ( LiteLLMIntegration, + _convert_message_parts, _input_callback, _success_callback, _failure_callback, @@ -753,3 +755,162 @@ def test_litellm_message_truncation(sentry_init, capture_events): assert "small message 4" in str(parsed_messages[0]) assert "small message 5" in str(parsed_messages[1]) assert tx["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 5 + + +IMAGE_DATA = b"fake_image_data_12345" +IMAGE_B64 = base64.b64encode(IMAGE_DATA).decode("utf-8") +IMAGE_DATA_URI = f"data:image/png;base64,{IMAGE_B64}" + + +def test_binary_content_encoding_image_url(sentry_init, capture_events): + sentry_init( + integrations=[LiteLLMIntegration(include_prompts=True)], + traces_sample_rate=1.0, + send_default_pii=True, + ) + events = capture_events() + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Look at this image:"}, + { + "type": "image_url", + "image_url": {"url": IMAGE_DATA_URI, "detail": "high"}, + }, + ], + } + ] + mock_response = MockCompletionResponse() + + with start_transaction(name="litellm test"): + kwargs = {"model": "gpt-4-vision-preview", "messages": messages} + _input_callback(kwargs) + _success_callback(kwargs, mock_response, datetime.now(), datetime.now()) + + (event,) = events + (span,) = event["spans"] + messages_data = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + + blob_item = next( + ( + item + for msg in messages_data + if "content" in msg + for item in msg["content"] + if item.get("type") == "blob" + ), + None, + ) + assert blob_item is not None + assert blob_item["modality"] == "image" + assert blob_item["mime_type"] == "data:image/png" + assert IMAGE_B64 in blob_item["content"] or "[Filtered]" in str( + blob_item["content"] + ) + + +def test_binary_content_encoding_mixed_content(sentry_init, capture_events): + sentry_init( + integrations=[LiteLLMIntegration(include_prompts=True)], + traces_sample_rate=1.0, + send_default_pii=True, + ) + events = capture_events() + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Here is an image:"}, + { + "type": "image_url", + "image_url": {"url": IMAGE_DATA_URI}, + }, + {"type": "text", "text": "What do you see?"}, + ], + } + ] + mock_response = MockCompletionResponse() + + with start_transaction(name="litellm test"): + kwargs = {"model": "gpt-4-vision-preview", "messages": messages} + _input_callback(kwargs) + _success_callback(kwargs, mock_response, datetime.now(), datetime.now()) + + (event,) = events + (span,) = event["spans"] + messages_data = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + + content_items = [ + item for msg in messages_data if "content" in msg for item in msg["content"] + ] + assert any(item.get("type") == "text" for item in content_items) + assert any(item.get("type") == "blob" for item in content_items) + + +def test_binary_content_encoding_uri_type(sentry_init, capture_events): + sentry_init( + integrations=[LiteLLMIntegration(include_prompts=True)], + traces_sample_rate=1.0, + send_default_pii=True, + ) + events = capture_events() + + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.jpg"}, + } + ], + } + ] + mock_response = MockCompletionResponse() + + with start_transaction(name="litellm test"): + kwargs = {"model": "gpt-4-vision-preview", "messages": messages} + _input_callback(kwargs) + _success_callback(kwargs, mock_response, datetime.now(), datetime.now()) + + (event,) = events + (span,) = event["spans"] + messages_data = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + + uri_item = next( + ( + item + for msg in messages_data + if "content" in msg + for item in msg["content"] + if item.get("type") == "uri" + ), + None, + ) + assert uri_item is not None + assert uri_item["uri"] == "https://example.com/image.jpg" + + +def test_convert_message_parts_direct(): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Hello"}, + { + "type": "image_url", + "image_url": {"url": IMAGE_DATA_URI}, + }, + ], + } + ] + converted = _convert_message_parts(messages) + blob_item = next( + item for item in converted[0]["content"] if item.get("type") == "blob" + ) + assert blob_item["modality"] == "image" + assert blob_item["mime_type"] == "data:image/png" + assert IMAGE_B64 in blob_item["content"]