FIX: embed only text pieces in multimodal writes instead of failing them

Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com>

Author: ManoharPaturi

pyrit/memory/memory_interface.py

@@ -2819,7 +2819,13 @@ def add_message_to_memory(self, *, request: Message) -> None:
         self.add_message_pieces_to_memory(message_pieces=message_pieces)
 
         if self.memory_embedding:
-            for piece in message_pieces:
+            # Embeddings power text similarity search and only describe text pieces.
+            # A multimodal message (e.g. text plus an image_path piece) must still
+            # persist instead of failing the whole write, and pieces flagged
+            # not_in_memory are never persisted, so no embedding row may reference them.
+            for piece in pieces_to_persist:
+                if piece.converted_value_data_type != "text":
+                    continue
                 embedding_entry = self.memory_embedding.generate_embedding_memory_data(message_piece=piece)
                 embedding_entries.append(embedding_entry)
 

tests/unit/memory/test_sqlite_memory.py

@@ -22,11 +22,20 @@
 from pyrit.common.singleton import Singleton
 from pyrit.converter.base64_converter import Base64Converter
 from pyrit.memory.alembic.versions.ab8f2c1a9d07_pre_alembic_release_schema import INITIAL_METADATA
+from pyrit.memory.memory_embedding import MemoryEmbedding
 from pyrit.memory.memory_models import EmbeddingDataEntry, PromptMemoryEntry
 from pyrit.memory.migration import run_schema_migrations
 from pyrit.memory.sqlite_memory import SQLiteMemory
 from pyrit.memory.storage.serializers import set_message_piece_sha256_async
-from pyrit.models import Conversation, MessagePiece, flatten_to_message_pieces
+from pyrit.models import (
+    Conversation,
+    EmbeddingData,
+    EmbeddingResponse,
+    EmbeddingUsageInformation,
+    Message,
+    MessagePiece,
+    flatten_to_message_pieces,
+)
 from pyrit.prompt_target.text_target import TextTarget
 from unit.mocks import get_sample_conversation_entries
 
@@ -456,6 +465,79 @@ def test_insert_embedding_entry(sqlite_instance):
         assert persisted_embedding_entry.embedding_type_name == "test_type"
 
 
+class _MockEmbeddingGenerator:
+    """Duck-typed ``EmbeddingSupport`` stand-in so no network call is needed."""
+
+    def generate_text_embedding(self, text: str, **kwargs) -> EmbeddingResponse:
+        return EmbeddingResponse(
+            model="mock_model",
+            object="mock_object",
+            usage=EmbeddingUsageInformation(prompt_tokens=0, total_tokens=0),
+            data=[EmbeddingData(embedding=[0.5], index=0, object="mock_object")],
+        )
+
+
+def test_add_multimodal_message_with_embedding_persists_and_skips_non_text_pieces(sqlite_instance):
+    """A multimodal message must persist when embeddings are enabled.
+
+    ``MemoryEmbedding.generate_embedding_memory_data`` raises for non-text pieces, so
+    embedding only the text pieces (instead of failing the whole write) keeps multimodal
+    sends working while text similarity search still covers the embeddable content.
+    """
+    sqlite_instance.memory_embedding = MemoryEmbedding(embedding_model=_MockEmbeddingGenerator())
+    sqlite_instance.add_conversation_to_memory(conversation=Conversation(conversation_id="multimodal-conversation"))
+
+    text_piece = MessagePiece(
+        role="user",
+        original_value="describe this image",
+        conversation_id="multimodal-conversation",
+    )
+    image_piece = MessagePiece(
+        role="user",
+        original_value="/tmp/image.png",
+        original_value_data_type="image_path",
+        converted_value="/tmp/image.png",
+        converted_value_data_type="image_path",
+        conversation_id="multimodal-conversation",
+    )
+
+    sqlite_instance.add_message_to_memory(request=Message(message_pieces=[text_piece, image_piece]))
+
+    persisted = sqlite_instance.get_message_pieces(conversation_id="multimodal-conversation")
+    assert len(persisted) == 2
+    with sqlite_instance.get_session() as session:
+        embedding_entries = session.query(EmbeddingDataEntry).all()
+        assert len(embedding_entries) == 1
+        assert embedding_entries[0].id == text_piece.id
+
+
+def test_add_message_with_embedding_skips_pieces_not_in_memory(sqlite_instance):
+    """Pieces flagged ``not_in_memory`` are never persisted, so they get no embedding row."""
+    sqlite_instance.memory_embedding = MemoryEmbedding(embedding_model=_MockEmbeddingGenerator())
+    sqlite_instance.add_conversation_to_memory(conversation=Conversation(conversation_id="ephemeral-conversation"))
+
+    persisted_piece = MessagePiece(
+        role="user",
+        original_value="persisted text",
+        conversation_id="ephemeral-conversation",
+    )
+    ephemeral_piece = MessagePiece(
+        role="user",
+        original_value="ephemeral text",
+        conversation_id="ephemeral-conversation",
+        not_in_memory=True,
+    )
+
+    sqlite_instance.add_message_to_memory(request=Message(message_pieces=[persisted_piece, ephemeral_piece]))
+
+    persisted = sqlite_instance.get_message_pieces(conversation_id="ephemeral-conversation")
+    assert [piece.original_value for piece in persisted] == ["persisted text"]
+    with sqlite_instance.get_session() as session:
+        embedding_entries = session.query(EmbeddingDataEntry).all()
+        assert len(embedding_entries) == 1
+        assert embedding_entries[0].id == persisted_piece.id
+
+
 def test_disable_embedding(sqlite_instance):
     sqlite_instance.disable_embedding()