Skip to content

Commit 749d8e5

Browse files
fix(integrations): google-genai: reworked gen_ai.request.messages extraction from parameters (#5275)
### Description Previously we only extracted only text parts were extracted. Now the full range of possibilities are covered. #### Issues Closes https://linear.app/getsentry/issue/TET-1638/redact-images-google-genai
1 parent ab5cdde commit 749d8e5

2 files changed

Lines changed: 1079 additions & 32 deletions

File tree

sentry_sdk/integrations/google_genai/utils.py

Lines changed: 310 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import inspect
33
from functools import wraps
44
from .consts import ORIGIN, TOOL_ATTRIBUTES_MAP, GEN_AI_SYSTEM
5+
from sentry_sdk._types import BLOB_DATA_SUBSTITUTE
56
from typing import (
67
cast,
78
TYPE_CHECKING,
@@ -12,13 +13,17 @@
1213
Optional,
1314
Union,
1415
TypedDict,
16+
Dict,
1517
)
1618

1719
import sentry_sdk
1820
from sentry_sdk.ai.utils import (
1921
set_data_normalized,
2022
truncate_and_annotate_messages,
2123
normalize_message_roles,
24+
redact_blob_message_parts,
25+
transform_google_content_part,
26+
get_modality_from_mime_type,
2227
)
2328
from sentry_sdk.consts import OP, SPANDATA
2429
from sentry_sdk.scope import should_send_default_pii
@@ -145,44 +150,303 @@ def get_model_name(model: "Union[str, Model]") -> str:
145150
return str(model)
146151

147152

148-
def extract_contents_text(contents: "ContentListUnion") -> "Optional[str]":
149-
"""Extract text from contents parameter which can have various formats."""
153+
def extract_contents_messages(contents: "ContentListUnion") -> "List[Dict[str, Any]]":
154+
"""Extract messages from contents parameter which can have various formats.
155+
156+
Returns a list of message dictionaries in the format:
157+
- System: {"role": "system", "content": "string"}
158+
- User/Assistant: {"role": "user"|"assistant", "content": [{"text": "...", "type": "text"}, ...]}
159+
"""
150160
if contents is None:
151-
return None
161+
return []
152162

153-
# Simple string case
163+
messages = []
164+
165+
# Handle string case
154166
if isinstance(contents, str):
155-
return contents
167+
return [{"role": "user", "content": contents}]
156168

157-
# List of contents or parts
169+
# Handle list case - process each item (non-recursive, flatten at top level)
158170
if isinstance(contents, list):
159-
texts = []
160171
for item in contents:
161-
# Recursively extract text from each item
162-
extracted = extract_contents_text(item)
163-
if extracted:
164-
texts.append(extracted)
165-
return " ".join(texts) if texts else None
172+
item_messages = extract_contents_messages(item)
173+
messages.extend(item_messages)
174+
return messages
166175

167-
# Dictionary case
176+
# Handle dictionary case (ContentDict)
168177
if isinstance(contents, dict):
169-
if "text" in contents:
170-
return contents["text"]
171-
# Try to extract from parts if present in dict
172-
if "parts" in contents:
173-
return extract_contents_text(contents["parts"])
178+
role = contents.get("role", "user")
179+
parts = contents.get("parts")
180+
181+
if parts:
182+
content_parts = []
183+
tool_messages = []
184+
185+
for part in parts:
186+
part_result = _extract_part_content(part)
187+
if part_result is None:
188+
continue
189+
190+
if isinstance(part_result, dict) and part_result.get("role") == "tool":
191+
# Tool message - add separately
192+
tool_messages.append(part_result)
193+
else:
194+
# Regular content part
195+
content_parts.append(part_result)
196+
197+
# Add main message if we have content parts
198+
if content_parts:
199+
# Normalize role: "model" -> "assistant"
200+
normalized_role = "assistant" if role == "model" else role or "user"
201+
messages.append({"role": normalized_role, "content": content_parts})
202+
203+
# Add tool messages
204+
messages.extend(tool_messages)
205+
elif "text" in contents:
206+
# Simple text in dict
207+
messages.append(
208+
{
209+
"role": role or "user",
210+
"content": [{"text": contents["text"], "type": "text"}],
211+
}
212+
)
213+
214+
return messages
215+
216+
# Handle Content object
217+
if hasattr(contents, "parts") and contents.parts:
218+
role = getattr(contents, "role", None) or "user"
219+
content_parts = []
220+
tool_messages = []
221+
222+
for part in contents.parts:
223+
part_result = _extract_part_content(part)
224+
if part_result is None:
225+
continue
226+
227+
if isinstance(part_result, dict) and part_result.get("role") == "tool":
228+
tool_messages.append(part_result)
229+
else:
230+
content_parts.append(part_result)
231+
232+
if content_parts:
233+
normalized_role = "assistant" if role == "model" else role
234+
messages.append({"role": normalized_role, "content": content_parts})
174235

175-
# Content object with parts - recurse into parts
176-
if getattr(contents, "parts", None):
177-
return extract_contents_text(contents.parts)
236+
messages.extend(tool_messages)
237+
return messages
238+
239+
# Handle Part object directly
240+
part_result = _extract_part_content(contents)
241+
if part_result:
242+
if isinstance(part_result, dict) and part_result.get("role") == "tool":
243+
return [part_result]
244+
else:
245+
return [{"role": "user", "content": [part_result]}]
246+
247+
# Handle PIL.Image.Image
248+
try:
249+
from PIL import Image as PILImage # type: ignore[import-not-found]
250+
251+
if isinstance(contents, PILImage.Image):
252+
blob_part = _extract_pil_image(contents)
253+
if blob_part:
254+
return [{"role": "user", "content": [blob_part]}]
255+
except ImportError:
256+
pass
257+
258+
# Handle File object
259+
if hasattr(contents, "uri") and hasattr(contents, "mime_type"):
260+
# File object
261+
file_uri = getattr(contents, "uri", None)
262+
mime_type = getattr(contents, "mime_type", None)
263+
if file_uri and mime_type:
264+
blob_part = {
265+
"type": "uri",
266+
"modality": get_modality_from_mime_type(mime_type),
267+
"mime_type": mime_type,
268+
"uri": file_uri,
269+
}
270+
return [{"role": "user", "content": [blob_part]}]
271+
272+
# Handle direct text attribute
273+
if hasattr(contents, "text") and contents.text:
274+
return [
275+
{"role": "user", "content": [{"text": str(contents.text), "type": "text"}]}
276+
]
277+
278+
return []
279+
280+
281+
def _extract_part_content(part: "Any") -> "Optional[dict[str, Any]]":
282+
"""Extract content from a Part object or dict.
178283
179-
# Direct text attribute
180-
if hasattr(contents, "text"):
181-
return contents.text
284+
Returns:
285+
- dict for content part (text/blob) or tool message
286+
- None if part should be skipped
287+
"""
288+
if part is None:
289+
return None
290+
291+
# Handle dict Part
292+
if isinstance(part, dict):
293+
# Check for function_response first (tool message)
294+
if "function_response" in part:
295+
return _extract_tool_message_from_part(part)
296+
297+
if part.get("text"):
298+
return {"text": part["text"], "type": "text"}
299+
300+
# Try using Google-specific transform for dict formats (inline_data, file_data)
301+
result = transform_google_content_part(part)
302+
if result is not None:
303+
# For inline_data with bytes data, substitute the content
304+
if "inline_data" in part:
305+
inline_data = part["inline_data"]
306+
if isinstance(inline_data, dict) and isinstance(
307+
inline_data.get("data"), bytes
308+
):
309+
result["content"] = BLOB_DATA_SUBSTITUTE
310+
return result
311+
312+
return None
313+
314+
# Handle Part object
315+
# Check for function_response (tool message)
316+
if hasattr(part, "function_response") and part.function_response:
317+
return _extract_tool_message_from_part(part)
318+
319+
# Handle text
320+
if hasattr(part, "text") and part.text:
321+
return {"text": part.text, "type": "text"}
322+
323+
# Handle file_data
324+
if hasattr(part, "file_data") and part.file_data:
325+
file_data = part.file_data
326+
file_uri = getattr(file_data, "file_uri", None)
327+
mime_type = getattr(file_data, "mime_type", None)
328+
if file_uri and mime_type:
329+
return {
330+
"type": "uri",
331+
"modality": get_modality_from_mime_type(mime_type),
332+
"mime_type": mime_type,
333+
"uri": file_uri,
334+
}
335+
336+
# Handle inline_data
337+
if hasattr(part, "inline_data") and part.inline_data:
338+
inline_data = part.inline_data
339+
data = getattr(inline_data, "data", None)
340+
mime_type = getattr(inline_data, "mime_type", None)
341+
if data and mime_type:
342+
if isinstance(data, bytes):
343+
return {
344+
"type": "blob",
345+
"mime_type": mime_type,
346+
"content": BLOB_DATA_SUBSTITUTE,
347+
}
182348

183349
return None
184350

185351

352+
def _extract_tool_message_from_part(part: "Any") -> "Optional[dict[str, Any]]":
353+
"""Extract tool message from a Part with function_response.
354+
355+
Returns:
356+
{"role": "tool", "content": {"toolCallId": "...", "toolName": "...", "output": "..."}}
357+
or None if not a valid tool message
358+
"""
359+
function_response = None
360+
361+
if isinstance(part, dict):
362+
function_response = part.get("function_response")
363+
elif hasattr(part, "function_response"):
364+
function_response = part.function_response
365+
366+
if not function_response:
367+
return None
368+
369+
# Extract fields from function_response
370+
tool_call_id = None
371+
tool_name = None
372+
output = None
373+
374+
if isinstance(function_response, dict):
375+
tool_call_id = function_response.get("id")
376+
tool_name = function_response.get("name")
377+
response_dict = function_response.get("response", {})
378+
# Prefer "output" key if present, otherwise use entire response
379+
output = response_dict.get("output", response_dict)
380+
else:
381+
# FunctionResponse object
382+
tool_call_id = getattr(function_response, "id", None)
383+
tool_name = getattr(function_response, "name", None)
384+
response_obj = getattr(function_response, "response", None)
385+
if response_obj is None:
386+
response_obj = {}
387+
if isinstance(response_obj, dict):
388+
output = response_obj.get("output", response_obj)
389+
else:
390+
output = response_obj
391+
392+
if not tool_name:
393+
return None
394+
395+
return {
396+
"role": "tool",
397+
"content": {
398+
"toolCallId": str(tool_call_id) if tool_call_id else None,
399+
"toolName": str(tool_name),
400+
"output": safe_serialize(output) if output is not None else None,
401+
},
402+
}
403+
404+
405+
def _extract_pil_image(image: "Any") -> "Optional[dict[str, Any]]":
406+
"""Extract blob part from PIL.Image.Image."""
407+
try:
408+
from PIL import Image as PILImage
409+
410+
if not isinstance(image, PILImage.Image):
411+
return None
412+
413+
# Get format, default to JPEG
414+
format_str = image.format or "JPEG"
415+
suffix = format_str.lower()
416+
mime_type = f"image/{suffix}"
417+
418+
return {
419+
"type": "blob",
420+
"mime_type": mime_type,
421+
"content": BLOB_DATA_SUBSTITUTE,
422+
}
423+
except Exception:
424+
return None
425+
426+
427+
def extract_contents_text(contents: "ContentListUnion") -> "Optional[str]":
428+
"""Extract text from contents parameter which can have various formats.
429+
430+
This is a compatibility function that extracts text from messages.
431+
For new code, use extract_contents_messages instead.
432+
"""
433+
messages = extract_contents_messages(contents)
434+
if not messages:
435+
return None
436+
437+
texts = []
438+
for message in messages:
439+
content = message.get("content")
440+
if isinstance(content, str):
441+
texts.append(content)
442+
elif isinstance(content, list):
443+
for part in content:
444+
if isinstance(part, dict) and part.get("type") == "text":
445+
texts.append(part.get("text", ""))
446+
447+
return " ".join(texts) if texts else None
448+
449+
186450
def _format_tools_for_span(
187451
tools: "Iterable[Tool | Callable[..., Any]]",
188452
) -> "Optional[List[dict[str, Any]]]":
@@ -457,14 +721,28 @@ def set_span_data_for_request(
457721
if config and hasattr(config, "system_instruction"):
458722
system_instruction = config.system_instruction
459723
if system_instruction:
460-
system_text = extract_contents_text(system_instruction)
461-
if system_text:
462-
messages.append({"role": "system", "content": system_text})
463-
464-
# Add user message
465-
contents_text = extract_contents_text(contents)
466-
if contents_text:
467-
messages.append({"role": "user", "content": contents_text})
724+
system_messages = extract_contents_messages(system_instruction)
725+
# System instruction should be a single system message
726+
# Extract text from all messages and combine into one system message
727+
system_texts = []
728+
for msg in system_messages:
729+
content = msg.get("content")
730+
if isinstance(content, list):
731+
# Extract text from content parts
732+
for part in content:
733+
if isinstance(part, dict) and part.get("type") == "text":
734+
system_texts.append(part.get("text", ""))
735+
elif isinstance(content, str):
736+
system_texts.append(content)
737+
738+
if system_texts:
739+
messages.append(
740+
{"role": "system", "content": " ".join(system_texts)}
741+
)
742+
743+
# Extract messages from contents
744+
contents_messages = extract_contents_messages(contents)
745+
messages.extend(contents_messages)
468746

469747
if messages:
470748
normalized_messages = normalize_message_roles(messages)

0 commit comments

Comments
 (0)