Update default batch_retrieve for RAG system classes (#441)

* default batch retrieve for rag systems

* coverage

* more coverage

* parametrize no encode

* add not implemented batch retrieval async rag system test

* changelog

Author: nerdai

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 
 ### Added
 
+- Update `batch_retrieve` for RAGSystems to use `batch_retrieve` of Knowledge Stores if implemented (#441)
 - Implement `batch_retrieve` for Qdrant sync/async knowledge stores (#439)
 - Add `batch_retrieve` to KnowledgeStore classes that raise `NotImplementedError` by default (#436)
 - Add batch methods for RAGSystem (#270)

src/fed_rag/core/no_encode_rag_system/_asynchronous.py

@@ -66,14 +66,18 @@ async def batch_retrieve(
         self, queries: list[str]
     ) -> list[list[SourceNode]]:
         """Batch retrieve from KnowledgeStore."""
-        # TODO: move this to knowledge store batch retrieve once implemented
-        raw_retrieval_tasks = [
-            self.knowledge_store.retrieve(
-                query=query, top_k=self.rag_config.top_k
+        try:
+            raw_retrieval_results = await self.knowledge_store.batch_retrieve(
+                queries=queries, top_k=self.rag_config.top_k
             )
-            for query in queries
-        ]
-        raw_retrieval_results = await asyncio.gather(*raw_retrieval_tasks)
+        except NotImplementedError:
+            raw_retrieval_tasks = [
+                self.knowledge_store.retrieve(
+                    query=query, top_k=self.rag_config.top_k
+                )
+                for query in queries
+            ]
+            raw_retrieval_results = await asyncio.gather(*raw_retrieval_tasks)
         return [
             [SourceNode(score=el[0], node=el[1]) for el in raw_result]
             for raw_result in raw_retrieval_results

src/fed_rag/core/no_encode_rag_system/_synchronous.py

@@ -63,13 +63,17 @@ def retrieve(self, query: str) -> list[SourceNode]:
 
     def batch_retrieve(self, queries: list[str]) -> list[list[SourceNode]]:
         """Batch retrieve from NoEncodeKnowledgeStore."""
-        # TODO: move this to knowledge store batch retrieve once implemented
-        raw_retrieval_results = [
-            self.knowledge_store.retrieve(
-                query=query, top_k=self.rag_config.top_k
+        try:
+            raw_retrieval_results = self.knowledge_store.batch_retrieve(
+                queries=queries, top_k=self.rag_config.top_k
             )
-            for query in queries
-        ]
+        except NotImplementedError:
+            raw_retrieval_results = [
+                self.knowledge_store.retrieve(
+                    query=query, top_k=self.rag_config.top_k
+                )
+                for query in queries
+            ]
         return [
             [SourceNode(score=el[0], node=el[1]) for el in raw_result]
             for raw_result in raw_retrieval_results

src/fed_rag/core/rag_system/_asynchronous.py

@@ -55,7 +55,6 @@ async def batch_query(self, queries: list[str]) -> list[RAGResponse]:
     async def retrieve(self, query: str) -> list[SourceNode]:
         """Retrieve from KnowledgeStore."""
         query_emb: list[float] = self.retriever.encode_query(query).tolist()
-        # TODO: move this to knowledge store batch retrieve once implemented
         raw_retrieval_result = await self.knowledge_store.retrieve(
             query_emb=query_emb, top_k=self.rag_config.top_k
         )
@@ -70,13 +69,19 @@ async def batch_retrieve(
         query_embs: list[list[float]] = self.retriever.encode_query(
             queries
         ).tolist()
-        raw_retrieval_tasks = [
-            self.knowledge_store.retrieve(
-                query_emb=query_emb, top_k=self.rag_config.top_k
+        try:
+            raw_retrieval_results = await self.knowledge_store.batch_retrieve(
+                query_embs=query_embs, top_k=self.rag_config.top_k
             )
-            for query_emb in query_embs
-        ]
-        raw_retrieval_results = await asyncio.gather(*raw_retrieval_tasks)
+        except NotImplementedError:
+            raw_retrieval_tasks = [
+                self.knowledge_store.retrieve(
+                    query_emb=query_emb, top_k=self.rag_config.top_k
+                )
+                for query_emb in query_embs
+            ]
+            raw_retrieval_results = await asyncio.gather(*raw_retrieval_tasks)
+
         return [
             [SourceNode(score=el[0], node=el[1]) for el in raw_result]
             for raw_result in raw_retrieval_results

src/fed_rag/core/rag_system/_synchronous.py

@@ -66,13 +66,18 @@ def batch_retrieve(self, queries: list[str]) -> list[list[SourceNode]]:
         query_embs: list[list[float]] = self.retriever.encode_query(
             queries
         ).tolist()
-        # TODO: move this to knowledge store batch retrieve once implemented
-        raw_retrieval_results = [
-            self.knowledge_store.retrieve(
-                query_emb=query_emb, top_k=self.rag_config.top_k
+        try:
+            raw_retrieval_results = self.knowledge_store.batch_retrieve(
+                query_embs=query_embs, top_k=self.rag_config.top_k
             )
-            for query_emb in query_embs
-        ]
+        except NotImplementedError:
+            raw_retrieval_results = [
+                self.knowledge_store.retrieve(
+                    query_emb=query_emb, top_k=self.rag_config.top_k
+                )
+                for query_emb in query_embs
+            ]
+
         return [
             [SourceNode(score=el[0], node=el[1]) for el in raw_result]
             for raw_result in raw_retrieval_results

src/fed_rag/knowledge_stores/in_memory.py

@@ -46,7 +46,6 @@ def cosine_sim(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
         similarities = similarities.to("cpu")
     similarities = similarities.tolist()[0]
     zipped = list(zip(nodes, similarities))
-    # scores.sort(key=lambda tup: tup[1], reverse=True)
     sorted_similarities = sorted(zipped, key=lambda row: row[1], reverse=True)
     return sorted_similarities[:top_k]
 
@@ -83,7 +82,6 @@ def load_nodes(self, nodes: list[KnowledgeNode]) -> None:
     def retrieve(
         self, query_emb: list[float], top_k: int = DEFAULT_TOP_K
     ) -> list[tuple[float, KnowledgeNode]]:
-        # all_nodes = list(self._data.values())
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         query_emb = torch.tensor(query_emb).to(device)
         if not torch.is_tensor(self._data_storage):

tests/rag_system/conftest.py

@@ -101,6 +101,38 @@ async def load(self) -> None:
         pass
 
 
+class DummyAsyncNoBatchRetrievalKnowledgeStore(BaseAsyncKnowledgeStore):
+    nodes: list[KnowledgeNode] = []
+
+    async def load_node(self, node: KnowledgeNode) -> None:
+        self.nodes.append(node)
+
+    async def retrieve(
+        self, query_emb: list[float], top_k: int
+    ) -> list[tuple[float, KnowledgeNode]]:
+        return []
+
+    async def batch_retrieve(
+        self, query_embs: list[list[float]], top_k: int
+    ) -> list[list[tuple[float, KnowledgeNode]]]:
+        raise NotImplementedError
+
+    async def delete_node(self, node_id: str) -> bool:
+        return True
+
+    async def clear(self) -> None:
+        self.nodes.clear()
+
+    async def count(self) -> int:
+        return len(self.nodes)
+
+    async def persist(self) -> None:
+        pass
+
+    async def load(self) -> None:
+        pass
+
+
 @pytest.fixture
 def mock_retriever() -> MockRetriever:
     return MockRetriever()

tests/rag_system/test_async_no_encode_rag_system.py

@@ -35,7 +35,47 @@ async def retrieve(
     async def batch_retrieve(
         self, queries: list[str], top_k: int
     ) -> list[list[tuple[float, KnowledgeNode]]]:
-        return [[]]
+        return [
+            [(ix, n) for ix, n in enumerate(self.nodes[:top_k])]
+            for jx in range(len(queries))
+        ]
+
+    async def delete_node(self, node_id: str) -> bool:
+        return True
+
+    async def clear(self) -> None:
+        self.nodes.clear()
+
+    async def count(self) -> int:
+        return len(self.nodes)
+
+    async def persist(self) -> None:
+        pass
+
+    async def load(self) -> None:
+        pass
+
+
+class DummyNoEncodeNoBatchRetrievalKnowledgeStore(
+    BaseAsyncNoEncodeKnowledgeStore
+):
+    nodes: list[KnowledgeNode] = []
+
+    async def load_node(self, node: KnowledgeNode) -> None:
+        self.nodes.append(node)
+
+    async def load_nodes(self, nodes: list[KnowledgeNode]) -> None:
+        await asyncio.gather(*(self.load_node(n) for n in nodes))
+
+    async def retrieve(
+        self, query: str, top_k: int
+    ) -> list[tuple[float, KnowledgeNode]]:
+        return [(ix, n) for ix, n in enumerate(self.nodes[:top_k])]
+
+    async def batch_retrieve(
+        self, queries: list[str], top_k: int
+    ) -> list[list[tuple[float, KnowledgeNode]]]:
+        raise NotImplementedError
 
     async def delete_node(self, node_id: str) -> bool:
         return True
@@ -64,6 +104,24 @@ async def dummy_store() -> BaseAsyncNoEncodeKnowledgeStore:
     return dummy_store
 
 
+@pytest.fixture()
+async def dummy_store_no_batch_retrieval() -> BaseAsyncNoEncodeKnowledgeStore:
+    dummy_store = DummyNoEncodeNoBatchRetrievalKnowledgeStore()
+    nodes = [
+        KnowledgeNode(node_type=NodeType.TEXT, text_content="Dummy text")
+        for _ in range(5)
+    ]
+    await dummy_store.load_nodes(nodes)
+    return dummy_store
+
+
+@pytest.fixture()
+def knowledge_store(
+    request: pytest.FixtureRequest,
+) -> BaseAsyncNoEncodeKnowledgeStore:
+    return request.getfixturevalue(request.param)
+
+
 def test_rag_system_init(
     mock_generator: BaseGenerator,
     dummy_store: BaseAsyncNoEncodeKnowledgeStore,
@@ -307,17 +365,22 @@ async def test_rag_system_format_context(
 
 
 @pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "knowledge_store",
+    ["dummy_store", "dummy_store_no_batch_retrieval"],
+    indirect=True,
+)
 async def test_rag_system_batch_retrieve(
+    knowledge_store: BaseAsyncNoEncodeKnowledgeStore,
     mock_generator: BaseGenerator,
-    dummy_store: BaseAsyncNoEncodeKnowledgeStore,
 ) -> None:
     # build rag system
     rag_config = RAGConfig(
         top_k=2,
     )
     rag_system = AsyncNoEncodeRAGSystem(
         generator=mock_generator,
-        knowledge_store=dummy_store,
+        knowledge_store=knowledge_store,
         rag_config=rag_config,
     )
 

tests/rag_system/test_async_rag_system.py

@@ -5,11 +5,34 @@
 
 from fed_rag import AsyncRAGSystem, RAGConfig
 from fed_rag.base.generator import BaseGenerator
+from fed_rag.base.knowledge_store import BaseAsyncKnowledgeStore
 from fed_rag.base.retriever import BaseRetriever
 from fed_rag.data_structures import KnowledgeNode, SourceNode
 from fed_rag.exceptions import RAGSystemError
 
-from .conftest import DummyAsyncKnowledgeStore, MockGenerator, MockRetriever
+from .conftest import (
+    DummyAsyncKnowledgeStore,
+    DummyAsyncNoBatchRetrievalKnowledgeStore,
+    MockGenerator,
+    MockRetriever,
+)
+
+
+@pytest.fixture()
+def dummy_store() -> BaseAsyncKnowledgeStore:
+    return DummyAsyncKnowledgeStore()
+
+
+@pytest.fixture()
+def dummy_store_no_batch_retrieval() -> BaseAsyncKnowledgeStore:
+    return DummyAsyncNoBatchRetrievalKnowledgeStore()
+
+
+@pytest.fixture()
+def knowledge_store(
+    request: pytest.FixtureRequest,
+) -> BaseAsyncKnowledgeStore:
+    return request.getfixturevalue(request.param)
 
 
 @pytest.mark.asyncio
@@ -208,9 +231,15 @@ async def test_rag_system_retrieve(
 
 
 @pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "knowledge_store",
+    ["dummy_store", "dummy_store_no_batch_retrieval"],
+    indirect=True,
+)
 @patch.object(MockRetriever, "encode_query")
 async def test_rag_system_batch_retrieve(
     mock_encode_query: MagicMock,
+    knowledge_store: BaseAsyncKnowledgeStore,
     mock_generator: BaseGenerator,
     mock_retriever: MockRetriever,
     knowledge_nodes: list[KnowledgeNode],
@@ -221,7 +250,6 @@ async def test_rag_system_batch_retrieve(
     )
 
     # build rag system
-    knowledge_store = DummyAsyncKnowledgeStore()
     await knowledge_store.load_nodes(nodes=knowledge_nodes)
     rag_config = RAGConfig(
         top_k=2,