Support for hybrid search in Azure AI vector store (#2408)
Co-authored-by: Deshraj Yadav <deshrajdry@gmail.com>
This commit is contained in:
@@ -196,8 +196,8 @@ class TestElasticsearchDB(unittest.TestCase):
|
||||
self.client_mock.search.return_value = mock_response
|
||||
|
||||
# Perform search
|
||||
query_vector = [0.1] * 1536
|
||||
results = self.es_db.search(query=query_vector, limit=5)
|
||||
vectors = [[0.1] * 1536]
|
||||
results = self.es_db.search(query="", vectors=vectors, limit=5)
|
||||
|
||||
# Verify search call
|
||||
self.client_mock.search.assert_called_once()
|
||||
@@ -210,7 +210,7 @@ class TestElasticsearchDB(unittest.TestCase):
|
||||
# Verify KNN query structure
|
||||
self.assertIn("knn", body)
|
||||
self.assertEqual(body["knn"]["field"], "vector")
|
||||
self.assertEqual(body["knn"]["query_vector"], query_vector)
|
||||
self.assertEqual(body["knn"]["query_vector"], vectors)
|
||||
self.assertEqual(body["knn"]["k"], 5)
|
||||
self.assertEqual(body["knn"]["num_candidates"], 10)
|
||||
|
||||
@@ -226,13 +226,13 @@ class TestElasticsearchDB(unittest.TestCase):
|
||||
self.es_db.custom_search_query.return_value = {"custom_key": "custom_value"}
|
||||
|
||||
# Perform search
|
||||
query_vector = [0.1] * 1536
|
||||
vectors = [[0.1] * 1536]
|
||||
limit = 5
|
||||
filters = {"key1": "value1"}
|
||||
self.es_db.search(query=query_vector, limit=limit, filters=filters)
|
||||
self.es_db.search(query="", vectors=vectors, limit=limit, filters=filters)
|
||||
|
||||
# Verify custom search query function was called
|
||||
self.es_db.custom_search_query.assert_called_once_with(query_vector, limit, filters)
|
||||
self.es_db.custom_search_query.assert_called_once_with(vectors, limit, filters)
|
||||
|
||||
# Verify custom search query was used
|
||||
self.client_mock.search.assert_called_once_with(index=self.es_db.collection_name, body={"custom_key": "custom_value"})
|
||||
|
||||
Reference in New Issue
Block a user