[Feature] Implement batch_retrieve for Qdrant knowledge stores (#439)

* implement batch_retrieve for sync

* implement batch retrieve for async

* fix broken tests

* fix tests async

* update qdrant cookbooks with batch retrieval usage

* changelog

Author: nerdai

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 
 ### Added
 
+- 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)
 - `AsyncQdrantKnowledgeStore` (#430)

docs/notebooks/integrations/qdrant.ipynb

@@ -50,7 +50,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "id": "36881aa5-3a6f-4358-985b-b829aac42826",
    "metadata": {},
    "outputs": [],
@@ -94,7 +94,7 @@
      "output_type": "stream",
      "text": [
       "Image 'qdrant/qdrant' already exists locally\n",
-      "Container started with ID: 33dec7bc01280a8b8343ff4da87e8822cbfdd3a242249d1935b581031ae48784\n",
+      "Container started with ID: 7a505aaa0d10843ece53e55fdf42458a41723c6cc7ddcacd1534e36a515db551\n",
       "Container status: running\n",
       "Container logs:\n",
       "           _                 _    \n",
@@ -107,16 +107,16 @@
       "Version: 1.14.0, build: 3617a011\n",
       "Access web UI at http://localhost:6333/dashboard\n",
       "\n",
-      "2025-05-21T19:20:02.141355Z  INFO storage::content_manager::consensus::persistent: Loading raft state from ./storage/raft_state.json    \n",
-      "2025-05-21T19:20:02.145747Z  INFO qdrant: Distributed mode disabled    \n",
-      "2025-05-21T19:20:02.145819Z  INFO qdrant: Telemetry reporting enabled, id: 92001ec9-eb22-4dbf-bcfd-5a800bac4ffb    \n",
-      "2025-05-21T19:20:02.145906Z  INFO qdrant: Inference service is not configured.    \n",
-      "2025-05-21T19:20:02.147806Z  INFO qdrant::actix: TLS disabled for REST API    \n",
-      "2025-05-21T19:20:02.147861Z  INFO qdrant::actix: Qdrant HTTP listening on 6333    \n",
-      "2025-05-21T19:20:02.147878Z  INFO actix_server::builder: Starting 11 workers\n",
-      "2025-05-21T19:20:02.147884Z  INFO actix_server::server: Actix runtime found; starting in Actix runtime\n",
-      "2025-05-21T19:20:02.151909Z  INFO qdrant::tonic: Qdrant gRPC listening on 6334    \n",
-      "2025-05-21T19:20:02.151920Z  INFO qdrant::tonic: TLS disabled for gRPC API    \n",
+      "2025-06-18T02:43:01.654204Z  INFO storage::content_manager::consensus::persistent: Loading raft state from ./storage/raft_state.json    \n",
+      "2025-06-18T02:43:01.658166Z  INFO qdrant: Distributed mode disabled    \n",
+      "2025-06-18T02:43:01.658223Z  INFO qdrant: Telemetry reporting enabled, id: d8e363b9-1c51-4b9a-afaf-18e5a03cb920    \n",
+      "2025-06-18T02:43:01.658285Z  INFO qdrant: Inference service is not configured.    \n",
+      "2025-06-18T02:43:01.660085Z  INFO qdrant::actix: TLS disabled for REST API    \n",
+      "2025-06-18T02:43:01.660146Z  INFO qdrant::actix: Qdrant HTTP listening on 6333    \n",
+      "2025-06-18T02:43:01.660163Z  INFO actix_server::builder: Starting 11 workers\n",
+      "2025-06-18T02:43:01.660175Z  INFO actix_server::server: Actix runtime found; starting in Actix runtime\n",
+      "2025-06-18T02:43:01.663848Z  INFO qdrant::tonic: Qdrant gRPC listening on 6334    \n",
+      "2025-06-18T02:43:01.663858Z  INFO qdrant::tonic: TLS disabled for gRPC API    \n",
       "\n"
      ]
     }
@@ -371,7 +371,7 @@
      "output_type": "stream",
      "text": [
       "Similarity score:  0.49984106\n",
-      "KnowledgeNode:  node_id='d641b713-62dd-46e7-86c2-4e3968ed5339' embedding=None node_type=<NodeType.TEXT: 'text'> text_content=\"title: History of marine biology\\nsection: James Cook\\ntext:  James Cook is well known for his voyages of exploration for the British Navy in which he mapped out a significant amount of the world's uncharted waters. Cook's explorations took him around the world twice and led to countless descriptions of previously unknown plants and animals. Cook's explorations influenced many others and led to a number of scientists examining marine life more closely. Among those influenced was Charles Darwin who went on to make many contributions of his own. \" image_content=None metadata={'id': '140'}\n"
+      "KnowledgeNode:  node_id='3e7e2e62-51ef-4da6-94e7-a1caef85c76e' embedding=None node_type=<NodeType.TEXT: 'text'> text_content=\"title: History of marine biology\\nsection: James Cook\\ntext:  James Cook is well known for his voyages of exploration for the British Navy in which he mapped out a significant amount of the world's uncharted waters. Cook's explorations took him around the world twice and led to countless descriptions of previously unknown plants and animals. Cook's explorations influenced many others and led to a number of scientists examining marine life more closely. Among those influenced was Charles Darwin who went on to make many contributions of his own. \" image_content=None metadata={'id': '140'}\n"
      ]
     }
    ],
@@ -380,6 +380,83 @@
     "print(\"KnowledgeNode: \", knowledge_node)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "1e7532c7-bbb2-4351-8ce1-107556ec35f8",
+   "metadata": {},
+   "source": [
+    "## Batch Retrieval"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "ad598746-9b13-4684-bbc3-611a1da4d886",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "another_query = \"Who is Charles Darwin?\"\n",
+    "another_query_emb = retriever.encode_query(another_query).tolist()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "74ca44a0-799e-4aa2-90da-1eb11b835ce2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch_retrieval_result = knowledge_store.batch_retrieve(\n",
+    "    query_embs=[query_emb, another_query_emb], top_k=1\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "ee08ec6f-fe6c-40be-b93b-2b4a163b02ff",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[(0.49984106,\n",
+       "  KnowledgeNode(node_id='3e7e2e62-51ef-4da6-94e7-a1caef85c76e', embedding=None, node_type=<NodeType.TEXT: 'text'>, text_content=\"title: History of marine biology\\nsection: James Cook\\ntext:  James Cook is well known for his voyages of exploration for the British Navy in which he mapped out a significant amount of the world's uncharted waters. Cook's explorations took him around the world twice and led to countless descriptions of previously unknown plants and animals. Cook's explorations influenced many others and led to a number of scientists examining marine life more closely. Among those influenced was Charles Darwin who went on to make many contributions of his own. \", image_content=None, metadata={'id': '140'}))]"
+      ]
+     },
+     "execution_count": 15,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Who is James Cook?\n",
+    "batch_retrieval_result[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "6bbd04a8-040b-4db3-833e-73db83cfdf47",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[(0.50851303,\n",
+       "  KnowledgeNode(node_id='27462ceb-8f25-4d4b-99b2-0b821697b37e', embedding=None, node_type=<NodeType.TEXT: 'text'>, text_content='title: History of marine biology\\nsection: Charles Darwin\\ntext:  Charles Darwin, best known for his theory of evolution, made many significant contributions to the early study of marine biology. He spent much of his time from 1831 to 1836 on the voyage of HMS Beagle collecting and studying specimens from a variety of marine organisms. It was also on this expedition where Darwin began to study coral reefs and their formation. He came up with the theory that the overall growth of corals is a balance between the growth of corals upward and the sinking of the sea floor. He then came up with the idea that wherever coral atolls would be found, the central island where the coral had started to grow would be gradually subsiding', image_content=None, metadata={'id': '141'}))]"
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Who is Charles Darwin?\n",
+    "batch_retrieval_result[1]"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "463bfac0-bf0d-453e-a47f-b136c069793a",
@@ -390,7 +467,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 17,
+   "id": "4f0d1315-f9eb-4306-a340-4d060511a27f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "knowledge_store.clear()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
    "id": "b2cf76bc-c0b6-4934-9e8a-5374ea16d34e",
    "metadata": {},
    "outputs": [],