fix: prevent in-place mutation of documents in Document Embedders (#9693)

* fix: prevent in-place mutation of documents after embeddings by using deepcopy

* Add tests

* use from dataclasses import replace instead of deepcopy

* Address PR comments

Author: abdokaseb

haystack/components/embedders/hugging_face_api_document_embedder.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from dataclasses import replace
 from typing import Any, Optional, Union
 
 from tqdm import tqdm
@@ -328,10 +329,11 @@ def run(self, documents: list[Document]):
 
         embeddings = self._embed_batch(texts_to_embed=texts_to_embed, batch_size=self.batch_size)
 
+        new_documents = []
         for doc, emb in zip(documents, embeddings):
-            doc.embedding = emb
+            new_documents.append(replace(doc, embedding=emb))
 
-        return {"documents": documents}
+        return {"documents": new_documents}
 
     @component.output_types(documents=list[Document])
     async def run_async(self, documents: list[Document]):
@@ -355,7 +357,8 @@ async def run_async(self, documents: list[Document]):
 
         embeddings = await self._embed_batch_async(texts_to_embed=texts_to_embed, batch_size=self.batch_size)
 
+        new_documents = []
         for doc, emb in zip(documents, embeddings):
-            doc.embedding = emb
+            new_documents.append(replace(doc, embedding=emb))
 
-        return {"documents": documents}
+        return {"documents": new_documents}

haystack/components/embedders/image/sentence_transformers_doc_image_embedder.py

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from copy import copy
+from dataclasses import replace
 from typing import Any, Literal, Optional
 
 from haystack import Document, component, default_from_dict, default_to_dict
@@ -281,10 +281,12 @@ def run(self, documents: list[Document]) -> dict[str, list[Document]]:
 
         docs_with_embeddings = []
         for doc, emb in zip(documents, embeddings):
-            copied_doc = copy(doc)
-            copied_doc.embedding = emb
             # we store this information for later inspection
-            copied_doc.meta["embedding_source"] = {"type": "image", "file_path_meta_field": self.file_path_meta_field}
-            docs_with_embeddings.append(copied_doc)
+            new_meta = {
+                **doc.meta,
+                "embedding_source": {"type": "image", "file_path_meta_field": self.file_path_meta_field},
+            }
+            new_doc = replace(doc, meta=new_meta, embedding=emb)
+            docs_with_embeddings.append(new_doc)
 
         return {"documents": docs_with_embeddings}

haystack/components/embedders/openai_document_embedder.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import os
+from dataclasses import replace
 from typing import Any, Optional
 
 from more_itertools import batched
@@ -307,11 +308,14 @@ def run(self, documents: list[Document]):
 
         doc_ids_to_embeddings, meta = self._embed_batch(texts_to_embed=texts_to_embed, batch_size=self.batch_size)
 
-        doc_id_to_document = {doc.id: doc for doc in documents}
-        for doc_id, emb in doc_ids_to_embeddings.items():
-            doc_id_to_document[doc_id].embedding = emb
+        new_documents = []
+        for doc in documents:
+            if doc.id in doc_ids_to_embeddings:
+                new_documents.append(replace(doc, embedding=doc_ids_to_embeddings[doc.id]))
+            else:
+                new_documents.append(replace(doc))
 
-        return {"documents": list(doc_id_to_document.values()), "meta": meta}
+        return {"documents": new_documents, "meta": meta}
 
     @component.output_types(documents=list[Document], meta=dict[str, Any])
     async def run_async(self, documents: list[Document]):
@@ -338,8 +342,11 @@ async def run_async(self, documents: list[Document]):
             texts_to_embed=texts_to_embed, batch_size=self.batch_size
         )
 
-        doc_id_to_document = {doc.id: doc for doc in documents}
-        for doc_id, emb in doc_ids_to_embeddings.items():
-            doc_id_to_document[doc_id].embedding = emb
+        new_documents = []
+        for doc in documents:
+            if doc.id in doc_ids_to_embeddings:
+                new_documents.append(replace(doc, embedding=doc_ids_to_embeddings[doc.id]))
+            else:
+                new_documents.append(replace(doc))
 
-        return {"documents": list(doc_id_to_document.values()), "meta": meta}
+        return {"documents": new_documents, "meta": meta}

haystack/components/embedders/sentence_transformers_document_embedder.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from dataclasses import replace
 from typing import Any, Literal, Optional
 
 from haystack import Document, component, default_from_dict, default_to_dict
@@ -257,7 +258,8 @@ def run(self, documents: list[Document]):
             **(self.encode_kwargs if self.encode_kwargs else {}),
         )
 
+        new_documents = []
         for doc, emb in zip(documents, embeddings):
-            doc.embedding = emb
+            new_documents.append(replace(doc, embedding=emb))
 
-        return {"documents": documents}
+        return {"documents": new_documents}

releasenotes/notes/fix-embedder-deepcopy-c21216149f700993.yaml

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    Prevented in-place mutation of input `Document` objects in all `DocumentEmbedder` components
+    by creating copies with `dataclasses.replace` before processing.

test/components/embedders/image/test_sentence_transformers_doc_image_embedder.py

@@ -209,13 +209,16 @@ def test_run(self, test_files_path):
 
         assert isinstance(result["documents"], list)
         assert len(result["documents"]) == len(documents)
-        for doc in result["documents"]:
-            assert isinstance(doc, Document)
-            assert isinstance(doc.embedding, list)
-            assert isinstance(doc.embedding[0], float)
-            assert "embedding_source" in doc.meta
-            assert doc.meta["embedding_source"]["type"] == "image"
-            assert "file_path_meta_field" in doc.meta["embedding_source"]
+        for doc, new_doc in zip(documents, result["documents"]):
+            assert doc.embedding is None
+            assert new_doc is not doc
+            assert isinstance(new_doc, Document)
+            assert isinstance(new_doc.embedding, list)
+            assert isinstance(new_doc.embedding[0], float)
+            assert "embedding_source" not in doc.meta
+            assert "embedding_source" in new_doc.meta
+            assert new_doc.meta["embedding_source"]["type"] == "image"
+            assert "file_path_meta_field" in new_doc.meta["embedding_source"]
 
     def test_run_no_warmup(self):
         embedder = SentenceTransformersDocumentImageEmbedder(model="model")
@@ -338,11 +341,14 @@ def test_live_run(self, test_files_path, monkeypatch):
 
         result = embedder.run(documents=documents)
         assert len(result["documents"]) == len(documents)
-        for doc in result["documents"]:
-            assert isinstance(doc, Document)
-            assert isinstance(doc.embedding, list)
-            assert len(doc.embedding) == 512
-            assert all(isinstance(x, float) for x in doc.embedding)
-            assert "embedding_source" in doc.meta
-            assert doc.meta["embedding_source"]["type"] == "image"
-            assert "file_path_meta_field" in doc.meta["embedding_source"]
+        for doc, new_doc in zip(documents, result["documents"]):
+            assert doc.embedding is None
+            assert new_doc is not doc
+            assert isinstance(new_doc, Document)
+            assert isinstance(new_doc.embedding, list)
+            assert len(new_doc.embedding) == 512
+            assert all(isinstance(x, float) for x in new_doc.embedding)
+            assert "embedding_source" not in doc.meta
+            assert "embedding_source" in new_doc.meta
+            assert new_doc.meta["embedding_source"]["type"] == "image"
+            assert "file_path_meta_field" in new_doc.meta["embedding_source"]

test/components/embedders/test_azure_document_embedder.py

@@ -265,11 +265,13 @@ def test_run(self):
 
         assert isinstance(documents_with_embeddings, list)
         assert len(documents_with_embeddings) == len(docs)
-        for doc in documents_with_embeddings:
-            assert isinstance(doc, Document)
-            assert isinstance(doc.embedding, list)
-            assert len(doc.embedding) == 1536
-            assert all(isinstance(x, float) for x in doc.embedding)
+        for doc, new_doc in zip(docs, documents_with_embeddings):
+            assert doc.embedding is None
+            assert new_doc is not doc
+            assert isinstance(new_doc, Document)
+            assert isinstance(new_doc.embedding, list)
+            assert len(new_doc.embedding) == 1536
+            assert all(isinstance(x, float) for x in new_doc.embedding)
 
         assert metadata["usage"]["prompt_tokens"] == 15
         assert metadata["usage"]["total_tokens"] == 15

test/components/embedders/test_hugging_face_api_document_embedder.py

@@ -287,7 +287,6 @@ def test_run(self, mock_check_valid_model):
             Document(content="I love cheese", meta={"topic": "Cuisine"}),
             Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}),
         ]
-
         with patch("huggingface_hub.InferenceClient.feature_extraction") as mock_embedding_patch:
             mock_embedding_patch.side_effect = mock_embedding_generation
 
@@ -316,11 +315,13 @@ def test_run(self, mock_check_valid_model):
 
         assert isinstance(documents_with_embeddings, list)
         assert len(documents_with_embeddings) == len(docs)
-        for doc in documents_with_embeddings:
-            assert isinstance(doc, Document)
-            assert isinstance(doc.embedding, list)
-            assert len(doc.embedding) == 384
-            assert all(isinstance(x, float) for x in doc.embedding)
+        for doc, new_doc in zip(docs, documents_with_embeddings):
+            assert doc.embedding is None
+            assert new_doc is not doc
+            assert isinstance(new_doc, Document)
+            assert isinstance(new_doc.embedding, list)
+            assert len(new_doc.embedding) == 384
+            assert all(isinstance(x, float) for x in new_doc.embedding)
 
     def test_run_custom_batch_size(self, mock_check_valid_model):
         docs = [

test/components/embedders/test_openai_document_embedder.py

@@ -281,11 +281,13 @@ def test_run(self):
 
         assert isinstance(documents_with_embeddings, list)
         assert len(documents_with_embeddings) == len(docs)
-        for doc in documents_with_embeddings:
-            assert isinstance(doc, Document)
-            assert isinstance(doc.embedding, list)
-            assert len(doc.embedding) == 1536
-            assert all(isinstance(x, float) for x in doc.embedding)
+        for doc, new_doc in zip(docs, documents_with_embeddings):
+            assert doc.embedding is None
+            assert new_doc is not doc
+            assert isinstance(new_doc, Document)
+            assert isinstance(new_doc.embedding, list)
+            assert len(new_doc.embedding) == 1536
+            assert all(isinstance(x, float) for x in new_doc.embedding)
 
         assert "text" in result["meta"]["model"] and "ada" in result["meta"]["model"], (
             "The model name does not contain 'text' and 'ada'"
@@ -311,11 +313,13 @@ async def test_run_async(self):
 
         assert isinstance(documents_with_embeddings, list)
         assert len(documents_with_embeddings) == len(docs)
-        for doc in documents_with_embeddings:
-            assert isinstance(doc, Document)
-            assert isinstance(doc.embedding, list)
-            assert len(doc.embedding) == 1536
-            assert all(isinstance(x, float) for x in doc.embedding)
+        for doc, new_doc in zip(docs, documents_with_embeddings):
+            assert doc.embedding is None
+            assert new_doc is not doc
+            assert isinstance(new_doc, Document)
+            assert isinstance(new_doc.embedding, list)
+            assert len(new_doc.embedding) == 1536
+            assert all(isinstance(x, float) for x in new_doc.embedding)
 
         assert "text" in result["meta"]["model"] and "ada" in result["meta"]["model"], (
             "The model name does not contain 'text' and 'ada'"

test/components/embedders/test_sentence_transformers_document_embedder.py

@@ -293,10 +293,12 @@ def test_run(self):
 
         assert isinstance(result["documents"], list)
         assert len(result["documents"]) == len(documents)
-        for doc in result["documents"]:
-            assert isinstance(doc, Document)
-            assert isinstance(doc.embedding, list)
-            assert isinstance(doc.embedding[0], float)
+        for doc, new_doc in zip(documents, result["documents"]):
+            assert new_doc is not doc
+            assert doc.embedding is None
+            assert isinstance(new_doc, Document)
+            assert isinstance(new_doc.embedding, list)
+            assert isinstance(new_doc.embedding[0], float)
 
     def test_run_wrong_input_format(self):
         embedder = SentenceTransformersDocumentEmbedder(model="model")