287 lines
10 KiB
Python
287 lines
10 KiB
Python
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,
|
|
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.utils import event_from_exception
|
|
|
|
if TYPE_CHECKING:
|
|
from typing import Any, Dict
|
|
from datetime import datetime
|
|
|
|
try:
|
|
import litellm # type: ignore[import-not-found]
|
|
except ImportError:
|
|
raise DidNotEnable("LiteLLM not installed")
|
|
|
|
|
|
def _get_metadata_dict(kwargs):
|
|
# type: (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 _input_callback(kwargs):
|
|
# type: (Dict[str, Any]) -> None
|
|
"""Handle the start of a request."""
|
|
integration = sentry_sdk.get_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":
|
|
operation = "embeddings"
|
|
else:
|
|
operation = "chat"
|
|
|
|
# Start a new span/transaction
|
|
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]
|
|
)
|
|
messages_data = truncate_and_annotate_messages(input_list, span, scope)
|
|
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:
|
|
scope = sentry_sdk.get_current_scope()
|
|
messages_data = truncate_and_annotate_messages(messages, span, scope)
|
|
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)
|
|
|
|
# Record LiteLLM-specific parameters
|
|
litellm_params = {
|
|
"api_base": kwargs.get("api_base"),
|
|
"api_version": kwargs.get("api_version"),
|
|
"custom_llm_provider": kwargs.get("custom_llm_provider"),
|
|
}
|
|
for key, value in litellm_params.items():
|
|
if value is not None:
|
|
set_data_normalized(span, f"gen_ai.litellm.{key}", value)
|
|
|
|
|
|
def _success_callback(kwargs, completion_response, start_time, end_time):
|
|
# type: (Dict[str, Any], Any, datetime, datetime) -> None
|
|
"""Handle successful completion."""
|
|
|
|
span = _get_metadata_dict(kwargs).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:
|
|
# Always finish the span and clean up
|
|
span.__exit__(None, None, None)
|
|
|
|
|
|
def _failure_callback(kwargs, exception, start_time, end_time):
|
|
# type: (Dict[str, Any], Exception, datetime, 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, include_prompts=True):
|
|
# type: (LiteLLMIntegration, bool) -> None
|
|
self.include_prompts = include_prompts
|
|
|
|
@staticmethod
|
|
def setup_once():
|
|
# type: () -> None
|
|
"""Set up LiteLLM callbacks for monitoring."""
|
|
litellm.input_callback = litellm.input_callback or []
|
|
if _input_callback not in litellm.input_callback:
|
|
litellm.input_callback.append(_input_callback)
|
|
|
|
litellm.success_callback = litellm.success_callback or []
|
|
if _success_callback not in litellm.success_callback:
|
|
litellm.success_callback.append(_success_callback)
|
|
|
|
litellm.failure_callback = litellm.failure_callback or []
|
|
if _failure_callback not in litellm.failure_callback:
|
|
litellm.failure_callback.append(_failure_callback)
|