removed query_vector as a parameter of query

Author: kkopczynski-cyborg

cyborgdb/client/encrypted_index.py

@@ -401,8 +401,7 @@ def delete(self, ids: List[str]) -> None:
 
     def query(
         self,
-        query_vector: Optional[Union[List[float], np.ndarray]] = None,
-        query_vectors: Optional[Union[np.ndarray, List[List[float]]]] = None,
+        query_vectors: Optional[Union[np.ndarray, List[List[float]], List[float]]] = None,
         query_contents: Optional[str] = None,
         top_k: int = 100,
         n_probes: int = 1,
@@ -423,16 +422,34 @@ def query(
 
             # Determine the correct vector input
             vector_list = None
+            is_single_query = False
 
-            if query_vector is not None or query_contents is not None:
-                if isinstance(query_vector, np.ndarray):
-                    if query_vector.ndim != 1:
-                        raise ValueError("Expected 1D NumPy array for `query_vector`.")
-                    vector_list = query_vector.tolist()
-                elif isinstance(query_vector, list):
-                    if query_vector and isinstance(query_vector[0], (list, np.ndarray)):
-                        raise ValueError("Received nested list in `query_vector`; did you mean to use `query_vectors`?")
-                    vector_list = list(map(float, query_vector))  # Ensure float type
+            if query_vectors is not None:
+                if isinstance(query_vectors, np.ndarray):
+                    if query_vectors.ndim == 1:
+                        # Single vector as 1D NumPy array
+                        is_single_query = True
+                        vector_list = query_vectors.tolist()
+                    elif query_vectors.ndim == 2:
+                        # Batch of vectors as 2D NumPy array
+                        vector_list = query_vectors.tolist()
+                    else:
+                        raise ValueError("Expected 1D or 2D NumPy array for `query_vectors`.")
+                elif isinstance(query_vectors, list):
+                    if not query_vectors:
+                        raise ValueError("Empty list provided for `query_vectors`.")
+                    if isinstance(query_vectors[0], (list, np.ndarray)):
+                        # Batch of vectors as list of lists
+                        vector_list = [list(map(float, v)) if isinstance(v, list) else v.tolist() for v in query_vectors]
+                    else:
+                        # Single vector as flat list
+                        is_single_query = True
+                        vector_list = list(map(float, query_vectors))
+                else:
+                    raise ValueError("Invalid type for `query_vectors`")
+
+            if is_single_query or query_contents is not None:
+                # Use QueryRequest for single vector or content-based query
                 query_request = QueryRequest(
                     index_key=self._key_to_hex(),
                     index_name=self._index_name,
@@ -444,16 +461,8 @@ def query(
                     filters=filters,
                     include=include,
                 )
-
-            elif query_vectors is not None:
-                if isinstance(query_vectors, list):
-                    query_vectors = np.array(query_vectors, dtype=np.float32)
-                if isinstance(query_vectors, np.ndarray):
-                    if query_vectors.ndim != 2:
-                        raise ValueError("Expected 2D NumPy array or list of lists for `query_vectors`.")
-                    vector_list = query_vectors.tolist()
-                else:
-                    raise ValueError("Invalid type for `query_vectors`")
+            else:
+                # Use BatchQueryRequest for multiple vectors
                 query_request = BatchQueryRequest(
                     index_key=self._key_to_hex(),
                     index_name=self._index_name,
@@ -466,9 +475,6 @@ def query(
                     include=include,
                 )
 
-            elif query_contents is None:
-                raise ValueError("You must provide `query_vector`, `query_vectors`, or `query_contents`.")
-
             request = Request(query_request)
 
             # Execute query via REST

cyborgdb/integrations/langchain.py

@@ -458,7 +458,7 @@ def _execute_query(
                 # Text query - generate embedding
                 embedding = self.get_embeddings(query)
                 results = self.index.query(
-                    query_vector=embedding,
+                    query_vectors=embedding,
                     top_k=k,
                     n_probes=n_probes,
                     filters=filter,
@@ -467,7 +467,7 @@ def _execute_query(
             else:
                 # Vector query
                 results = self.index.query(
-                    query_vector=query,
+                    query_vectors=query,
                     top_k=k,
                     n_probes=n_probes,
                     filters=filter,

tests/test_client.py

@@ -54,7 +54,7 @@ def test_upsert_and_query(self):
 
         # Query a vector
         query_vector = np.random.rand(dimension).astype(np.float32)
-        results = self.index.query(query_vector=query_vector, top_k=10)
+        results = self.index.query(query_vectors=query_vector, top_k=10)
 
         # Check results
         self.assertEqual(len(results[0]), 10)