Spaces:
Build error
Build error
File size: 3,393 Bytes
3932407 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
import pytest
from .helpers.collection_setup import drop_collection
from .helpers.helpers import request_with_validation
def basic_collection_setup(
collection_name='test_collection',
on_disk_vectors=False,
):
response = request_with_validation(
api='/collections/{collection_name}',
method="DELETE",
path_params={'collection_name': collection_name},
)
assert response.ok
response = request_with_validation(
api='/collections/{collection_name}',
method="PUT",
path_params={'collection_name': collection_name},
body={
"vectors": {
"size": 2,
"distance": "Euclid",
"on_disk": on_disk_vectors,
}
}
)
assert response.ok
response = request_with_validation(
api='/collections/{collection_name}',
method="GET",
path_params={'collection_name': collection_name},
)
assert response.ok
response = request_with_validation(
api='/collections/{collection_name}/points',
method="PUT",
path_params={'collection_name': collection_name},
query_params={'wait': 'true'},
body={
"points": [
{
"id": 1,
"vector": [0.0, 0.0],
"payload": {"city": "Berlin"}
},
{
"id": 2,
"vector": [0.0, 1.0],
"payload": {"city": ["Berlin", "London"]}
},
{
"id": 3,
"vector": [-1., -1.],
"payload": {"city": ["Berlin", "Moscow"]}
},
]
}
)
assert response.ok
@pytest.fixture(autouse=True, scope="module")
def setup(on_disk_vectors, collection_name):
basic_collection_setup(collection_name=collection_name, on_disk_vectors=on_disk_vectors)
yield
drop_collection(collection_name=collection_name)
def test_search_with_threshold(collection_name):
response = request_with_validation(
api='/collections/{collection_name}/points/search',
method="POST",
path_params={'collection_name': collection_name},
body={
"vector": [1., 1.],
"limit": 3
}
)
assert response.ok
assert len(response.json()['result']) == 3
assert response.json()['result'][0]['id'] == 2
assert response.json()['result'][1]['id'] == 1
assert response.json()['result'][2]['id'] == 3
assert response.json()['result'][0]['score'] - 1.0 < 0.0001
assert response.json()['result'][1]['score'] - 1.414214 < 0.0001
assert response.json()['result'][2]['score'] - 2.828427 < 0.0001
response = request_with_validation(
api='/collections/{collection_name}/points/search',
method="POST",
path_params={'collection_name': collection_name},
body={
"vector": [1., 1.],
"limit": 3,
"score_threshold": 1.5
}
)
assert response.ok
assert len(response.json()['result']) == 2
assert response.json()['result'][0]['id'] == 2
assert response.json()['result'][1]['id'] == 1
assert response.json()['result'][0]['score'] - 1.0 < 0.0001
assert response.json()['result'][1]['score'] - 1.414214 < 0.0001
|