Skip to content

Commit d197318

Browse files
ericapisaniclaude
andcommitted
fix(google_genai): Redact binary data in inline_data and fix multi-part message extraction
Redact binary/byte data that appears in `inline_data`, which fixes a failing `checkBinaryRedaction` assertion in the AI testing framework. Also includes changes for the following: - Properly handling lists of part-like items (merging into single multi-part user message) - Handling bare inline_data dicts that aren't wrapped in Part objects - Always substituting blob data (both bytes and base64 strings) - Moving PIL import to module level with availability flag to reduce all the dynamic imports of the module within the code Fixes PY-2287 and #5965 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 4a056c0 commit d197318

File tree

2 files changed

+226
-49
lines changed

2 files changed

+226
-49
lines changed

sentry_sdk/integrations/google_genai/utils.py

Lines changed: 116 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
event_from_exception,
3232
safe_serialize,
3333
)
34-
from google.genai.types import GenerateContentConfig, Part, Content
34+
from google.genai.types import GenerateContentConfig, Part, Content, PartDict
3535
from itertools import chain
3636

3737
if TYPE_CHECKING:
@@ -47,6 +47,18 @@
4747
ContentUnion,
4848
)
4949

50+
_is_PIL_available = False
51+
try:
52+
from PIL import Image as PILImage # type: ignore[import-not-found]
53+
54+
_is_PIL_available = True
55+
except ImportError:
56+
pass
57+
58+
# Keys to use when checking to see if a dict provided by the user
59+
# is Part-like (as opposed to a Content or multi-turn conversation entry).
60+
_PART_DICT_KEYS = PartDict.__optional_keys__
61+
5062

5163
class UsageData(TypedDict):
5264
"""Structure for token usage data."""
@@ -169,12 +181,23 @@ def extract_contents_messages(contents: "ContentListUnion") -> "List[Dict[str, A
169181
if isinstance(contents, str):
170182
return [{"role": "user", "content": contents}]
171183

172-
# Handle list case - process each item (non-recursive, flatten at top level)
184+
# Handle list case
173185
if isinstance(contents, list):
174-
for item in contents:
175-
item_messages = extract_contents_messages(item)
176-
messages.extend(item_messages)
177-
return messages
186+
if contents and all(_is_part_like(item) for item in contents):
187+
# All items are parts — merge into a single multi-part user message
188+
content_parts = []
189+
for item in contents:
190+
part = _extract_part_from_item(item)
191+
if part is not None:
192+
content_parts.append(part)
193+
194+
return [{"role": "user", "content": content_parts}]
195+
else:
196+
# Multi-turn conversation or mixed content types
197+
for item in contents:
198+
item_messages = extract_contents_messages(item)
199+
messages.extend(item_messages)
200+
return messages
178201

179202
# Handle dictionary case (ContentDict)
180203
if isinstance(contents, dict):
@@ -206,13 +229,23 @@ def extract_contents_messages(contents: "ContentListUnion") -> "List[Dict[str, A
206229
# Add tool messages
207230
messages.extend(tool_messages)
208231
elif "text" in contents:
209-
# Simple text in dict
210232
messages.append(
211233
{
212-
"role": role or "user",
234+
"role": role,
213235
"content": [{"text": contents["text"], "type": "text"}],
214236
}
215237
)
238+
elif "inline_data" in contents:
239+
# The "data" will always be bytes (or bytes within a string),
240+
# so if this is present, it's safe to automatically substitute with the placeholder
241+
messages.append(
242+
{
243+
"inline_data": {
244+
"mime_type": contents["inline_data"].get("mime_type", ""),
245+
"data": BLOB_DATA_SUBSTITUTE,
246+
}
247+
}
248+
)
216249

217250
return messages
218251

@@ -248,15 +281,10 @@ def extract_contents_messages(contents: "ContentListUnion") -> "List[Dict[str, A
248281
return [{"role": "user", "content": [part_result]}]
249282

250283
# Handle PIL.Image.Image
251-
try:
252-
from PIL import Image as PILImage # type: ignore[import-not-found]
253-
254-
if isinstance(contents, PILImage.Image):
255-
blob_part = _extract_pil_image(contents)
256-
if blob_part:
257-
return [{"role": "user", "content": [blob_part]}]
258-
except ImportError:
259-
pass
284+
if _is_PIL_available and isinstance(contents, PILImage.Image):
285+
blob_part = _extract_pil_image(contents)
286+
if blob_part:
287+
return [{"role": "user", "content": [blob_part]}]
260288

261289
# Handle File object
262290
if hasattr(contents, "uri") and hasattr(contents, "mime_type"):
@@ -310,11 +338,9 @@ def _extract_part_content(part: "Any") -> "Optional[dict[str, Any]]":
310338
if result is not None:
311339
# For inline_data with bytes data, substitute the content
312340
if "inline_data" in part:
313-
inline_data = part["inline_data"]
314-
if isinstance(inline_data, dict) and isinstance(
315-
inline_data.get("data"), bytes
316-
):
317-
result["content"] = BLOB_DATA_SUBSTITUTE
341+
# inline_data.data will always be bytes, or a string containing base64-encoded bytes,
342+
# so can automatically substitute without further checks
343+
result["content"] = BLOB_DATA_SUBSTITUTE
318344
return result
319345

320346
return None
@@ -357,18 +383,11 @@ def _extract_part_content(part: "Any") -> "Optional[dict[str, Any]]":
357383
if mime_type is None:
358384
mime_type = ""
359385

360-
# Handle both bytes (binary data) and str (base64-encoded data)
361-
if isinstance(data, bytes):
362-
content = BLOB_DATA_SUBSTITUTE
363-
else:
364-
# For non-bytes data (e.g., base64 strings), use as-is
365-
content = data
366-
367386
return {
368387
"type": "blob",
369388
"modality": get_modality_from_mime_type(mime_type),
370389
"mime_type": mime_type,
371-
"content": content,
390+
"content": BLOB_DATA_SUBSTITUTE,
372391
}
373392

374393
return None
@@ -429,25 +448,78 @@ def _extract_tool_message_from_part(part: "Any") -> "Optional[dict[str, Any]]":
429448

430449
def _extract_pil_image(image: "Any") -> "Optional[dict[str, Any]]":
431450
"""Extract blob part from PIL.Image.Image."""
432-
try:
433-
from PIL import Image as PILImage
451+
if not _is_PIL_available or not isinstance(image, PILImage.Image):
452+
return None
434453

435-
if not isinstance(image, PILImage.Image):
436-
return None
454+
# Get format, default to JPEG
455+
format_str = image.format or "JPEG"
456+
suffix = format_str.lower()
457+
mime_type = f"image/{suffix}"
458+
459+
return {
460+
"type": "blob",
461+
"modality": get_modality_from_mime_type(mime_type),
462+
"mime_type": mime_type,
463+
"content": BLOB_DATA_SUBSTITUTE,
464+
}
437465

438-
# Get format, default to JPEG
439-
format_str = image.format or "JPEG"
440-
suffix = format_str.lower()
441-
mime_type = f"image/{suffix}"
442466

467+
def _is_part_like(item: "Any") -> bool:
468+
"""Check if item is a part-like value (PartUnionDict) rather than a Content/multi-turn entry."""
469+
if isinstance(item, (str, Part)):
470+
return True
471+
if isinstance(item, (list, Content)):
472+
return False
473+
if isinstance(item, dict):
474+
if "role" in item or "parts" in item:
475+
return False
476+
# Part objects that came in as plain dicts
477+
return bool(_PART_DICT_KEYS & item.keys())
478+
# File objects
479+
if hasattr(item, "uri"):
480+
return True
481+
# PIL.Image
482+
if _is_PIL_available and isinstance(item, PILImage.Image):
483+
return True
484+
return False
485+
486+
487+
def _extract_part_from_item(item: "Any") -> "Optional[dict[str, Any]]":
488+
"""Convert a single part-like item to a content part dict."""
489+
if isinstance(item, str):
490+
return {"text": item, "type": "text"}
491+
492+
# Handle bare inline_data dicts directly to preserve the raw format
493+
if isinstance(item, dict) and "inline_data" in item:
443494
return {
444-
"type": "blob",
445-
"modality": get_modality_from_mime_type(mime_type),
446-
"mime_type": mime_type,
447-
"content": BLOB_DATA_SUBSTITUTE,
495+
"inline_data": {
496+
"mime_type": item["inline_data"].get("mime_type", ""),
497+
"data": BLOB_DATA_SUBSTITUTE,
498+
}
448499
}
449-
except Exception:
450-
return None
500+
501+
# For other dicts and Part objects, use existing _extract_part_content
502+
result = _extract_part_content(item)
503+
if result is not None:
504+
return result
505+
506+
# PIL.Image
507+
if _is_PIL_available and isinstance(item, PILImage.Image):
508+
return _extract_pil_image(item)
509+
510+
# File objects
511+
if hasattr(item, "uri") and hasattr(item, "mime_type"):
512+
file_uri = getattr(item, "uri", None)
513+
mime_type = getattr(item, "mime_type", None) or ""
514+
if file_uri is not None:
515+
return {
516+
"type": "uri",
517+
"modality": get_modality_from_mime_type(mime_type),
518+
"mime_type": mime_type,
519+
"uri": file_uri,
520+
}
521+
522+
return None
451523

452524

453525
def extract_contents_text(contents: "ContentListUnion") -> "Optional[str]":

tests/integrations/google_genai/test_google_genai.py

Lines changed: 110 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -942,11 +942,9 @@ def test_google_genai_message_truncation(
942942
assert isinstance(parsed_messages, list)
943943
assert len(parsed_messages) == 1
944944
assert parsed_messages[0]["role"] == "user"
945-
assert small_content in parsed_messages[0]["content"]
946945

947-
assert (
948-
event["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 2
949-
)
946+
# What "small content" becomes because the large message used the entire character limit
947+
assert "..." in parsed_messages[0]["content"][1]["text"]
950948

951949

952950
# Sample embed content API response JSON
@@ -1595,6 +1593,12 @@ def test_generate_content_with_function_response(
15951593

15961594
mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON)
15971595

1596+
# Conversation with the function call from the model
1597+
function_call = genai_types.FunctionCall(
1598+
name="get_weather",
1599+
args={"location": "Paris"},
1600+
)
1601+
15981602
# Conversation with function response (tool result)
15991603
function_response = genai_types.FunctionResponse(
16001604
id="call_123", name="get_weather", response={"output": "Sunny, 72F"}
@@ -1603,6 +1607,9 @@ def test_generate_content_with_function_response(
16031607
genai_types.Content(
16041608
role="user", parts=[genai_types.Part(text="What's the weather in Paris?")]
16051609
),
1610+
genai_types.Content(
1611+
role="model", parts=[genai_types.Part(function_call=function_call)]
1612+
),
16061613
genai_types.Content(
16071614
role="user", parts=[genai_types.Part(function_response=function_response)]
16081615
),
@@ -1708,7 +1715,13 @@ def test_generate_content_with_part_object_directly(
17081715
def test_generate_content_with_list_of_dicts(
17091716
sentry_init, capture_events, mock_genai_client
17101717
):
1711-
"""Test generate_content with list of dict format inputs."""
1718+
"""
1719+
Test generate_content with list of dict format inputs.
1720+
1721+
We only keep (and assert) the last dict in `content` because we've made popping the last message a form of
1722+
message truncation to keep the span size within limits. If we were following OTEL conventions, all 3 dicts
1723+
would be present.
1724+
"""
17121725
sentry_init(
17131726
integrations=[GoogleGenAIIntegration(include_prompts=True)],
17141727
traces_sample_rate=1.0,
@@ -1788,6 +1801,98 @@ def test_generate_content_with_dict_inline_data(
17881801
assert messages[0]["content"][1]["content"] == BLOB_DATA_SUBSTITUTE
17891802

17901803

1804+
def test_generate_content_without_parts_property_inline_data(
1805+
sentry_init, capture_events, mock_genai_client
1806+
):
1807+
sentry_init(
1808+
integrations=[GoogleGenAIIntegration(include_prompts=True)],
1809+
traces_sample_rate=1.0,
1810+
send_default_pii=True,
1811+
)
1812+
events = capture_events()
1813+
1814+
mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON)
1815+
1816+
contents = [
1817+
{"text": "What's in this image?"},
1818+
{"inline_data": {"data": b"fake_binary_data", "mime_type": "image/gif"}},
1819+
]
1820+
1821+
with mock.patch.object(
1822+
mock_genai_client._api_client, "request", return_value=mock_http_response
1823+
):
1824+
with start_transaction(name="google_genai"):
1825+
mock_genai_client.models.generate_content(
1826+
model="gemini-1.5-flash", contents=contents, config=create_test_config()
1827+
)
1828+
1829+
(event,) = events
1830+
invoke_span = event["spans"][0]
1831+
1832+
messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES])
1833+
1834+
assert len(messages) == 1
1835+
1836+
assert len(messages[0]["content"]) == 2
1837+
assert messages[0]["role"] == "user"
1838+
assert messages[0]["content"][0] == {
1839+
"text": "What's in this image?",
1840+
"type": "text",
1841+
}
1842+
assert messages[0]["content"][1]["inline_data"]
1843+
1844+
assert messages[0]["content"][1]["inline_data"]["data"] == BLOB_DATA_SUBSTITUTE
1845+
assert messages[0]["content"][1]["inline_data"]["mime_type"] == "image/gif"
1846+
1847+
1848+
def test_generate_content_without_parts_property_inline_data_and_binary_data_within_string(
1849+
sentry_init, capture_events, mock_genai_client
1850+
):
1851+
sentry_init(
1852+
integrations=[GoogleGenAIIntegration(include_prompts=True)],
1853+
traces_sample_rate=1.0,
1854+
send_default_pii=True,
1855+
)
1856+
events = capture_events()
1857+
1858+
mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON)
1859+
1860+
contents = [
1861+
{"text": "What's in this image?"},
1862+
{
1863+
"inline_data": {
1864+
"data": "iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs+9AAAAFUlEQVR42mP8z8BQz0AEYBxVSF+FABJADveWkH6oAAAAAElFTkSuQmCC",
1865+
"mime_type": "image/png",
1866+
}
1867+
},
1868+
]
1869+
1870+
with mock.patch.object(
1871+
mock_genai_client._api_client, "request", return_value=mock_http_response
1872+
):
1873+
with start_transaction(name="google_genai"):
1874+
mock_genai_client.models.generate_content(
1875+
model="gemini-1.5-flash", contents=contents, config=create_test_config()
1876+
)
1877+
1878+
(event,) = events
1879+
invoke_span = event["spans"][0]
1880+
1881+
messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES])
1882+
assert len(messages) == 1
1883+
assert messages[0]["role"] == "user"
1884+
1885+
assert len(messages[0]["content"]) == 2
1886+
assert messages[0]["content"][0] == {
1887+
"text": "What's in this image?",
1888+
"type": "text",
1889+
}
1890+
assert messages[0]["content"][1]["inline_data"]
1891+
1892+
assert messages[0]["content"][1]["inline_data"]["data"] == BLOB_DATA_SUBSTITUTE
1893+
assert messages[0]["content"][1]["inline_data"]["mime_type"] == "image/png"
1894+
1895+
17911896
# Tests for extract_contents_messages function
17921897
def test_extract_contents_messages_none():
17931898
"""Test extract_contents_messages with None input"""

0 commit comments

Comments
 (0)