|
4 | 4 | import pytest |
5 | 5 |
|
6 | 6 | import sentry_sdk |
7 | | -from sentry_sdk._types import AnnotatedValue |
| 7 | +from sentry_sdk._types import ( |
| 8 | + AnnotatedValue, |
| 9 | + SENSITIVE_DATA_SUBSTITUTE, |
| 10 | + BLOB_DATA_SUBSTITUTE, |
| 11 | +) |
8 | 12 | from sentry_sdk.ai.monitoring import ai_track |
9 | 13 | from sentry_sdk.ai.utils import ( |
10 | 14 | MAX_GEN_AI_MESSAGE_BYTES, |
|
13 | 17 | truncate_and_annotate_messages, |
14 | 18 | truncate_messages_by_size, |
15 | 19 | _find_truncation_index, |
| 20 | + redact_blob_message_parts, |
16 | 21 | ) |
17 | 22 | from sentry_sdk.serializer import serialize |
18 | 23 | from sentry_sdk.utils import safe_serialize |
@@ -425,6 +430,49 @@ def __init__(self): |
425 | 430 | assert isinstance(result, list) |
426 | 431 | assert result[0] == large_messages[-len(result)] |
427 | 432 |
|
| 433 | + def test_preserves_original_messages_with_blobs(self): |
| 434 | + """Test that truncate_and_annotate_messages doesn't mutate the original messages""" |
| 435 | + |
| 436 | + class MockSpan: |
| 437 | + def __init__(self): |
| 438 | + self.span_id = "test_span_id" |
| 439 | + self.data = {} |
| 440 | + |
| 441 | + def set_data(self, key, value): |
| 442 | + self.data[key] = value |
| 443 | + |
| 444 | + class MockScope: |
| 445 | + def __init__(self): |
| 446 | + self._gen_ai_original_message_count = {} |
| 447 | + |
| 448 | + messages = [ |
| 449 | + { |
| 450 | + "role": "user", |
| 451 | + "content": [ |
| 452 | + {"text": "What's in this image?", "type": "text"}, |
| 453 | + { |
| 454 | + "type": "blob", |
| 455 | + "modality": "image", |
| 456 | + "content": "data:image/jpeg;base64,original_content", |
| 457 | + }, |
| 458 | + ], |
| 459 | + } |
| 460 | + ] |
| 461 | + |
| 462 | + original_blob_content = messages[0]["content"][1]["content"] |
| 463 | + |
| 464 | + span = MockSpan() |
| 465 | + scope = MockScope() |
| 466 | + |
| 467 | + # This should NOT mutate the original messages |
| 468 | + result = truncate_and_annotate_messages(messages, span, scope) |
| 469 | + |
| 470 | + # Verify original is unchanged |
| 471 | + assert messages[0]["content"][1]["content"] == original_blob_content |
| 472 | + |
| 473 | + # Verify result has redacted content |
| 474 | + assert result[0]["content"][1]["content"] == BLOB_DATA_SUBSTITUTE |
| 475 | + |
428 | 476 |
|
429 | 477 | class TestClientAnnotation: |
430 | 478 | def test_client_wraps_truncated_messages_in_annotated_value(self, large_messages): |
@@ -542,3 +590,170 @@ def __init__(self): |
542 | 590 | assert isinstance(messages_value, AnnotatedValue) |
543 | 591 | assert messages_value.metadata["len"] == stored_original_length |
544 | 592 | assert len(messages_value.value) == len(truncated_messages) |
| 593 | + |
| 594 | + |
| 595 | +class TestRedactBlobMessageParts: |
| 596 | + def test_redacts_single_blob_content(self): |
| 597 | + """Test that blob content is redacted without mutating original messages""" |
| 598 | + messages = [ |
| 599 | + { |
| 600 | + "role": "user", |
| 601 | + "content": [ |
| 602 | + { |
| 603 | + "text": "How many ponies do you see in the image?", |
| 604 | + "type": "text", |
| 605 | + }, |
| 606 | + { |
| 607 | + "type": "blob", |
| 608 | + "modality": "image", |
| 609 | + "mime_type": "image/jpeg", |
| 610 | + "content": "data:image/jpeg;base64,/9j/4AAQSkZJRg==", |
| 611 | + }, |
| 612 | + ], |
| 613 | + } |
| 614 | + ] |
| 615 | + |
| 616 | + # Save original blob content for comparison |
| 617 | + original_blob_content = messages[0]["content"][1]["content"] |
| 618 | + |
| 619 | + result = redact_blob_message_parts(messages) |
| 620 | + |
| 621 | + # Original messages should be UNCHANGED |
| 622 | + assert messages[0]["content"][1]["content"] == original_blob_content |
| 623 | + |
| 624 | + # Result should have redacted content |
| 625 | + assert ( |
| 626 | + result[0]["content"][0]["text"] |
| 627 | + == "How many ponies do you see in the image?" |
| 628 | + ) |
| 629 | + assert result[0]["content"][0]["type"] == "text" |
| 630 | + assert result[0]["content"][1]["type"] == "blob" |
| 631 | + assert result[0]["content"][1]["modality"] == "image" |
| 632 | + assert result[0]["content"][1]["mime_type"] == "image/jpeg" |
| 633 | + assert result[0]["content"][1]["content"] == BLOB_DATA_SUBSTITUTE |
| 634 | + |
| 635 | + def test_redacts_multiple_blob_parts(self): |
| 636 | + """Test that multiple blob parts are redacted without mutation""" |
| 637 | + messages = [ |
| 638 | + { |
| 639 | + "role": "user", |
| 640 | + "content": [ |
| 641 | + {"text": "Compare these images", "type": "text"}, |
| 642 | + { |
| 643 | + "type": "blob", |
| 644 | + "modality": "image", |
| 645 | + "mime_type": "image/jpeg", |
| 646 | + "content": "data:image/jpeg;base64,first_image", |
| 647 | + }, |
| 648 | + { |
| 649 | + "type": "blob", |
| 650 | + "modality": "image", |
| 651 | + "mime_type": "image/png", |
| 652 | + "content": "data:image/png;base64,second_image", |
| 653 | + }, |
| 654 | + ], |
| 655 | + } |
| 656 | + ] |
| 657 | + |
| 658 | + original_first = messages[0]["content"][1]["content"] |
| 659 | + original_second = messages[0]["content"][2]["content"] |
| 660 | + |
| 661 | + result = redact_blob_message_parts(messages) |
| 662 | + |
| 663 | + # Original should be unchanged |
| 664 | + assert messages[0]["content"][1]["content"] == original_first |
| 665 | + assert messages[0]["content"][2]["content"] == original_second |
| 666 | + |
| 667 | + # Result should be redacted |
| 668 | + assert result[0]["content"][0]["text"] == "Compare these images" |
| 669 | + assert result[0]["content"][1]["content"] == BLOB_DATA_SUBSTITUTE |
| 670 | + assert result[0]["content"][2]["content"] == BLOB_DATA_SUBSTITUTE |
| 671 | + |
| 672 | + def test_redacts_blobs_in_multiple_messages(self): |
| 673 | + """Test that blob parts are redacted across multiple messages without mutation""" |
| 674 | + messages = [ |
| 675 | + { |
| 676 | + "role": "user", |
| 677 | + "content": [ |
| 678 | + {"text": "First message", "type": "text"}, |
| 679 | + { |
| 680 | + "type": "blob", |
| 681 | + "modality": "image", |
| 682 | + "content": "data:image/jpeg;base64,first", |
| 683 | + }, |
| 684 | + ], |
| 685 | + }, |
| 686 | + { |
| 687 | + "role": "assistant", |
| 688 | + "content": "I see the image.", |
| 689 | + }, |
| 690 | + { |
| 691 | + "role": "user", |
| 692 | + "content": [ |
| 693 | + {"text": "Second message", "type": "text"}, |
| 694 | + { |
| 695 | + "type": "blob", |
| 696 | + "modality": "image", |
| 697 | + "content": "data:image/jpeg;base64,second", |
| 698 | + }, |
| 699 | + ], |
| 700 | + }, |
| 701 | + ] |
| 702 | + |
| 703 | + original_first = messages[0]["content"][1]["content"] |
| 704 | + original_second = messages[2]["content"][1]["content"] |
| 705 | + |
| 706 | + result = redact_blob_message_parts(messages) |
| 707 | + |
| 708 | + # Original should be unchanged |
| 709 | + assert messages[0]["content"][1]["content"] == original_first |
| 710 | + assert messages[2]["content"][1]["content"] == original_second |
| 711 | + |
| 712 | + # Result should be redacted |
| 713 | + assert result[0]["content"][1]["content"] == BLOB_DATA_SUBSTITUTE |
| 714 | + assert result[1]["content"] == "I see the image." # Unchanged |
| 715 | + assert result[2]["content"][1]["content"] == BLOB_DATA_SUBSTITUTE |
| 716 | + |
| 717 | + def test_no_blobs_returns_original_list(self): |
| 718 | + """Test that messages without blobs are returned as-is (performance optimization)""" |
| 719 | + messages = [ |
| 720 | + {"role": "user", "content": "Simple text message"}, |
| 721 | + {"role": "assistant", "content": "Simple response"}, |
| 722 | + ] |
| 723 | + |
| 724 | + result = redact_blob_message_parts(messages) |
| 725 | + |
| 726 | + # Should return the same list object when no blobs present |
| 727 | + assert result is messages |
| 728 | + |
| 729 | + def test_handles_non_dict_messages(self): |
| 730 | + """Test that non-dict messages are handled gracefully""" |
| 731 | + messages = [ |
| 732 | + "string message", |
| 733 | + {"role": "user", "content": "text"}, |
| 734 | + None, |
| 735 | + 123, |
| 736 | + ] |
| 737 | + |
| 738 | + result = redact_blob_message_parts(messages) |
| 739 | + |
| 740 | + # Should return same list since no blobs |
| 741 | + assert result is messages |
| 742 | + |
| 743 | + def test_handles_non_dict_content_items(self): |
| 744 | + """Test that non-dict content items in arrays are handled""" |
| 745 | + messages = [ |
| 746 | + { |
| 747 | + "role": "user", |
| 748 | + "content": [ |
| 749 | + "string item", |
| 750 | + {"text": "text item", "type": "text"}, |
| 751 | + None, |
| 752 | + ], |
| 753 | + } |
| 754 | + ] |
| 755 | + |
| 756 | + result = redact_blob_message_parts(messages) |
| 757 | + |
| 758 | + # Should return same list since no blobs |
| 759 | + assert result is messages |
0 commit comments