lazy dependency for milvus

Signed-off-by: ChengZi <chen.zhang@zilliz.com>
This commit is contained in:
ChengZi 2024-09-25 11:47:28 +08:00
parent e5f037f86d
commit 4d92fdcec9
2 changed files with 28 additions and 53 deletions

View file

@ -1,4 +1,7 @@
import random
import pytest
from metagpt.document_store.milvus_store import MilvusConnection, MilvusStore
seed_value = 42
@ -19,6 +22,7 @@ def assert_almost_equal(actual, expected):
assert abs(actual - expected) <= delta, f"{actual} is not within {delta} of {expected}"
@pytest.mark.skip() # Skip because the pymilvus dependency is not installed by default
def test_milvus_store():
milvus_connection = MilvusConnection(uri="./milvus_local.db")
milvus_store = MilvusStore(milvus_connection)
@ -33,11 +37,7 @@ def test_milvus_store():
first_result = search_results[0]
assert first_result["id"] == "doc_0"
search_results_with_filter = milvus_store.search(
collection_name,
query=[1.0] * 8,
filter={"rand_number": 1}
)
search_results_with_filter = milvus_store.search(collection_name, query=[1.0] * 8, filter={"rand_number": 1})
assert len(search_results_with_filter) > 0
assert search_results_with_filter[0]["id"] == "doc_1"