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

[Feature] Return score when doing search in vectorDB (#1060)

Browse Source 
Co-authored-by: Deven Patel <deven298@yahoo.com>
This commit is contained in:
Deven Patel
2023-12-29 15:56:12 +05:30
committed by GitHub
parent 19d80914df
commit c0aafd38c9
12 changed files with 72 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
tests/vectordb/test_chroma_db.py
View File

@@ -342,8 +342,22 @@ def test_chroma_db_collection_query(app_with_settings):
input_query=[0, 0, 0], where={}, n_results=2, skip_embedding=True, citations=True
)
expected_value_with_citations = [
("document", {"url": "url_1", "doc_id": "doc_id_1"}),
("document2", {"url": "url_2", "doc_id": "doc_id_2"}),
(
"document",
{
"url": "url_1",
"doc_id": "doc_id_1",
"score": 0.0,
},
),
(
"document2",
{
"url": "url_2",
"doc_id": "doc_id_2",
"score": 1.0,
},
),
]
assert data_with_citations == expected_value_with_citations

4
tests/vectordb/test_elasticsearch_db.py
View File

@@ -66,8 +66,8 @@ class TestEsDB(unittest.TestCase):
results_with_citations = self.db.query(query, n_results=2, where={}, skip_embedding=False, citations=True)
expected_results_with_citations = [
("This is a document.", {"url": "url_1", "doc_id": "doc_id_1"}),
("This is another document.", {"url": "url_2", "doc_id": "doc_id_2"}),
("This is a document.", {"url": "url_1", "doc_id": "doc_id_1", "score": 0.9}),
("This is another document.", {"url": "url_2", "doc_id": "doc_id_2", "score": 0.8}),
]
self.assertEqual(results_with_citations, expected_results_with_citations)

30
tests/vectordb/test_zilliz_db.py
View File

@@ -123,7 +123,14 @@ class TestZillizDBCollection:
# Mock the MilvusClient search method
with patch.object(zilliz_db.client, "search") as mock_search:
# Mock the search result
mock_search.return_value = [[{"entity": {"text": "result_doc", "url": "url_1", "doc_id": "doc_id_1"}}]]
mock_search.return_value = [
[
{
"distance": 0.5,
"entity": {"text": "result_doc", "url": "url_1", "doc_id": "doc_id_1", "embeddings": [1, 2, 3]},
}
]
]
# Call the query method with skip_embedding=True
query_result = zilliz_db.query(input_query=["query_text"], n_results=1, where={}, skip_embedding=True)
@@ -133,7 +140,7 @@ class TestZillizDBCollection:
collection_name=mock_config.collection_name,
data=["query_text"],
limit=1,
output_fields=["text", "url", "doc_id"],
output_fields=["*"],
)
# Assert that the query result matches the expected result
@@ -147,11 +154,11 @@ class TestZillizDBCollection:
collection_name=mock_config.collection_name,
data=["query_text"],
limit=1,
output_fields=["text", "url", "doc_id"],
output_fields=["*"],
)
assert query_result_with_citations == [
("result_doc", {"text": "result_doc", "url": "url_1", "doc_id": "doc_id_1"})
("result_doc", {"text": "result_doc", "url": "url_1", "doc_id": "doc_id_1", "score": 0.5})
]
@patch("embedchain.vectordb.zilliz.MilvusClient", autospec=True)
@@ -177,7 +184,14 @@ class TestZillizDBCollection:
mock_embedder.embedding_fn.return_value = ["query_vector"]
# Mock the search result
mock_search.return_value = [[{"entity": {"text": "result_doc", "url": "url_1", "doc_id": "doc_id_1"}}]]
mock_search.return_value = [
[
{
"distance": 0.0,
"entity": {"text": "result_doc", "url": "url_1", "doc_id": "doc_id_1", "embeddings": [1, 2, 3]},
}
]
]
# Call the query method with skip_embedding=False
query_result = zilliz_db.query(input_query=["query_text"], n_results=1, where={}, skip_embedding=False)
@@ -187,7 +201,7 @@ class TestZillizDBCollection:
collection_name=mock_config.collection_name,
data=["query_vector"],
limit=1,
output_fields=["text", "url", "doc_id"],
output_fields=["*"],
)
# Assert that the query result matches the expected result
@@ -201,9 +215,9 @@ class TestZillizDBCollection:
collection_name=mock_config.collection_name,
data=["query_vector"],
limit=1,
output_fields=["text", "url", "doc_id"],
output_fields=["*"],
)
assert query_result_with_citations == [
("result_doc", {"text": "result_doc", "url": "url_1", "doc_id": "doc_id_1"})
("result_doc", {"text": "result_doc", "url": "url_1", "doc_id": "doc_id_1", "score": 0.0})
]