mirror of
https://github.com/qdrant/qdrant.git
synced 2026-07-23 11:11:00 -05:00
6f23975e3e
* `QDRANT_HOST` variable is used by all openapi tests, removed its copies from
`test_multi_vector_uint8.py`, `test_multi_vector_unnamed.py`, `test_query.py`,
`test_shard_snapshot.py`, and `test_snapshot.py`
* Added optional `QDRANT_HOST_HEADERS` env variable to set custom HTTP headers
in order to reach Qdrant instance behind a reverse proxy. The content of the
variable is JSON payload, e.g. `{"host": "qdrant.local"}`
* Adapted `./tests/basic_*.sh` scripts to aware of `QDRANT_HOST_HEADERS` env variable
497 lines
15 KiB
Python
497 lines
15 KiB
Python
from math import isclose
|
|
import pytest
|
|
import requests
|
|
import os
|
|
|
|
from .helpers.collection_setup import basic_collection_setup, drop_collection
|
|
from .helpers.helpers import distribution_based_score_fusion, reciprocal_rank_fusion, request_with_validation, \
|
|
qdrant_host_headers
|
|
from .helpers.settings import QDRANT_HOST
|
|
|
|
collection_name = "test_query"
|
|
|
|
|
|
@pytest.fixture(autouse=True, scope="module")
|
|
def setup(on_disk_vectors):
|
|
basic_collection_setup(collection_name=collection_name, on_disk_vectors=on_disk_vectors)
|
|
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/index",
|
|
method="PUT",
|
|
path_params={"collection_name": collection_name},
|
|
body={"field_name": "price", "field_schema": "float"},
|
|
)
|
|
assert response.ok
|
|
yield
|
|
drop_collection(collection_name=collection_name)
|
|
|
|
|
|
def test_query_validation():
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"prefetch": [
|
|
{
|
|
"query": {
|
|
"recommend": {
|
|
"positive": [1]
|
|
},
|
|
}
|
|
}
|
|
]
|
|
},
|
|
)
|
|
assert not response.ok, response.text
|
|
assert response.json()["status"]["error"] == ("Bad request: A query is needed to merge the prefetches. Can't have prefetches without defining a query.")
|
|
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"score_threshold": 10,
|
|
},
|
|
)
|
|
assert not response.ok, response.text
|
|
assert response.json()["status"]["error"] == ("Bad request: A query is needed to use the score_threshold. Can't have score_threshold without defining a query.")
|
|
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"score_threshold": 10,
|
|
"query": {
|
|
"order_by": {
|
|
"key": "price",
|
|
}
|
|
}
|
|
},
|
|
)
|
|
assert not response.ok, response.text
|
|
assert response.json()["status"]["error"] == ("Bad request: Can't use score_threshold with an order_by query.")
|
|
|
|
|
|
# raw query to bypass local validation
|
|
response = requests.post(f"{QDRANT_HOST}/collections/{collection_name}/points/query",
|
|
headers=qdrant_host_headers(),
|
|
json={
|
|
"query": {
|
|
"recommend": {
|
|
"positive": [1]
|
|
},
|
|
},
|
|
"limit": 0,
|
|
},
|
|
)
|
|
assert not response.ok, response.text
|
|
assert response.json()["status"]["error"] == ("Validation error in JSON body: [internal.limit: value 0 invalid, must be 1 or larger]")
|
|
|
|
|
|
def root_and_rescored_query(query, limit=None, with_payload=None):
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"query": query,
|
|
"limit": limit,
|
|
"with_payload": with_payload,
|
|
},
|
|
)
|
|
assert response.ok
|
|
root_query_result = response.json()["result"]["points"]
|
|
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"prefetch": {
|
|
"limit": 1000,
|
|
},
|
|
"query": query,
|
|
"with_payload": with_payload,
|
|
},
|
|
)
|
|
assert response.ok
|
|
nested_query_result = response.json()["result"]["points"]
|
|
|
|
assert root_query_result == nested_query_result
|
|
return root_query_result
|
|
|
|
|
|
def test_basic_search():
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/search",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"vector": [0.1, 0.2, 0.3, 0.4],
|
|
"limit": 10,
|
|
},
|
|
)
|
|
assert response.ok
|
|
search_result = response.json()["result"]
|
|
|
|
default_query_result = root_and_rescored_query([0.1, 0.2, 0.3, 0.4])
|
|
|
|
nearest_query_result = root_and_rescored_query({"nearest": [0.1, 0.2, 0.3, 0.4]})
|
|
|
|
assert search_result == default_query_result
|
|
assert search_result == nearest_query_result
|
|
|
|
|
|
def test_basic_scroll():
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/scroll",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={},
|
|
)
|
|
assert response.ok
|
|
scroll_result = response.json()["result"]["points"]
|
|
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"with_payload": True,
|
|
},
|
|
)
|
|
assert response.ok
|
|
query_result = response.json()["result"]["points"]
|
|
|
|
for record, scored_point in zip(scroll_result, query_result):
|
|
assert record.get("id") == scored_point.get("id")
|
|
assert record.get("payload") == scored_point.get("payload")
|
|
|
|
def test_basic_scroll_offset():
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/scroll",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"offset": 2, # skip first record, start at id 2
|
|
},
|
|
)
|
|
assert response.ok
|
|
scroll_result = response.json()["result"]["points"]
|
|
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"with_payload": True,
|
|
"offset": 1, # skip one record
|
|
},
|
|
)
|
|
assert response.ok
|
|
query_result = response.json()["result"]["points"]
|
|
|
|
for record, scored_point in zip(scroll_result, query_result):
|
|
assert record.get("id") == scored_point.get("id")
|
|
assert record.get("payload") == scored_point.get("payload")
|
|
|
|
def test_basic_recommend_avg():
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/recommend",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"positive": [1, 2, 3, 4], # ids
|
|
"negative": [3], # ids
|
|
"limit": 10,
|
|
},
|
|
)
|
|
assert response.ok
|
|
recommend_result = response.json()["result"]
|
|
|
|
query_result = root_and_rescored_query(
|
|
{
|
|
"recommend": {"positive": [1, 2, 3, 4], "negative": [3]}, # ids
|
|
}
|
|
)
|
|
|
|
assert recommend_result == query_result
|
|
|
|
|
|
def test_basic_recommend_best_score():
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/recommend",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"positive": [1, 2, 3, 4], # ids
|
|
"negative": [3], # ids
|
|
"limit": 10,
|
|
"strategy": "best_score",
|
|
},
|
|
)
|
|
assert response.ok
|
|
recommend_result = response.json()["result"]
|
|
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"query": {
|
|
"recommend": {
|
|
"positive": [1, 2, 3, 4], # ids
|
|
"negative": [3], # ids
|
|
"strategy": "best_score",
|
|
},
|
|
}
|
|
},
|
|
)
|
|
assert response.ok
|
|
query_result = response.json()["result"]["points"]
|
|
|
|
assert recommend_result == query_result
|
|
|
|
|
|
def test_basic_discover():
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/discover",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"target": 2, # ids
|
|
"context": [{"positive": 3, "negative": 4}], # ids
|
|
"limit": 10,
|
|
},
|
|
)
|
|
assert response.ok
|
|
discover_result = response.json()["result"]
|
|
|
|
query_result = root_and_rescored_query(
|
|
{
|
|
"discover": {
|
|
"target": 2,
|
|
"context": [{"positive": 3, "negative": 4}],
|
|
}
|
|
},
|
|
)
|
|
|
|
assert discover_result == query_result
|
|
|
|
|
|
def test_basic_context():
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/discover",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"context": [{"positive": 2, "negative": 4}],
|
|
"limit": 100,
|
|
},
|
|
)
|
|
assert response.ok
|
|
context_result = response.json()["result"]
|
|
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"query": {
|
|
"context": [{"positive": 2, "negative": 4}],
|
|
},
|
|
"limit": 100,
|
|
},
|
|
)
|
|
assert response.ok
|
|
query_result = response.json()["result"]["points"]
|
|
|
|
assert set([p["id"] for p in context_result]) == set([p["id"] for p in query_result])
|
|
|
|
|
|
def test_basic_order_by():
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/scroll",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"order_by": "price",
|
|
},
|
|
)
|
|
assert response.ok, response.text
|
|
scroll_result = response.json()["result"]["points"]
|
|
|
|
query_result = root_and_rescored_query({"order_by": "price"}, with_payload=True)
|
|
|
|
for record, scored_point in zip(scroll_result, query_result):
|
|
assert record.get("id") == scored_point.get("id")
|
|
assert record.get("payload") == scored_point.get("payload")
|
|
|
|
|
|
def test_basic_random_query():
|
|
ids_lists = set()
|
|
for _ in range(4):
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"query": { "sample": "random" }
|
|
},
|
|
)
|
|
assert response.ok, response.text
|
|
|
|
points = response.json()["result"]["points"]
|
|
assert len(points) == 10
|
|
assert set(point["id"] for point in points) == set(range(1, 11))
|
|
|
|
ids = str([point["id"] for point in points])
|
|
ids_lists.add(ids)
|
|
|
|
# check the order of ids are different at least once
|
|
assert len(ids_lists) > 1
|
|
|
|
|
|
def test_basic_rrf():
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/search",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"vector": [0.1, 0.2, 0.3, 0.4],
|
|
"limit": 10,
|
|
},
|
|
)
|
|
assert response.ok
|
|
search_result_1 = response.json()["result"]
|
|
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/search",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"vector": [0.5, 0.6, 0.7, 0.8],
|
|
"limit": 10,
|
|
},
|
|
)
|
|
assert response.ok
|
|
search_result_2 = response.json()["result"]
|
|
|
|
rrf_expected = reciprocal_rank_fusion([search_result_1, search_result_2], limit=10)
|
|
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"prefetch": [
|
|
{ "query": [0.1, 0.2, 0.3, 0.4] },
|
|
{ "query": [0.5, 0.6, 0.7, 0.8] },
|
|
],
|
|
"query": {"fusion": "rrf"},
|
|
},
|
|
)
|
|
assert response.ok, response.json()
|
|
rrf_result = response.json()["result"]["points"]
|
|
|
|
def get_id(x):
|
|
return x["id"]
|
|
|
|
# rrf order is not deterministic with same scores, so we need to sort by id
|
|
for expected, result in zip(sorted(rrf_expected, key=get_id), sorted(rrf_result, key=get_id)):
|
|
assert expected["id"] == result["id"]
|
|
assert expected.get("payload") == result.get("payload")
|
|
assert isclose(expected["score"], result["score"], rel_tol=1e-5)
|
|
|
|
|
|
def test_basic_dbsf():
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/search",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"vector": [0.1, 0.2, 0.3, 0.4],
|
|
"limit": 10,
|
|
},
|
|
)
|
|
assert response.ok
|
|
search_result_1 = response.json()["result"]
|
|
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/search",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"vector": [0.5, 0.6, 0.7, 0.8],
|
|
"limit": 10,
|
|
},
|
|
)
|
|
assert response.ok
|
|
search_result_2 = response.json()["result"]
|
|
|
|
dbsf_expected = distribution_based_score_fusion([search_result_1, search_result_2], limit=10)
|
|
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"prefetch": [
|
|
{ "query": [0.1, 0.2, 0.3, 0.4] },
|
|
{ "query": [0.5, 0.6, 0.7, 0.8] },
|
|
],
|
|
"query": {"fusion": "dbsf"},
|
|
},
|
|
)
|
|
assert response.ok, response.json()
|
|
dbsf_result = response.json()["result"]["points"]
|
|
|
|
for point, expected in zip(dbsf_result, dbsf_expected):
|
|
assert point["id"] == expected["id"]
|
|
assert point.get("payload") == expected.get("payload")
|
|
assert isclose(point["score"], expected["score"], rel_tol=1e-5)
|
|
|
|
|
|
@pytest.mark.parametrize("body", [
|
|
{
|
|
"prefetch": [
|
|
{ "query": [0.1, 0.2, 0.3, 0.4] },
|
|
{ "query": [0.5, 0.6, 0.7, 0.8] },
|
|
],
|
|
"query": {"fusion": "rrf"},
|
|
},
|
|
{ "query": [0.1, 0.2, 0.3, 0.4] }
|
|
])
|
|
def test_score_threshold(body):
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
**body
|
|
},
|
|
)
|
|
assert response.ok, response.json()
|
|
points = response.json()["result"]["points"]
|
|
|
|
assert len(points) == 8
|
|
score_threshold = points[3]["score"]
|
|
|
|
response = request_with_validation(
|
|
api="/collections/{collection_name}/points/query",
|
|
method="POST",
|
|
path_params={"collection_name": collection_name},
|
|
body={
|
|
"score_threshold": score_threshold,
|
|
**body
|
|
},
|
|
)
|
|
assert response.ok, response.json()
|
|
points = response.json()["result"]["points"]
|
|
|
|
assert len(points) < 8
|
|
for point in points:
|
|
assert point["score"] >= score_threshold
|