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

[Feature] Update db.query to return source of context (#831)

Browse Source
This commit is contained in:
Deven Patel
2023-10-25 22:20:32 -07:00
committed by GitHub
parent a27eeb3255
commit d77e8da3f3
13 changed files with 195 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/vectordb/test_zilliz_db.py
View File

@@ -123,7 +123,7 @@ 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"}}]]
mock_search.return_value = [[{"entity": {"text": "result_doc", "url": "url_1", "doc_id": "doc_id_1"}}]]
# 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,11 +133,11 @@ class TestZillizDBCollection:
collection_name=mock_config.collection_name,
data=["query_text"],
limit=1,
output_fields=["text"],
output_fields=["text", "url", "doc_id"],
)
# Assert that the query result matches the expected result
assert query_result == ["result_doc"]
assert query_result == [("result_doc", "url_1", "doc_id_1")]
@patch("embedchain.vectordb.zilliz.MilvusClient", autospec=True)
@patch("embedchain.vectordb.zilliz.connections", autospec=True)
@@ -162,7 +162,7 @@ class TestZillizDBCollection:
mock_embedder.embedding_fn.return_value = ["query_vector"]
# Mock the search result
mock_search.return_value = [[{"entity": {"text": "result_doc"}}]]
mock_search.return_value = [[{"entity": {"text": "result_doc", "url": "url_1", "doc_id": "doc_id_1"}}]]
# Call the query method with skip_embedding=False
query_result = zilliz_db.query(input_query=["query_text"], n_results=1, where={}, skip_embedding=False)
@@ -172,8 +172,8 @@ class TestZillizDBCollection:
collection_name=mock_config.collection_name,
data=["query_vector"],
limit=1,
output_fields=["text"],
output_fields=["text", "url", "doc_id"],
)
# Assert that the query result matches the expected result
assert query_result == ["result_doc"]
assert query_result == [("result_doc", "url_1", "doc_id_1")]