feat: allow to limit how many elements retrieve (qdrant) (#14904)

run-llama · Jul 24, 2024 · 479d72a · 479d72a
1 parent 34dec27
commit 479d72a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/...s/vector_stores/llama-index-vector-stores-qdrant/llama_index/vector_stores/qdrant/base.py b/...s/vector_stores/llama-index-vector-stores-qdrant/llama_index/vector_stores/qdrant/base.py
@@ -298,6 +298,7 @@ def get_nodes(
         self,
         node_ids: Optional[List[str]] = None,
         filters: Optional[MetadataFilters] = None,
+        limit: Optional[int] = None,
     ) -> List[BaseNode]:
         """
         Get nodes from the index.
@@ -316,6 +317,10 @@ def get_nodes(
                     has_id=node_ids,
                 )
             ]
+            # If we pass a node_ids list,
+            # we can limit the search to only those nodes
+            # or less if limit is provided
+            limit = len(node_ids) if limit is None else min(len(node_ids), limit)
 
         if filters is not None:
             filter = self._build_subfilter(filters)
@@ -337,7 +342,7 @@ def get_nodes(
 
         response = self._client.scroll(
             collection_name=self.collection_name,
-            limit=9999,
+            limit=limit or 9999,
             scroll_filter=filter,
         )
 
@@ -347,6 +352,7 @@ async def aget_nodes(
         self,
         node_ids: Optional[List[str]] = None,
         filters: Optional[MetadataFilters] = None,
+        limit: Optional[int] = None,
     ) -> List[BaseNode]:
         """
         Asynchronous method to get nodes from the index.
@@ -365,6 +371,10 @@ async def aget_nodes(
                     has_id=node_ids,
                 )
             ]
+            # If we pass a node_ids list,
+            # we can limit the search to only those nodes
+            # or less if limit is provided
+            limit = len(node_ids) if limit is None else min(len(node_ids), limit)
 
         if filters is not None:
             filter = self._build_subfilter(filters)
@@ -377,7 +387,7 @@ async def aget_nodes(
 
         response = await self._aclient.scroll(
             collection_name=self.collection_name,
-            limit=9999,
+            limit=limit or 9999,
             scroll_filter=filter,
         )
 

diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-qdrant/pyproject.toml b/llama-index-integrations/vector_stores/llama-index-vector-stores-qdrant/pyproject.toml
@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-vector-stores-qdrant"
 readme = "README.md"
-version = "0.2.13"
+version = "0.2.14"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<3.13"