Skip to content

Commit 999667f

Browse files
authored
fix(google_genai): Redact binary data in inline_data and fix multi-part message extraction (#5977)
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
1 parent 822b244 commit 999667f

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
@@ -941,11 +941,9 @@ def test_google_genai_message_truncation(
941941
assert isinstance(parsed_messages, list)
942942
assert len(parsed_messages) == 1
943943
assert parsed_messages[0]["role"] == "user"
944-
assert small_content in parsed_messages[0]["content"]
945944

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

950948

951949
# Sample embed content API response JSON
@@ -1594,6 +1592,12 @@ def test_generate_content_with_function_response(
15941592

15951593
mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON)
15961594

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

17891802

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

0 commit comments

Comments
 (0)