|
2 | 2 | import inspect |
3 | 3 | from functools import wraps |
4 | 4 | from .consts import ORIGIN, TOOL_ATTRIBUTES_MAP, GEN_AI_SYSTEM |
| 5 | +from sentry_sdk._types import BLOB_DATA_SUBSTITUTE |
5 | 6 | from typing import ( |
6 | 7 | cast, |
7 | 8 | TYPE_CHECKING, |
|
12 | 13 | Optional, |
13 | 14 | Union, |
14 | 15 | TypedDict, |
| 16 | + Dict, |
15 | 17 | ) |
16 | 18 |
|
17 | 19 | import sentry_sdk |
18 | 20 | from sentry_sdk.ai.utils import ( |
19 | 21 | set_data_normalized, |
20 | 22 | truncate_and_annotate_messages, |
21 | 23 | normalize_message_roles, |
| 24 | + redact_blob_message_parts, |
| 25 | + transform_google_content_part, |
| 26 | + get_modality_from_mime_type, |
22 | 27 | ) |
23 | 28 | from sentry_sdk.consts import OP, SPANDATA |
24 | 29 | from sentry_sdk.scope import should_send_default_pii |
@@ -145,44 +150,303 @@ def get_model_name(model: "Union[str, Model]") -> str: |
145 | 150 | return str(model) |
146 | 151 |
|
147 | 152 |
|
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 | + """ |
150 | 160 | if contents is None: |
151 | | - return None |
| 161 | + return [] |
152 | 162 |
|
153 | | - # Simple string case |
| 163 | + messages = [] |
| 164 | + |
| 165 | + # Handle string case |
154 | 166 | if isinstance(contents, str): |
155 | | - return contents |
| 167 | + return [{"role": "user", "content": contents}] |
156 | 168 |
|
157 | | - # List of contents or parts |
| 169 | + # Handle list case - process each item (non-recursive, flatten at top level) |
158 | 170 | if isinstance(contents, list): |
159 | | - texts = [] |
160 | 171 | 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 |
166 | 175 |
|
167 | | - # Dictionary case |
| 176 | + # Handle dictionary case (ContentDict) |
168 | 177 | 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}) |
174 | 235 |
|
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. |
178 | 283 |
|
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 | + } |
182 | 348 |
|
183 | 349 | return None |
184 | 350 |
|
185 | 351 |
|
| 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 | + |
186 | 450 | def _format_tools_for_span( |
187 | 451 | tools: "Iterable[Tool | Callable[..., Any]]", |
188 | 452 | ) -> "Optional[List[dict[str, Any]]]": |
@@ -457,14 +721,28 @@ def set_span_data_for_request( |
457 | 721 | if config and hasattr(config, "system_instruction"): |
458 | 722 | system_instruction = config.system_instruction |
459 | 723 | 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) |
468 | 746 |
|
469 | 747 | if messages: |
470 | 748 | normalized_messages = normalize_message_roles(messages) |
|
0 commit comments