mirror of
https://github.com/qdrant/qdrant.git
synced 2026-07-23 11:11:00 -05:00
328 lines
10 KiB
Python
328 lines
10 KiB
Python
import pytest
|
|
|
|
from .helpers.collection_setup import basic_collection_setup, drop_collection
|
|
from .helpers.helpers import request_with_validation
|
|
|
|
@pytest.fixture(scope='module', autouse=True)
|
|
def collection_name2(collection_name):
|
|
return f"{collection_name}_2"
|
|
|
|
|
|
@pytest.fixture(autouse=True, scope="module")
|
|
def setup(on_disk_vectors, collection_name, collection_name2):
|
|
basic_collection_setup(collection_name=collection_name, on_disk_vectors=on_disk_vectors)
|
|
yield
|
|
drop_collection(collection_name=collection_name)
|
|
# delete potential auxiliary collection as well
|
|
drop_collection(collection_name=collection_name2)
|
|
|
|
|
|
def test_recommend_with_wrong_vector_size(on_disk_vectors, collection_name, collection_name2):
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}',
|
|
method="DELETE",
|
|
path_params={'collection_name': collection_name2},
|
|
)
|
|
assert response.ok
|
|
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}',
|
|
method="PUT",
|
|
path_params={'collection_name': collection_name2},
|
|
body={
|
|
"vectors": {
|
|
"other": {
|
|
"size": 5,
|
|
"distance": "Dot",
|
|
"on_disk": on_disk_vectors,
|
|
}
|
|
}
|
|
}
|
|
)
|
|
assert response.ok
|
|
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}/points',
|
|
method="PUT",
|
|
path_params={'collection_name': collection_name2},
|
|
query_params={'wait': 'true'},
|
|
body={
|
|
"points": [
|
|
{
|
|
"id": 1,
|
|
"vector": {"other": [1.0, 0.0, 0.0, 0.0, 0.0]},
|
|
},
|
|
{
|
|
"id": "00000000-0000-0000-0000-000000000000",
|
|
"vector": {"other": [0.0, 1.0, 0.0, 0.0, 0.0]},
|
|
}
|
|
]
|
|
}
|
|
)
|
|
assert response.ok, response.text
|
|
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}/points/recommend',
|
|
method="POST",
|
|
path_params={'collection_name': collection_name},
|
|
body={
|
|
"limit": 3,
|
|
"positive": [1],
|
|
"with_vector": False,
|
|
"with_payload": True,
|
|
"lookup_from": {
|
|
"collection": collection_name2,
|
|
"vector": "other"
|
|
}
|
|
}
|
|
)
|
|
assert response.status_code == 400, response.text
|
|
|
|
|
|
def test_recommend_from_another_collection(on_disk_vectors, collection_name, collection_name2):
|
|
# Create another collection with the same vector size.
|
|
# Use vectors from the second collection to search in the first collection.
|
|
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}',
|
|
method="DELETE",
|
|
path_params={'collection_name': collection_name2},
|
|
)
|
|
assert response.ok
|
|
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}',
|
|
method="PUT",
|
|
path_params={'collection_name': collection_name2},
|
|
body={
|
|
"vectors": {
|
|
"other": {
|
|
"size": 4,
|
|
"distance": "Dot",
|
|
"on_disk": on_disk_vectors,
|
|
}
|
|
}
|
|
}
|
|
)
|
|
assert response.ok
|
|
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}/points',
|
|
method="PUT",
|
|
path_params={'collection_name': collection_name2},
|
|
query_params={'wait': 'true'},
|
|
body={
|
|
"points": [
|
|
{
|
|
"id": 1,
|
|
"vector": {"other": [1.0, 0.0, 0.0, 0.0]},
|
|
},
|
|
{
|
|
"id": "00000000-0000-0000-0000-000000000000",
|
|
"vector": {"other": [0.0, 1.0, 0.0, 0.0]},
|
|
}
|
|
]
|
|
}
|
|
)
|
|
assert response.ok, response.text
|
|
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}/points/recommend',
|
|
method="POST",
|
|
path_params={'collection_name': collection_name},
|
|
body={
|
|
"limit": 3,
|
|
"positive": [1],
|
|
"with_vector": False,
|
|
"with_payload": True,
|
|
"lookup_from": {
|
|
"collection": collection_name2,
|
|
"vector": "other"
|
|
}
|
|
}
|
|
)
|
|
assert response.ok, response.text
|
|
assert len(response.json()['result']) == 3
|
|
# vector with the largest 1st element
|
|
assert response.json()['result'][0]['id'] == 8
|
|
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}/points/recommend/batch',
|
|
method="POST",
|
|
path_params={'collection_name': collection_name},
|
|
body={
|
|
"searches": [
|
|
{
|
|
"limit": 3,
|
|
"positive": [1],
|
|
"with_vector": False,
|
|
"with_payload": True,
|
|
"lookup_from": {
|
|
"collection": collection_name2,
|
|
"vector": "other"
|
|
}
|
|
},
|
|
{
|
|
"limit": 3,
|
|
"positive": ["00000000-0000-0000-0000-000000000000"],
|
|
"with_vector": False,
|
|
"with_payload": True,
|
|
"lookup_from": {
|
|
"collection": collection_name2,
|
|
"vector": "other"
|
|
}
|
|
},
|
|
{
|
|
"limit": 3,
|
|
"positive": [1],
|
|
"with_vector": False,
|
|
"with_payload": True,
|
|
}
|
|
]
|
|
}
|
|
)
|
|
assert response.ok, response.text
|
|
assert len(response.json()['result']) == 3
|
|
assert len(response.json()['result'][0]) == 3
|
|
assert len(response.json()['result'][1]) == 3
|
|
assert len(response.json()['result'][2]) == 3
|
|
# vector with the largest 1st element
|
|
assert response.json()['result'][0][0]['id'] == 8
|
|
# vector with the largest 2nd element
|
|
assert response.json()['result'][1][0]['id'] == 7
|
|
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}/points/recommend',
|
|
method="POST",
|
|
path_params={'collection_name': collection_name},
|
|
body={
|
|
"limit": 3,
|
|
"positive": [1],
|
|
"lookup_from": {
|
|
"collection": "unknown_collection",
|
|
"vector": "other"
|
|
}
|
|
}
|
|
)
|
|
assert response.status_code == 404
|
|
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}/points/recommend',
|
|
method="POST",
|
|
path_params={'collection_name': collection_name},
|
|
body={
|
|
"limit": 3,
|
|
"positive": [1],
|
|
"lookup_from": {
|
|
"collection": collection_name2,
|
|
"vector": "unknown_vector"
|
|
}
|
|
}
|
|
)
|
|
assert response.status_code == 400, response.text
|
|
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}/points/recommend',
|
|
method="POST",
|
|
path_params={'collection_name': collection_name},
|
|
body={
|
|
"limit": 3,
|
|
"positive": [2],
|
|
"lookup_from": {
|
|
"collection": collection_name2,
|
|
"vector": "unknown_vector"
|
|
}
|
|
}
|
|
)
|
|
assert response.status_code == 404, response.text
|
|
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}',
|
|
method="DELETE",
|
|
path_params={'collection_name': collection_name2},
|
|
)
|
|
assert response.ok
|
|
|
|
|
|
def test_recommend_lookup(collection_name, collection_name2):
|
|
# delete lookup collection if exists
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}',
|
|
method="DELETE",
|
|
path_params={'collection_name': collection_name2},
|
|
)
|
|
assert response.ok, response.text
|
|
|
|
# re-create lookup collection
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}',
|
|
method="PUT",
|
|
path_params={'collection_name': collection_name2},
|
|
body={
|
|
"vectors": {
|
|
"other": {
|
|
"size": 4,
|
|
"distance": "Dot",
|
|
}
|
|
}
|
|
}
|
|
)
|
|
assert response.ok, response.text
|
|
|
|
# insert vectors to lookup collection
|
|
response = request_with_validation(
|
|
api='/collections/{collection_name}/points',
|
|
method="PUT",
|
|
path_params={'collection_name': collection_name2},
|
|
query_params={'wait': 'true'},
|
|
body={
|
|
"points": [
|
|
{
|
|
"id": 1,
|
|
"vector": {"other": [1.0, 0.0, 0.0, 0.0]},
|
|
},
|
|
{
|
|
"id": 2,
|
|
"vector": {"other": [0.0, 0.0, 0.0, 2.0]},
|
|
},
|
|
]
|
|
}
|
|
)
|
|
assert response.ok, response.text
|
|
|
|
# check recommend by id + lookup_from
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/recommend",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"positive": [1],
|
|
"negative": [2],
|
|
"limit": 10,
|
|
"lookup_from": {
|
|
"collection": collection_name2,
|
|
"vector": "other"
|
|
}
|
|
},
|
|
)
|
|
assert response.ok, response.text
|
|
recommend_result_by_id = response.json()["result"]
|
|
|
|
# check recommend by vector + lookup_from
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/recommend",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"positive": [[1.0, 0.0, 0.0, 0.0]],
|
|
"negative": [[0.0, 0.0, 0.0, 2.0]],
|
|
"limit": 10,
|
|
},
|
|
)
|
|
assert response.ok, response.text
|
|
recommend_result_by_vector = response.json()["result"]
|
|
|
|
# results should be equivalent
|
|
assert recommend_result_by_id == recommend_result_by_vector, \
|
|
f"recommend_result_by_id: {recommend_result_by_id}, recommend_result_by_vector: {recommend_result_by_vector}"
|