change qdrant_store.py

This commit is contained in:
hezz 2023-08-13 21:19:43 +08:00
parent dc38226b59
commit bb7ea6c398

View file

@ -43,151 +43,17 @@ def test_milvus_store():
assert qdrant_store.has_collection("Book") is True
qdrant_store.add("Book", points)
results = qdrant_store.search("Book", query=[1.0, 1.0])
assert results == [
{
"id": 2,
"version": 0,
"score": 0.999106722578389,
"payload": {"color": "red", "rand_number": 2},
"vector": None,
},
{
"id": 7,
"version": 0,
"score": 0.9961650411397226,
"payload": {"color": "red", "rand_number": 7},
"vector": None,
},
{
"id": 1,
"version": 0,
"score": 0.9946351526856256,
"payload": {"color": "red", "rand_number": 1},
"vector": None,
},
{
"id": 5,
"version": 0,
"score": 0.9297466022881021,
"payload": {"color": "red", "rand_number": 5},
"vector": None,
},
{
"id": 8,
"version": 0,
"score": 0.9100373450784073,
"payload": {"color": "red", "rand_number": 8},
"vector": None,
},
{
"id": 6,
"version": 0,
"score": 0.7944306996390111,
"payload": {"color": "red", "rand_number": 6},
"vector": None,
},
{
"id": 3,
"version": 0,
"score": 0.7723528053480722,
"payload": {"color": "red", "rand_number": 3},
"vector": None,
},
{
"id": 4,
"version": 0,
"score": 0.755163629383033,
"payload": {"color": "red", "rand_number": 4},
"vector": None,
},
{
"id": 0,
"version": 0,
"score": 0.73420337995255,
"payload": {"color": "red", "rand_number": 0},
"vector": None,
},
{
"id": 9,
"version": 0,
"score": 0.7127610621127889,
"payload": {"color": "red", "rand_number": 9},
"vector": None,
},
]
assert results[0]["id"] == 2
assert results[0]["score"] == 0.999106722578389
assert results[1]["score"] == 7
assert results[1]["score"] == 0.9961650411397226
results = qdrant_store.search("Book", query=[1.0, 1.0], return_vector=True)
assert results == [
{
"id": 2,
"version": 0,
"score": 0.999106722578389,
"payload": {"color": "red", "rand_number": 2},
"vector": [0.7363563179969788, 0.6765939593315125],
},
{
"id": 7,
"version": 0,
"score": 0.9961650411397226,
"payload": {"color": "red", "rand_number": 7},
"vector": [0.7662628889083862, 0.6425272226333618],
},
{
"id": 1,
"version": 0,
"score": 0.9946351526856256,
"payload": {"color": "red", "rand_number": 1},
"vector": [0.7764601111412048, 0.6301664113998413],
},
{
"id": 5,
"version": 0,
"score": 0.9297466022881021,
"payload": {"color": "red", "rand_number": 5},
"vector": [0.39707326889038086, 0.9177868962287903],
},
{
"id": 8,
"version": 0,
"score": 0.9100373450784073,
"payload": {"color": "red", "rand_number": 8},
"vector": [0.35037919878959656, 0.9366079568862915],
},
{
"id": 6,
"version": 0,
"score": 0.7944306996390111,
"payload": {"color": "red", "rand_number": 6},
"vector": [0.13228265941143036, 0.991212010383606],
},
{
"id": 3,
"version": 0,
"score": 0.7723528053480722,
"payload": {"color": "red", "rand_number": 3},
"vector": [0.9952857494354248, 0.0969860628247261],
},
{
"id": 4,
"version": 0,
"score": 0.755163629383033,
"payload": {"color": "red", "rand_number": 4},
"vector": [0.9975154995918274, 0.07044714689254761],
},
{
"id": 0,
"version": 0,
"score": 0.73420337995255,
"payload": {"color": "red", "rand_number": 0},
"vector": [0.9992359280586243, 0.03908444941043854],
},
{
"id": 9,
"version": 0,
"score": 0.7127610621127889,
"payload": {"color": "red", "rand_number": 9},
"vector": [0.9999677538871765, 0.00802854634821415],
},
]
assert results[0]["id"] == 2
assert results[0]["score"] == 0.999106722578389
assert results[0]["vector"] == [0.7363563179969788, 0.6765939593315125]
assert results[1]["score"] == 7
assert results[1]["score"] == 0.9961650411397226
assert results[1]["vector"] == [0.7662628889083862, 0.6425272226333618]
results = qdrant_store.search(
"Book",
query=[1.0, 1.0],
@ -195,22 +61,10 @@ def test_milvus_store():
must=[FieldCondition(key="rand_number", range=Range(gte=8))]
),
)
assert results == [
{
"id": 8,
"version": 0,
"score": 0.9100373450784073,
"payload": {"color": "red", "rand_number": 8},
"vector": None,
},
{
"id": 9,
"version": 0,
"score": 0.7127610621127889,
"payload": {"color": "red", "rand_number": 9},
"vector": None,
},
]
assert results[0]["id"] == 8
assert results[0]["score"] == 0.9100373450784073
assert results[1]["id"] == 9
assert results[1]["score"] == 0.7127610621127889
results = qdrant_store.search(
"Book",
query=[1.0, 1.0],
@ -219,19 +73,5 @@ def test_milvus_store():
),
return_vector=True,
)
assert results == [
{
"id": 8,
"version": 0,
"score": 0.9100373450784073,
"payload": {"color": "red", "rand_number": 8},
"vector": [0.35037919878959656, 0.9366079568862915],
},
{
"id": 9,
"version": 0,
"score": 0.7127610621127889,
"payload": {"color": "red", "rand_number": 9},
"vector": [0.9999677538871765, 0.00802854634821415],
},
]
assert results[0]["vector"] == [0.35037919878959656, 0.9366079568862915]
assert results[1]["vector"] == [0.9999677538871765, 0.00802854634821415]