Repository: deepset-ai/haystack

Issue:
Add `return_embedding` param to `FilterRetriever`
**Is your feature request related to a problem? Please describe.**
It's not possible to remove the embeddings from the Documents that the FilterRetriever returns. These embeddings fill up the memory and context limit for LLMs if retrieved docs are passed directly to an LLM in Agent

**Describe the solution you'd like**
have `return_embedding` param for FilterRetriever like we have in other retrievers

**Describe alternatives you've considered**
Based on the architectures, there are different approaches to eliminate this, such as using `outputs_to_string` in Agent

**Additional context**
N/A



Discussion:
Maybe to add to this, not only should we add this option to the FilterRetriever but we should also add it to `InMemoryDocumentStore.bm25_retrieval` since often times we use BM25 retrieval even on documents that have embeddings. 

One thing I noticed is that `filter_documents` used by `FilterRetriever` is a part of our `DocumentStore` protocol. So we will need to think if we add `return_embedding` to the protocol and force all downstream DocumentStore's to update or find a more flexible way to expand the protocol. 

**Update:** 
As an alternative we could follow a similar approach as is used by the `OpenSearchDocumentStore` where we add a `return_embedding` parameter to the init method of `InMemoryDocumentStore` and use that within the necessary methods like bm25 retrieval, filter documents, and embedding retrieval
hello here,

Happy to help with this issue :)
@sjrl 

> Update:
> As an alternative we could follow a similar approach as is used by the OpenSearchDocumentStore where we add a
> return_embedding parameter to the init method of InMemoryDocumentStore and use that within the necessary methods 
> like bm25 retrieval, filter documents, and embedding retrieval

Totally agree with this 🤝 , then we can be controlling it from a centralized position in the document store 👍 



Patch:
diff --git a/haystack/document_stores/in_memory/document_store.py b/haystack/document_stores/in_memory/document_store.py
index b5dc67c444..67d4b1d635 100644
--- a/haystack/document_stores/in_memory/document_store.py
+++ b/haystack/document_stores/in_memory/document_store.py
@@ -68,6 +68,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         embedding_similarity_function: Literal["dot_product", "cosine"] = "dot_product",
         index: Optional[str] = None,
         async_executor: Optional[ThreadPoolExecutor] = None,
+        return_embedding: bool = True,
     ):
         """
         Initializes the DocumentStore.
@@ -85,6 +86,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         :param async_executor:
             Optional ThreadPoolExecutor to use for async calls. If not provided, a single-threaded
             executor will be initialized and used.
+        :param return_embedding: Whether to return the embedding of the retrieved Documents. Default is True.
         """
         self.bm25_tokenization_regex = bm25_tokenization_regex
         self.tokenizer = re.compile(bm25_tokenization_regex).findall
@@ -118,6 +120,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
             if async_executor is None
             else async_executor
         )
+        self.return_embedding = return_embedding
 
     def __del__(self):
         """
@@ -355,6 +358,7 @@ def to_dict(self) -> Dict[str, Any]:
             bm25_parameters=self.bm25_parameters,
             embedding_similarity_function=self.embedding_similarity_function,
             index=self.index,
+            return_embedding=self.return_embedding,
         )
 
     @classmethod
@@ -426,8 +430,15 @@ def filter_documents(self, filters: Optional[Dict[str, Any]] = None) -> List[Doc
                 raise ValueError(
                     "Invalid filter syntax. See https://docs.haystack.deepset.ai/docs/metadata-filtering for details."
                 )
-            return [doc for doc in self.storage.values() if document_matches_filter(filters=filters, document=doc)]
-        return list(self.storage.values())
+            docs = [doc for doc in self.storage.values() if document_matches_filter(filters=filters, document=doc)]
+        else:
+            docs = list(self.storage.values())
+
+        if not self.return_embedding:
+            for doc in docs:
+                doc.embedding = None
+
+        return docs
 
     def write_documents(self, documents: List[Document], policy: DuplicatePolicy = DuplicatePolicy.NONE) -> int:
         """
@@ -542,7 +553,12 @@ def bm25_retrieval(
 
             doc_fields = doc.to_dict()
             doc_fields["score"] = score
+
+            if not self.return_embedding and "embedding" in doc_fields:
+                doc_fields.pop("embedding")
+
             return_document = Document.from_dict(doc_fields)
+
             return_documents.append(return_document)
 
         return return_documents
@@ -553,7 +569,7 @@ def embedding_retrieval(  # pylint: disable=too-many-positional-arguments
         filters: Optional[Dict[str, Any]] = None,
         top_k: int = 10,
         scale_score: bool = False,
-        return_embedding: bool = False,
+        return_embedding: Optional[bool] = False,
     ) -> List[Document]:
         """
         Retrieves documents that are most similar to the query embedding using a vector similarity metric.
@@ -562,14 +578,24 @@ def embedding_retrieval(  # pylint: disable=too-many-positional-arguments
         :param filters: A dictionary with filters to narrow down the search space.
         :param top_k: The number of top documents to retrieve. Default is 10.
         :param scale_score: Whether to scale the scores of the retrieved Documents. Default is False.
-        :param return_embedding: Whether to return the embedding of the retrieved Documents. Default is False.
+        :param return_embedding: Whether to return the embedding of the retrieved Documents.
+            If not provided, the value of the `return_embedding` parameter set at component
+            initialization will be used. Default is False.
         :returns: A list of the top_k documents most relevant to the query.
         """
         if len(query_embedding) == 0 or not isinstance(query_embedding[0], float):
             raise ValueError("query_embedding should be a non-empty list of floats.")
 
-        filters = filters or {}
-        all_documents = self.filter_documents(filters=filters)
+        if filters:
+            if "operator" not in filters and "conditions" not in filters:
+                raise ValueError(
+                    "Invalid filter syntax. See https://docs.haystack.deepset.ai/docs/metadata-filtering for details."
+                )
+            all_documents = [
+                doc for doc in self.storage.values() if document_matches_filter(filters=filters, document=doc)
+            ]
+        else:
+            all_documents = list(self.storage.values())
 
         documents_with_embeddings = [doc for doc in all_documents if doc.embedding is not None]
         if len(documents_with_embeddings) == 0:
@@ -587,12 +613,14 @@ def embedding_retrieval(  # pylint: disable=too-many-positional-arguments
             embedding=query_embedding, documents=documents_with_embeddings, scale_score=scale_score
         )
 
+        resolved_return_embedding = self.return_embedding if return_embedding is None else return_embedding
+
         # create Documents with the similarity score for the top k results
         top_documents = []
         for doc, score in sorted(zip(documents_with_embeddings, scores), key=lambda x: x[1], reverse=True)[:top_k]:
             doc_fields = doc.to_dict()
             doc_fields["score"] = score
-            if return_embedding is False:
+            if resolved_return_embedding is False:
                 doc_fields["embedding"] = None
             top_documents.append(Document.from_dict(doc_fields))
 
diff --git a/releasenotes/notes/update-in-memory-document-store-17f555695caf9d52.yaml b/releasenotes/notes/update-in-memory-document-store-17f555695caf9d52.yaml
new file mode 100644
index 0000000000..92748412d3
--- /dev/null
+++ b/releasenotes/notes/update-in-memory-document-store-17f555695caf9d52.yaml
@@ -0,0 +1,6 @@
+---
+enhancements:
+  - |
+    - Added `return_embedding` parameter inside `InMemoryDocumentStore::init` method.
+    - Updated methods `bm25_retrieval`, and `filter_documents` to use `self.return_embedding` to determine whether embeddings are returned.
+    - Updated tests (test_in_memory & test_in_memory_embedding_retriever) to reflect the changes in the `InMemoryDocumentStore`.


Test patch:
diff --git a/test/components/retrievers/test_sentence_window_retriever.py b/test/components/retrievers/test_sentence_window_retriever.py
index e4ebaead95..65e21360fd 100644
--- a/test/components/retrievers/test_sentence_window_retriever.py
+++ b/test/components/retrievers/test_sentence_window_retriever.py
@@ -76,6 +76,7 @@ def test_to_dict(self):
                         "bm25_tokenization_regex": "(?u)\\b\\w\\w+\\b",
                         "embedding_similarity_function": "dot_product",
                         "index": ANY,
+                        "return_embedding": True,
                     },
                 },
                 "window_size": 3,
diff --git a/test/document_stores/test_in_memory.py b/test/document_stores/test_in_memory.py
index 118336eb54..eb48cb8373 100644
--- a/test/document_stores/test_in_memory.py
+++ b/test/document_stores/test_in_memory.py
@@ -43,6 +43,7 @@ def test_to_dict(self):
                 "bm25_parameters": {},
                 "embedding_similarity_function": "dot_product",
                 "index": store.index,
+                "return_embedding": True,
             },
         }
 
@@ -53,6 +54,7 @@ def test_to_dict_with_custom_init_parameters(self):
             bm25_parameters={"key": "value"},
             embedding_similarity_function="cosine",
             index="my_cool_index",
+            return_embedding=True,
         )
         data = store.to_dict()
         assert data == {
@@ -63,6 +65,7 @@ def test_to_dict_with_custom_init_parameters(self):
                 "bm25_parameters": {"key": "value"},
                 "embedding_similarity_function": "cosine",
                 "index": "my_cool_index",
+                "return_embedding": True,
             },
         }
 
@@ -250,6 +253,47 @@ def test_bm25_retrieval_default_filter(self, document_store: InMemoryDocumentSto
         results = document_store.bm25_retrieval(query="doesn't matter, top_k is 10", top_k=10)
         assert len(results) == 0
 
+    def test_embedding_retrieval_return_embedding_false_on_store(self):
+        # Initialize InMemoryDocumentStore with return_embedding=False
+        docstore = InMemoryDocumentStore(embedding_similarity_function="cosine", return_embedding=False)
+        docs = [
+            Document(content="Hello world", embedding=[0.1, 0.2, 0.3, 0.4]),
+            Document(content="Haystack supports multiple languages", embedding=[1.0, 1.0, 1.0, 1.0]),
+        ]
+        docstore.write_documents(docs)
+
+        # embedding_retrieval should not return embeddings in the documents
+        results = docstore.embedding_retrieval(query_embedding=[0.1, 0.1, 0.1, 0.1], top_k=2)
+        assert all(doc.embedding is None for doc in results)
+
+        # bm25_retrieval should also not return embeddings
+        bm25_results = docstore.bm25_retrieval(query="languages", top_k=2)
+        assert all(doc.embedding is None for doc in bm25_results)
+
+        # filter_documents should not return embeddings
+        filtered_docs = docstore.filter_documents()
+        assert all(doc.embedding is None for doc in filtered_docs)
+
+    def test_embedding_retrieval_override_return_embedding(self):
+        docstore = InMemoryDocumentStore(embedding_similarity_function="cosine", return_embedding=False)
+        docs = [
+            Document(content="Hello world", embedding=[0.1, 0.2, 0.3, 0.4]),
+            Document(content="Haystack supports multiple languages", embedding=[1.0, 1.0, 1.0, 1.0]),
+        ]
+        docstore.write_documents(docs)
+
+        # Overriding return_embedding to True should return embeddings
+        # Query for the embedding that matches both documents by cosine similarity
+        results_with_embedding = docstore.embedding_retrieval(
+            query_embedding=[0.1, 0.2, 0.3, 0.4], top_k=2, return_embedding=True
+        )
+
+        # Assert that the retrieved documents have the expected embeddings
+        assert len(results_with_embedding) == 2
+
+        assert results_with_embedding[0].embedding in ([1.0, 1.0, 1.0, 1.0], [0.1, 0.2, 0.3, 0.4])
+        assert results_with_embedding[1].embedding in ([1.0, 1.0, 1.0, 1.0], [0.1, 0.2, 0.3, 0.4])
+
     def test_embedding_retrieval(self):
         docstore = InMemoryDocumentStore(embedding_similarity_function="cosine")
         # Tests if the embedding retrieval method returns the correct document based on the input query embedding.