qdrant/tests/consensus_tests/fixtures.py

import random
from typing import List

import requests

from consensus_tests.assertions import assert_http_ok

CITIES = ["London", "New York", "Paris", "Tokyo", "Berlin", "Rome", "Madrid", "Moscow"]

# dense vector sizing
DENSE_VECTOR_SIZE = 4

# sparse vector sizing
SPARSE_VECTOR_SIZE = 1000
SPARSE_VECTOR_DENSITY = 0.1


# Generate a random dense vector
def random_dense_vector():
    return [random.random() for _ in range(DENSE_VECTOR_SIZE)]


# Generate a random sparse vector
def random_sparse_vector():
    num_non_zero = int(SPARSE_VECTOR_SIZE * SPARSE_VECTOR_DENSITY)
    indices: List[int] = random.sample(range(SPARSE_VECTOR_SIZE), num_non_zero)
    values: List[float] = [round(random.random(), 6) for _ in range(num_non_zero)]
    return {"indices": indices, "values": values}


def upsert_points(
        peer_url,
        points,
        collection_name="test_collection",
        wait="true",
        ordering="weak",
        shard_key=None,
        headers={},
) -> requests.Response:
    return requests.put(
        f"{peer_url}/collections/{collection_name}/points?wait={wait}&ordering={ordering}",
        json={
            "points": points,
            "shard_key": shard_key,
        },
        headers=headers,
    )


def update_points_vector(
        peer_url,
        points,
        collection_name="test_collection",
        wait="true",
):
    r_batch = requests.put(
        f"{peer_url}/collections/{collection_name}/points/vectors?wait={wait}",
        json={
            "points": [{"id": x, "vector": {"": random_dense_vector()}} for x in points],
        },
    )
    assert_http_ok(r_batch)
    return r_batch.json()


def update_points_payload(
        peer_url,
        points,
        collection_name="test_collection",
        wait="true",
        shard_key=None,
):
    r_batch = requests.post(
        f"{peer_url}/collections/{collection_name}/points/payload?wait={wait}",
        json={
            "points": points,
            "payload": {"city": random.choice(CITIES)},
            "shard_key": shard_key,
        },
    )
    assert_http_ok(r_batch)
    return r_batch.json()


def upsert_random_points(
    peer_url,
    num,
    collection_name="test_collection",
    fail_on_error=True,
    offset=0,
    batch_size=None,
    wait="true",
    ordering="weak",
    with_sparse_vector=True,
    shard_key=None,
    num_cities=None,
    headers={},
    extra_payload=None,
    timeout=None,
):
    extra_payload = extra_payload or {}

    def get_vector():
        # Create points in first peer's collection
        vector = {
            "": random_dense_vector(),
        }
        if with_sparse_vector:
            vector["sparse-text"] = random_sparse_vector()

        return vector

    while num > 0:
        size = num if batch_size is None else min(num, batch_size)

        r_batch = requests.put(
            f"{peer_url}/collections/{collection_name}/points?wait={wait}&ordering={ordering}",
            json={
                "points": [
                    {
                        "id": i + offset,
                        "vector": get_vector(),
                        "payload": {
                            **extra_payload,
                            "city": random.choice(CITIES[:num_cities]) if num_cities is not None else random.choice(CITIES)
                        },
                    }
                    for i in range(size)
                ],
                "shard_key": shard_key,
            },
            headers=headers,
            timeout=timeout,
        )
        if fail_on_error:
            assert_http_ok(r_batch)

        num -= size
        offset += size


def create_collection(
    peer_url,
    collection="test_collection",
    shard_number=1,
    replication_factor=1,
    write_consistency_factor=1,
    timeout=10,
    sharding_method=None,
    indexing_threshold=20000,
    headers={},
    strict_mode=None,
    sparse_vectors=True,
    default_segment_number=None,
    on_disk_payload=None,
    read_fan_out_factor=None,
    fail_on_error=True,
):
    payload = {
        "vectors": {"size": DENSE_VECTOR_SIZE, "distance": "Dot"},
        "shard_number": shard_number,
        "replication_factor": replication_factor,
        "write_consistency_factor": write_consistency_factor,
        "read_fan_out_factor": read_fan_out_factor,
        "sharding_method": sharding_method,
        "optimizers_config": {
            "indexing_threshold": indexing_threshold,
            "default_segment_number": default_segment_number,
        },
        "strict_mode_config": strict_mode,
        "on_disk_payload": on_disk_payload,
    }

    if sparse_vectors:
       payload["sparse_vectors"] = {"sparse-text": {}}

    # Create collection in peer_url
    r_batch = requests.put(
        f"{peer_url}/collections/{collection}?timeout={timeout}",
        json=payload,
        headers=headers,
    )
    if fail_on_error:
        assert_http_ok(r_batch)


def drop_collection(peer_url, collection="test_collection", timeout=10, headers={}):
    # Delete collection in peer_url
    r_delete = requests.delete(
        f"{peer_url}/collections/{collection}?timeout={timeout}", headers=headers
    )
    assert_http_ok(r_delete)


def create_field_index(
    peer_url,
    collection="test_collection",
    field_name="city",
    field_schema="keyword",
    headers={},
):
    # Create field index in peer_url
    r_batch = requests.put(
        f"{peer_url}/collections/{collection}/index",
        json={
            "field_name": field_name,
            "field_schema": field_schema,
        },
        headers=headers,
        params={"wait": "true"}
    )
    assert_http_ok(r_batch)


def search(peer_url, vector, city, collection="test_collection"):
    q = {
        "vector": vector,
        "top": 10,
        "with_vector": False,
        "with_payload": True,
        "filter": {"must": [{"key": "city", "match": {"value": city}}]},
    }
    r_search = requests.post(f"{peer_url}/collections/{collection}/points/search", json=q)
    assert_http_ok(r_search)
    return r_search.json()["result"]

def scroll(peer_url, city, collection="test_collection"):
    q = {
        "with_vector": False,
        "with_payload": True,
        "filter": {"must": [{"key": "city", "match": {"value": city}}]},
    }
    r_search = requests.post(f"{peer_url}/collections/{collection}/points/scroll", json=q)
    assert_http_ok(r_search)
    return r_search.json()["result"]["points"]

def count_counts(peer_url, collection="test_collection"):
    r_search = requests.post(
        f"{peer_url}/collections/{collection}/points/count",
        json={
            "exact": True,
        },
    )
    assert_http_ok(r_search)
    return r_search.json()["result"]["count"]


def set_strict_mode(peer_id, collection_name, strict_mode_config):
    requests.patch(
        f"{peer_id}/collections/{collection_name}",
        json={
            "strict_mode_config": strict_mode_config,
        },
    ).raise_for_status()


def get_telemetry_hw_info(peer_url, collection):
    r_search = requests.get(
        f"{peer_url}/telemetry", params="details_level=3"
    )
    assert_http_ok(r_search)
    hw = r_search.json()["result"]["hardware"]["collection_data"]
    if collection in hw:
        return hw[collection]
    else:
        return None

def get_telemetry_collections(peer_url):
    r_search = requests.get(
        f"{peer_url}/telemetry", params="details_level=3"
    )
    assert_http_ok(r_search)
    return r_search.json()["result"]['collections']['collections']