klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Support for hybrid search in Azure AI vector store (#2408)

Browse Source 
Co-authored-by: Deshraj Yadav <deshrajdry@gmail.com>
This commit is contained in:
Dev Khant
2025-03-20 22:57:00 +05:30
committed by GitHub
parent 8b9a8e5825
commit 8e6a08aa83
24 changed files with 275 additions and 294 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/vector_stores/test_chroma.py
View File

@@ -35,12 +35,11 @@ def test_search_vectors(chromadb_instance, mock_chromadb_client):
}
chromadb_instance.collection.query.return_value = mock_result
query = [[0.1, 0.2, 0.3]]
results = chromadb_instance.search(query=query, limit=2)
vectors = [[0.1, 0.2, 0.3]]
results = chromadb_instance.search(query="", vectors=vectors, limit=2)
chromadb_instance.collection.query.assert_called_once_with(query_embeddings=query, where=None, n_results=2)
chromadb_instance.collection.query.assert_called_once_with(query_embeddings=vectors, where=None, n_results=2)
print(results, type(results))
assert len(results) == 2
assert results[0].id == "id1"
assert results[0].score == 0.1

12
tests/vector_stores/test_elasticsearch.py
View File

@@ -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"})

6
tests/vector_stores/test_opensearch.py
View File

@@ -126,15 +126,15 @@ class TestOpenSearchDB(unittest.TestCase):
def test_search(self):
mock_response = {"hits": {"hits": [{"_id": "id1", "_score": 0.8, "_source": {"vector": [0.1] * 1536, "metadata": {"key1": "value1"}}}]}}
self.client_mock.search.return_value = mock_response
query_vector = [0.1] * 1536
results = self.os_db.search(query=query_vector, limit=5)
vectors = [[0.1] * 1536]
results = self.os_db.search(query="", vectors=vectors, limit=5)
self.client_mock.search.assert_called_once()
search_args = self.client_mock.search.call_args[1]
self.assertEqual(search_args["index"], "test_collection")
body = search_args["body"]
self.assertIn("knn", body["query"])
self.assertIn("vector", body["query"]["knn"])
self.assertEqual(body["query"]["knn"]["vector"]["vector"], query_vector)
self.assertEqual(body["query"]["knn"]["vector"]["vector"], vectors)
self.assertEqual(body["query"]["knn"]["vector"]["k"], 5)
self.assertEqual(len(results), 1)
self.assertEqual(results[0].id, "id1")

6
tests/vector_stores/test_qdrant.py
View File

@@ -50,15 +50,15 @@ class TestQdrant(unittest.TestCase):
self.assertEqual(points[0].payload, payloads[0])
def test_search(self):
query_vector = [0.1, 0.2]
vectors = [[0.1, 0.2]]
mock_point = MagicMock(id=str(uuid.uuid4()), score=0.95, payload={"key": "value"})
self.client_mock.query_points.return_value = MagicMock(points=[mock_point])
results = self.qdrant.search(query=query_vector, limit=1)
results = self.qdrant.search(query="", vectors=vectors, limit=1)
self.client_mock.query_points.assert_called_once_with(
collection_name="test_collection",
query=query_vector,
query=vectors,
query_filter=None,
limit=1,
)

6
tests/vector_stores/test_supabase.py
View File

@@ -77,12 +77,12 @@ def test_search_vectors(supabase_instance, mock_collection):
]
mock_collection.query.return_value = mock_results
query = [0.1, 0.2, 0.3]
vectors = [[0.1, 0.2, 0.3]]
filters = {"category": "test"}
results = supabase_instance.search(query=query, limit=2, filters=filters)
results = supabase_instance.search(query="", vectors=vectors, limit=2, filters=filters)
mock_collection.query.assert_called_once_with(
data=query,
data=vectors,
limit=2,
filters={"category": {"$eq": "test"}},
include_metadata=True,

8
tests/vector_stores/test_vertex_ai_vector_search.py
View File

@@ -73,12 +73,12 @@ def test_insert_vectors(vector_store, mock_vertex_ai):
def test_search_vectors(vector_store, mock_vertex_ai):
"""Test searching vectors with filters"""
query = [0.1, 0.2, 0.3]
vectors = [[0.1, 0.2, 0.3]]
filters = {"user_id": "test_user"}
mock_datapoint = Mock()
mock_datapoint.datapoint_id = "test-id"
mock_datapoint.feature_vector = query
mock_datapoint.feature_vector = vectors
mock_restrict = Mock()
mock_restrict.namespace = "user_id"
@@ -96,11 +96,11 @@ def test_search_vectors(vector_store, mock_vertex_ai):
mock_vertex_ai['endpoint'].find_neighbors.return_value = [[mock_neighbor]]
results = vector_store.search(query=query, filters=filters, limit=1)
results = vector_store.search(query="", vectors=vectors, filters=filters, limit=1)
mock_vertex_ai['endpoint'].find_neighbors.assert_called_once_with(
deployed_index_id=vector_store.deployment_index_id,
queries=[query],
queries=[vectors],
num_neighbors=1,
filter=[Namespace("user_id", ["test_user"], [])],
return_full_datapoint=True

4
tests/vector_stores/test_weaviate.py
View File

@@ -147,8 +147,8 @@ class TestWeaviateDB(unittest.TestCase):
self.client_mock.collections.get.return_value.query.hybrid = mock_hybrid
mock_hybrid.return_value = mock_response
query_vector = [0.1] * 1536
results = self.weaviate_db.search(query=query_vector, limit=5)
vectors = [[0.1] * 1536]
results = self.weaviate_db.search(query="", vectors=vectors, limit=5)
mock_hybrid.assert_called_once()