Skip to content
Merged
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
271 changes: 1 addition & 270 deletions sentry_sdk/ai/utils.py
Original file line number Diff line number Diff line change
@@ -1,24 +1,16 @@
import inspect
import json
from copy import deepcopy
from typing import TYPE_CHECKING

from sentry_sdk._types import BLOB_DATA_SUBSTITUTE
from sentry_sdk.ai.consts import DATA_URL_BASE64_REGEX

if TYPE_CHECKING:
from typing import Any, Dict, List, Optional, Tuple, Union
Comment thread
alexander-alderman-webb marked this conversation as resolved.
from typing import Any, Dict, Optional, Tuple, Union

from sentry_sdk.tracing import Span

import sentry_sdk
from sentry_sdk.traces import StreamedSpan
from sentry_sdk.utils import logger

MAX_GEN_AI_MESSAGE_BYTES = 20_000 # 20KB
# Maximum characters when only a single message is left after bytes truncation
MAX_SINGLE_MESSAGE_CONTENT_CHARS = 10_000


class GEN_AI_ALLOWED_MESSAGE_ROLES:
SYSTEM = "system"
Expand Down Expand Up @@ -431,38 +423,6 @@ def transform_content_part(
return None


def transform_message_content(content: "Any") -> "Any":
"""
Transform message content, handling both string content and list of content blocks.

For list content, each item is transformed using transform_content_part().
Items that cannot be transformed (return None) are kept as-is.

Args:
content: Message content - can be a string, list of content blocks, or other

Returns:
- String content: returned as-is
- List content: list with each transformable item converted to standardized format
- Other: returned as-is
"""
if isinstance(content, str):
return content

if isinstance(content, (list, tuple)):
transformed = []
for item in content:
if isinstance(item, dict):
result = transform_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)
return transformed

return content


def _normalize_data(data: "Any", unpack: bool = True) -> "Any":
# convert pydantic data (e.g. OpenAI v1+) to json compatible format
if hasattr(data, "model_dump"):
Expand Down Expand Up @@ -536,235 +496,6 @@ def normalize_message_roles(messages: "list[dict[str, Any]]") -> "list[dict[str,
return normalized_messages


def _truncate_single_message_content_if_present(
message: "Dict[str, Any]", max_chars: int
) -> "Dict[str, Any]":
"""
Truncate a message's content to at most `max_chars` characters and append an
ellipsis if truncation occurs.
"""
if not isinstance(message, dict) or "content" not in message:
return message
content = message["content"]

if isinstance(content, str):
if len(content) <= max_chars:
return message
message["content"] = content[:max_chars] + "..."
return message

if isinstance(content, list):
remaining = max_chars
for item in content:
if isinstance(item, dict) and "text" in item:
text = item["text"]
if isinstance(text, str):
if len(text) > remaining:
item["text"] = text[:remaining] + "..."
remaining = 0
else:
remaining -= len(text)
return message

return message


def _find_truncation_index(messages: "List[Dict[str, Any]]", max_bytes: int) -> int:
"""
Find the index of the first message that would exceed the max bytes limit.
Compute the individual message sizes, and return the index of the first message from the back
of the list that would exceed the max bytes limit.
"""
running_sum = 0
for idx in range(len(messages) - 1, -1, -1):
size = len(json.dumps(messages[idx], separators=(",", ":")).encode("utf-8"))
running_sum += size
if running_sum > max_bytes:
return idx + 1

return 0


def _is_image_type_with_blob_content(item: "Dict[str, Any]") -> bool:
"""
Some content blocks contain an image_url property with base64 content as its value.
This is used to identify those while not leading to unnecessary copying of data when the image URL does not contain base64 content.
"""
if item.get("type") != "image_url":
return False

image_url_val = item.get("image_url")
image_url = (
image_url_val.get("url", "")
if isinstance(image_url_val, dict)
else (image_url_val or "")
)
data_url_match = DATA_URL_BASE64_REGEX.match(image_url)

return bool(data_url_match)


def redact_blob_message_parts(
messages: "List[Dict[str, Any]]",
) -> "List[Dict[str, Any]]":
"""
Redact blob message parts from the messages by replacing blob content with "[Filtered]".

This function creates a deep copy of messages that contain blob content to avoid
mutating the original message dictionaries. Messages without blob content are
returned as-is to minimize copying overhead.

e.g:
{
"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,..."
}
]
}
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": "[Filtered]"
}
]
}
"""

# First pass: check if any message contains blob content
has_blobs = False
for message in messages:
if not isinstance(message, dict):
continue
content = message.get("content")
if isinstance(content, list):
for item in content:
if isinstance(item, dict) and (
item.get("type") == "blob" or _is_image_type_with_blob_content(item)
):
has_blobs = True
break
if has_blobs:
break

# If no blobs found, return original messages to avoid unnecessary copying
if not has_blobs:
return messages

# Deep copy messages to avoid mutating the original
messages_copy = deepcopy(messages)

# Second pass: redact blob content in the copy
for message in messages_copy:
if not isinstance(message, dict):
continue

content = message.get("content")
if isinstance(content, list):
for item in content:
if isinstance(item, dict):
if item.get("type") == "blob":
item["content"] = BLOB_DATA_SUBSTITUTE
elif _is_image_type_with_blob_content(item):
if isinstance(item["image_url"], dict):
item["image_url"]["url"] = BLOB_DATA_SUBSTITUTE
else:
item["image_url"] = BLOB_DATA_SUBSTITUTE

return messages_copy


def truncate_messages_by_size(
messages: "List[Dict[str, Any]]",
max_bytes: int = MAX_GEN_AI_MESSAGE_BYTES,
max_single_message_chars: int = MAX_SINGLE_MESSAGE_CONTENT_CHARS,
) -> "Tuple[List[Dict[str, Any]], int]":
"""
Returns a truncated messages list, consisting of
- the last message, with its content truncated to `max_single_message_chars` characters,
if the last message's size exceeds `max_bytes` bytes; otherwise,
- the maximum number of messages, starting from the end of the `messages` list, whose total
serialized size does not exceed `max_bytes` bytes.

In the single message case, the serialized message size may exceed `max_bytes`, because
truncation is based only on character count in that case.
"""
serialized_json = json.dumps(messages, separators=(",", ":"))
current_size = len(serialized_json.encode("utf-8"))

if current_size <= max_bytes:
return messages, 0

truncation_index = _find_truncation_index(messages, max_bytes)
if truncation_index < len(messages):
truncated_messages = messages[truncation_index:]
else:
truncation_index = len(messages) - 1
truncated_messages = messages[-1:]

if len(truncated_messages) == 1:
truncated_messages[0] = _truncate_single_message_content_if_present(
deepcopy(truncated_messages[0]), max_chars=max_single_message_chars
)

return truncated_messages, truncation_index


def truncate_and_annotate_messages(
messages: "Optional[List[Dict[str, Any]]]",
span: "Any",
scope: "Any",
max_single_message_chars: int = MAX_SINGLE_MESSAGE_CONTENT_CHARS,
) -> "Optional[List[Dict[str, Any]]]":
if not messages:
return None

messages = redact_blob_message_parts(messages)

truncated_message = _truncate_single_message_content_if_present(
deepcopy(messages[-1]), max_chars=max_single_message_chars
)
if len(messages) > 1:
scope._gen_ai_original_message_count[span.span_id] = len(messages)

return [truncated_message]


def truncate_and_annotate_embedding_inputs(
messages: "Optional[List[Dict[str, Any]]]",
span: "Any",
scope: "Any",
max_bytes: int = MAX_GEN_AI_MESSAGE_BYTES,
) -> "Optional[List[Dict[str, Any]]]":
if not messages:
return None

messages = redact_blob_message_parts(messages)

truncated_messages, removed_count = truncate_messages_by_size(messages, max_bytes)
if removed_count > 0:
scope._gen_ai_original_message_count[span.span_id] = len(messages)

return truncated_messages


def set_conversation_id(conversation_id: str) -> None:
"""
Set the conversation_id in the scope.
Expand Down
42 changes: 8 additions & 34 deletions sentry_sdk/integrations/langchain.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,13 @@
normalize_message_roles,
set_data_normalized,
transform_content_part,
truncate_and_annotate_messages,
)
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.traces import StreamedSpan
from sentry_sdk.tracing_utils import (
_get_value,
has_span_streaming_enabled,
)
from sentry_sdk.utils import (
capture_internal_exceptions,
Expand Down Expand Up @@ -425,17 +423,11 @@ def on_llm_start(
for prompt in prompts
]

scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_messages(normalized_messages, span, scope)
if not has_span_streaming_enabled(client.options)
else normalized_messages
)
if messages_data is not None:
if normalized_messages is not None:
set_data_normalized(
span,
SPANDATA.GEN_AI_REQUEST_MESSAGES,
messages_data,
normalized_messages,
unpack=False,
)

Expand Down Expand Up @@ -537,17 +529,11 @@ def on_chat_model_start(
)
normalized_messages = normalize_message_roles(normalized_messages)

scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_messages(normalized_messages, span, scope)
if not has_span_streaming_enabled(client.options)
else normalized_messages
)
if messages_data is not None:
if normalized_messages is not None:
set_data_normalized(
span,
SPANDATA.GEN_AI_REQUEST_MESSAGES,
messages_data,
normalized_messages,
unpack=False,
)

Expand Down Expand Up @@ -1158,17 +1144,11 @@ def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
if input is not None and record_inputs:
normalized_messages = normalize_message_roles([input])

scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_messages(normalized_messages, span, scope)
if not has_span_streaming_enabled(client.options)
else normalized_messages
)
if messages_data is not None:
if normalized_messages is not None:
set_data_normalized(
span,
SPANDATA.GEN_AI_REQUEST_MESSAGES,
messages_data,
normalized_messages,
unpack=False,
)

Expand Down Expand Up @@ -1220,17 +1200,11 @@ def new_stream(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
if input is not None and record_inputs:
normalized_messages = normalize_message_roles([input])

scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_messages(normalized_messages, span, scope)
if not has_span_streaming_enabled(client.options)
else normalized_messages
)
if messages_data is not None:
if normalized_messages is not None:
set_data_normalized(
span,
SPANDATA.GEN_AI_REQUEST_MESSAGES,
messages_data,
normalized_messages,
unpack=False,
)

Expand Down
Loading
Loading