Allow passing query_filter param to Vector Search

onyxgenai/embed.py

@@ -45,12 +45,14 @@ def _onyx_embed(
             print("Failed to get embedding:", response.status_code, response.text)
             return None
 
-    def _onyx_vector_search(self, query: str, collection_name: str, limit: int):
+    def _onyx_vector_search(
+        self, query: str, collection_name: str, limit: int, query_filter=None
+    ):
         url = f"{self.svc_url}/vector-store/search"
         payload = {
             "query_vector": query,
             "collection_name": collection_name,
-            "kwargs": {"limit": limit},
+            "kwargs": {"limit": limit, "query_filter": query_filter},
         }
 
         response = requests.post(url, json=payload)
@@ -183,17 +185,18 @@ def embed_images(
 
         return results
 
-    def vector_search(self, query, collection_name, limit=3):
+    def vector_search(self, query, collection_name, limit=3, query_filter=None):
         """
         Search for vectors in the collection
         Args:
             query (str): The query vector
             collection_name (str): The name of the collection
             limit (int): The number of results to return (default 3)
+            query_filter (dict): The query filter (default None)
         Returns:
             list: The vector search results
         """
-        return self._onyx_vector_search(query, collection_name, limit)
+        return self._onyx_vector_search(query, collection_name, limit, query_filter)
 
     def get_collections(self):
         """

onyxgenai/extract.py

@@ -129,7 +129,6 @@ def tokenize_sentences(text):
 
 
 if __name__ == "__main__":
-
     args = sys.argv
     fp = args[1]
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "onyxgenai"
-version = "1.0.4"
+version = "1.0.5"
 description = "Python SDK for interacting with Onyx Generative AI Services"
 readme = "README.md"
 license = { file = "LICENSE" }

tests/test_embed.py

@@ -56,7 +56,7 @@ def test_onyx_vector_search(mock_post):
 
 
 @patch("requests.post")
-def test_onyx_vector_search_with_limit(mock_post):
+def test_onyx_vector_search_with_kwargs(mock_post):
     svc_url = "http://localhost:8000"
     client = EmbeddingClient(svc_url)
     mock_response = {"results": ["result1", "result2", "result3"]}
@@ -66,7 +66,8 @@ def test_onyx_vector_search_with_limit(mock_post):
     query = "sample query"
     collection_name = "test_collection"
     limit = 10
-    result = client.vector_search(query, collection_name, limit)
+    query_filter = {"field": "value"}
+    result = client.vector_search(query, collection_name, limit, query_filter)
 
     assert result == mock_response["results"]