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

[Improvement] return all the metadata when citations flag is True (#1059)

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

5
tests/vectordb/test_chroma_db.py
View File

@@ -341,7 +341,10 @@ def test_chroma_db_collection_query(app_with_settings):
data_with_citations = app_with_settings.db.query(
input_query=[0, 0, 0], where={}, n_results=2, skip_embedding=True, citations=True
)
expected_value_with_citations = [("document", "url_1", "doc_id_1"), ("document2", "url_2", "doc_id_2")]
expected_value_with_citations = [
("document", {"url": "url_1", "doc_id": "doc_id_1"}),
("document2", {"url": "url_2", "doc_id": "doc_id_2"}),
]
assert data_with_citations == expected_value_with_citations
app_with_settings.db.reset()

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_1", "doc_id_1"),
("This is another document.", "url_2", "doc_id_2"),
("This is a document.", {"url": "url_1", "doc_id": "doc_id_1"}),
("This is another document.", {"url": "url_2", "doc_id": "doc_id_2"}),
]
self.assertEqual(results_with_citations, expected_results_with_citations)

8
tests/vectordb/test_zilliz_db.py
View File

@@ -150,7 +150,9 @@ class TestZillizDBCollection:
output_fields=["text", "url", "doc_id"],
)
assert query_result_with_citations == [("result_doc", "url_1", "doc_id_1")]
assert query_result_with_citations == [
("result_doc", {"text": "result_doc", "url": "url_1", "doc_id": "doc_id_1"})
]
@patch("embedchain.vectordb.zilliz.MilvusClient", autospec=True)
@patch("embedchain.vectordb.zilliz.connections", autospec=True)
@@ -202,4 +204,6 @@ class TestZillizDBCollection:
output_fields=["text", "url", "doc_id"],
)
assert query_result_with_citations == [("result_doc", "url_1", "doc_id_1")]
assert query_result_with_citations == [
("result_doc", {"text": "result_doc", "url": "url_1", "doc_id": "doc_id_1"})
]