qdrant/tests/openapi/test_query.py

from math import isclose
import pytest
import requests

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


@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)

    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(collection_name):
    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(collection_name, 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(collection_name):
    default_query_result = root_and_rescored_query(collection_name, [0.1, 0.2, 0.3, 0.4])

    nearest_query_result = root_and_rescored_query(collection_name, {"nearest": [0.1, 0.2, 0.3, 0.4]})

    assert default_query_result == nearest_query_result


# Test basic search with huge limit, it must not panic with allocation failure
# See: <https://github.com/qdrant/qdrant/issues/5483>
def test_basic_search_high_limit(collection_name):
    default_query_result = root_and_rescored_query(
        collection_name,
        [0.1, 0.2, 0.3, 0.4],
        limit=18446744073709551615, # u64::MAX
    )

    nearest_query_result = root_and_rescored_query(
        collection_name,
        {"nearest": [0.1, 0.2, 0.3, 0.4]},
        limit=18446744073709551615, # u64::MAX
    )

    assert default_query_result == nearest_query_result


def test_basic_scroll(collection_name):
    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(collection_name):
    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(collection_name):
    query_result = root_and_rescored_query(collection_name,
        {
            "recommend": {"positive": [1, 2, 3, 4], "negative": [3]},  # ids
        }
    )

    assert len(query_result) > 0


def test_basic_recommend_best_score(collection_name):
    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 len(query_result) > 0


def test_basic_discover(collection_name):
    query_result = root_and_rescored_query(collection_name,
        {
            "discover": {
                "target": 2,
                "context": [{"positive": 3, "negative": 4}],
            }
        },
    )

    assert len(query_result) > 0


def test_basic_context(collection_name):
    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 len(query_result) > 0


def test_basic_order_by(collection_name):
    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(collection_name, {"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(collection_name):
    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(collection_name):
    response = request_with_validation(
        api="/collections/{collection_name}/points/query",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "query": [0.1, 0.2, 0.3, 0.4],
            "limit": 10,
        },
    )
    assert response.ok
    search_result_1 = response.json()["result"]["points"]

    response = request_with_validation(
        api="/collections/{collection_name}/points/query",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "query": [0.5, 0.6, 0.7, 0.8],
            "limit": 10,
        },
    )
    assert response.ok
    search_result_2 = response.json()["result"]["points"]

    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_weighted_rrf(collection_name):
    """Test RRF with weights parameter."""
    response = request_with_validation(
        api="/collections/{collection_name}/points/query",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "query": [0.1, 0.2, 0.3, 0.4],
            "limit": 10,
        },
    )
    assert response.ok
    search_result_1 = response.json()["result"]["points"]

    response = request_with_validation(
        api="/collections/{collection_name}/points/query",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "query": [0.5, 0.6, 0.7, 0.8],
            "limit": 10,
        },
    )
    assert response.ok
    search_result_2 = response.json()["result"]["points"]

    # Test with weights [3.0, 1.0] - first source has 3x weight
    weights = [3.0, 1.0]
    rrf_expected = reciprocal_rank_fusion([search_result_1, search_result_2], limit=10, weights=weights)

    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": {
                "rrf": {
                    "weights": weights
                }
            },
        },
    )
    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)

    # Test with custom k parameter as well
    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": {
                "rrf": {
                    "k": 60,
                    "weights": [80, 20]
                }
            },
        },
    )
    assert response.ok, response.json()
    # Just verify it returns results without error
    assert len(response.json()["result"]["points"]) > 0


def test_basic_dbsf(collection_name):
    response = request_with_validation(
        api="/collections/{collection_name}/points/query",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "query": [0.1, 0.2, 0.3, 0.4],
            "limit": 10,
        },
    )
    assert response.ok
    search_result_1 = response.json()["result"]["points"]

    response = request_with_validation(
        api="/collections/{collection_name}/points/query",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "query": [0.5, 0.6, 0.7, 0.8],
            "limit": 10,
        },
    )
    assert response.ok
    search_result_2 = response.json()["result"]["points"]

    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)


def test_nearest_with_mmr(collection_name):
    # Regular nearest neighbor search
    response = request_with_validation(
        api="/collections/{collection_name}/points/query",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "query": {
                "nearest": [0.35, 0.08, 0.11, 0.47]
            },
        },
    )
    assert response.ok, response.json()
    search_result = response.json()["result"]["points"]

    response = request_with_validation(
        api="/collections/{collection_name}/points/query",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "query": {
                "nearest": [0.35, 0.08, 0.11, 0.47],
                "mmr": {
                    "diversity": 0.5,
                    "candidates_limit": 100
                }
            },
        },
    )
    assert response.ok, response.json()
    mmr_result = response.json()["result"]["points"]

    # Assert that both results are in different order
    search_ids = [point["id"] for point in search_result]
    mmr_ids = [point["id"] for point in mmr_result]
    assert search_ids != mmr_ids, "MMR should produce different ordering than regular search"

    # We didn't request vectors nor payloads
    for point in mmr_result:
        assert point.get("payload") is None
        assert point.get("vector") is None

    # Run request with default parameters
    response = request_with_validation(
        api="/collections/{collection_name}/points/query",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "query": {
                "nearest": [0.35, 0.08, 0.11, 0.47],
                "mmr": {}
            },
        },
    )
    assert response.ok, response.json()
    mmr_default_result = response.json()["result"]["points"]

    # Assert that both results are equal
    assert mmr_result == mmr_default_result, "MMR with explicit vs implicit defaults should produce the same output"


@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, collection_name):
    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

def test_query_missing_lookup_from_collection(collection_name):
    missing_collection = "missing_lookup_from_collection"
    lookup_from = {"collection": missing_collection, "vector": "default"}

    response = request_with_validation(
        api="/collections/{collection_name}/points/query",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "query": [0.1, 0.2, 0.3, 0.4],
            "limit": 3,
            "lookup_from": lookup_from,
        },
    )
    assert response.status_code == 404, response.text
    assert missing_collection in response.json()["status"]["error"]

    response = request_with_validation(
        api="/collections/{collection_name}/points/query/batch",
        method="POST",
        path_params={"collection_name": collection_name},
        body={
            "searches": [
                {
                    "query": [0.1, 0.2, 0.3, 0.4],
                    "limit": 3,
                    "lookup_from": lookup_from,
                }
            ]
        },
    )
    assert response.status_code == 404, response.text
    assert missing_collection in response.json()["status"]["error"]

    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],
                    "limit": 3,
                    "lookup_from": lookup_from,
                }
            ],
            "query": {"fusion": "rrf"},
            "limit": 3,
        },
    )
    assert response.status_code == 404, response.text
    assert missing_collection in response.json()["status"]["error"]