← Back
Editing: litellm.py
import copy from typing import TYPE_CHECKING import sentry_sdk from sentry_sdk import consts from sentry_sdk.ai.monitoring import record_token_usage from sentry_sdk.ai.utils import ( get_start_span_function, set_data_normalized, transform_openai_content_part, truncate_and_annotate_embedding_inputs, truncate_and_annotate_messages, ) from sentry_sdk.consts import SPANDATA from sentry_sdk.integrations import DidNotEnable, Integration from sentry_sdk.scope import should_send_default_pii from sentry_sdk.tracing_utils import ( has_span_streaming_enabled, should_truncate_gen_ai_input, ) from sentry_sdk.utils import event_from_exception if TYPE_CHECKING: from datetime import datetime from typing import Any, Dict, List try: import litellm # type: ignore[import-not-found] from litellm import failure_callback, input_callback, success_callback except ImportError: raise DidNotEnable("LiteLLM not installed") def _get_metadata_dict(kwargs: "Dict[str, Any]") -> "Dict[str, Any]": """Get the metadata dictionary from the kwargs.""" litellm_params = kwargs.setdefault("litellm_params", {}) # we need this weird little dance, as metadata might be set but may be None initially metadata = litellm_params.get("metadata") if metadata is None: metadata = {} litellm_params["metadata"] = metadata 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 using the OpenAI-specific transformer (LiteLLM uses OpenAI's message format). Deep copies messages to avoid mutating original kwargs. """ # Deep copy to avoid mutating original messages from kwargs messages = copy.deepcopy(messages) for message in messages: if not isinstance(message, dict): continue content = message.get("content") if isinstance(content, (list, tuple)): transformed = [] for item in content: if isinstance(item, dict): result = transform_openai_content_part(item) # If transformation succeeded, use the result; otherwise keep original transformed.append(result if result is not None else item) else: transformed.append(item) message["content"] = transformed return messages def _input_callback(kwargs: "Dict[str, Any]") -> None: """Handle the start of a request.""" client = sentry_sdk.get_client() integration = client.get_integration(LiteLLMIntegration) if integration is None: return # Get key parameters full_model = kwargs.get("model", "") try: model, provider, _, _ = litellm.get_llm_provider(full_model) except Exception: 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" # Start a new span/transaction if has_span_streaming_enabled(client.options): span = sentry_sdk.traces.start_span( name=f"{operation} {model}", attributes={ "sentry.op": ( consts.OP.GEN_AI_CHAT if operation == "chat" else consts.OP.GEN_AI_EMBEDDINGS ), "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}", origin=LiteLLMIntegration.origin, ) span.__enter__() # Store span for later _get_metadata_dict(kwargs)["_sentry_span"] = span # Set basic data set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, provider) set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, operation) # Record input/messages if allowed if should_send_default_pii() and integration.include_prompts: if operation == "embeddings": # For embeddings, look for the 'input' parameter embedding_input = kwargs.get("input") if embedding_input: scope = sentry_sdk.get_current_scope() # Normalize to list format input_list = ( embedding_input if isinstance(embedding_input, list) else [embedding_input] ) client = sentry_sdk.get_client() messages_data = ( truncate_and_annotate_embedding_inputs(input_list, span, scope) if should_truncate_gen_ai_input(client.options) else input_list ) if messages_data is not None: set_data_normalized( span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, messages_data, unpack=False, ) else: # For chat, look for the 'messages' parameter messages = kwargs.get("messages", []) if messages: client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() messages = _convert_message_parts(messages) messages_data = ( truncate_and_annotate_messages(messages, span, scope) if should_truncate_gen_ai_input(client.options) else messages ) if messages_data is not None: set_data_normalized( span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False, ) # Record other parameters params = { "model": SPANDATA.GEN_AI_REQUEST_MODEL, "stream": SPANDATA.GEN_AI_RESPONSE_STREAMING, "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE, "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P, } for key, attribute in params.items(): value = kwargs.get(key) if value is not None: set_data_normalized(span, attribute, value) async def _async_input_callback(kwargs: "Dict[str, Any]") -> None: return _input_callback(kwargs) def _success_callback( kwargs: "Dict[str, Any]", completion_response: "Any", start_time: "datetime", end_time: "datetime", ) -> None: """Handle successful completion.""" metadata = _get_metadata_dict(kwargs) span = metadata.get("_sentry_span") if span is None: return integration = sentry_sdk.get_client().get_integration(LiteLLMIntegration) if integration is None: return try: # Record model information if hasattr(completion_response, "model"): set_data_normalized( span, SPANDATA.GEN_AI_RESPONSE_MODEL, completion_response.model ) # Record response content if allowed if should_send_default_pii() and integration.include_prompts: if hasattr(completion_response, "choices"): response_messages = [] for choice in completion_response.choices: if hasattr(choice, "message"): if hasattr(choice.message, "model_dump"): response_messages.append(choice.message.model_dump()) elif hasattr(choice.message, "dict"): response_messages.append(choice.message.dict()) else: # Fallback for basic message objects msg = {} if hasattr(choice.message, "role"): msg["role"] = choice.message.role if hasattr(choice.message, "content"): msg["content"] = choice.message.content if hasattr(choice.message, "tool_calls"): msg["tool_calls"] = choice.message.tool_calls response_messages.append(msg) if response_messages: set_data_normalized( span, SPANDATA.GEN_AI_RESPONSE_TEXT, response_messages ) # Record token usage if hasattr(completion_response, "usage"): usage = completion_response.usage record_token_usage( span, input_tokens=getattr(usage, "prompt_tokens", None), output_tokens=getattr(usage, "completion_tokens", None), total_tokens=getattr(usage, "total_tokens", None), ) finally: is_streaming = kwargs.get("stream") # Callback is fired multiple times when streaming a response. # Streaming flag checked at https://github.com/BerriAI/litellm/blob/33c3f13443eaf990ac8c6e3da78bddbc2b7d0e7a/litellm/litellm_core_utils/litellm_logging.py#L1603 if ( is_streaming is not True or "complete_streaming_response" in kwargs or "async_complete_streaming_response" in kwargs ): span = metadata.pop("_sentry_span", None) if span is not None: span.__exit__(None, None, None) async def _async_success_callback( kwargs: "Dict[str, Any]", completion_response: "Any", start_time: "datetime", end_time: "datetime", ) -> None: return _success_callback( kwargs, completion_response, start_time, end_time, ) def _failure_callback( kwargs: "Dict[str, Any]", exception: Exception, start_time: "datetime", end_time: "datetime", ) -> None: """Handle request failure.""" span = _get_metadata_dict(kwargs).get("_sentry_span") if span is None: return try: # Capture the exception event, hint = event_from_exception( exception, client_options=sentry_sdk.get_client().options, mechanism={"type": "litellm", "handled": False}, ) sentry_sdk.capture_event(event, hint=hint) finally: # Always finish the span and clean up span.__exit__(type(exception), exception, None) class LiteLLMIntegration(Integration): """ LiteLLM integration for Sentry. This integration automatically captures LiteLLM API calls and sends them to Sentry for monitoring and error tracking. It supports all 100+ LLM providers that LiteLLM supports, including OpenAI, Anthropic, Google, Cohere, and many others. Features: - Automatic exception capture for all LiteLLM calls - Token usage tracking across all providers - Provider detection and attribution - Input/output message capture (configurable) - Streaming response support - Cost tracking integration Usage: ```python import litellm import sentry_sdk # Initialize Sentry with the LiteLLM integration sentry_sdk.init( dsn="your-dsn", send_default_pii=True integrations=[ sentry_sdk.integrations.LiteLLMIntegration( include_prompts=True # Set to False to exclude message content ) ] ) # All LiteLLM calls will now be monitored response = litellm.completion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello!"}] ) ``` Configuration: - include_prompts (bool): Whether to include prompts and responses in spans. Defaults to True. Set to False to exclude potentially sensitive data. """ identifier = "litellm" origin = f"auto.ai.{identifier}" def __init__(self: "LiteLLMIntegration", include_prompts: bool = True) -> None: self.include_prompts = include_prompts @staticmethod def setup_once() -> None: """Set up LiteLLM callbacks for monitoring.""" litellm.input_callback = input_callback or [] if _input_callback not in litellm.input_callback: litellm.input_callback.append(_input_callback) if _async_input_callback not in litellm.input_callback: litellm.input_callback.append(_async_input_callback) litellm.success_callback = success_callback or [] if _success_callback not in litellm.success_callback: litellm.success_callback.append(_success_callback) if _async_success_callback not in litellm.success_callback: litellm.success_callback.append(_async_success_callback) litellm.failure_callback = failure_callback or [] if _failure_callback not in litellm.failure_callback: litellm.failure_callback.append(_failure_callback)
Save File
Cancel