import random import uuid import numpy as np from qdrant_client._pydantic_compat import construct from qdrant_client.http import models from qdrant_client.http.models import SparseVector from qdrant_client.local.sparse import validate_sparse_vector from tests.fixtures.payload import one_random_payload_please def random_vectors( vector_sizes: dict[str, int] | int, ) -> models.VectorStruct: if isinstance(vector_sizes, int): return np.random.random(vector_sizes).round(3).tolist() elif isinstance(vector_sizes, dict): vectors = {} for vector_name, vector_size in vector_sizes.items(): vectors[vector_name] = np.random.random(vector_size).round(3).tolist() return vectors else: raise ValueError("vector_sizes must be int or dict") def random_multivectors(vector_sizes: dict[str, int] | int) -> models.VectorStruct: if isinstance(vector_sizes, int): vec_count = random.randint(1, 10) return generate_random_multivector(vector_sizes, vec_count) elif isinstance(vector_sizes, dict): vectors = {} for vector_name, vector_size in vector_sizes.items(): vec_count = random.randint(1, 10) vectors[vector_name] = generate_random_multivector(vector_size, vec_count) return vectors else: raise ValueError("vector_sizes must be int or dict") def generate_random_multivector(vec_size: int, vec_count: int) -> list[list[float]]: multivec = [] for _ in range(vec_count): multivec.append(np.random.random(vec_size).round(3).tolist()) return multivec # Generate random sparse vector with given size and density # The density is the probability of non-zero value over the whole vector def generate_random_sparse_vector(size: int, density: float) -> SparseVector: num_non_zero = int(size * density) indices: list[int] = random.sample(range(size), num_non_zero) values: list[float] = [round(random.random(), 6) for _ in range(num_non_zero)] sparse_vector = SparseVector(indices=indices, values=values) validate_sparse_vector(sparse_vector) return sparse_vector def generate_random_sparse_vector_uneven(size: int, density: float) -> SparseVector: if random.random() > 0.5: size = int(size * 0.3) return generate_random_sparse_vector(size, density) def generate_random_sparse_vector_list( num_vectors: int, vector_size: int, vector_density: float ) -> list[SparseVector]: sparse_vector_list = [] for _ in range(num_vectors): sparse_vector = generate_random_sparse_vector(vector_size, vector_density) sparse_vector_list.append(sparse_vector) return sparse_vector_list def random_sparse_vectors( vector_sizes: dict[str, int], even: bool = True, ) -> models.VectorStruct: vectors = {} for vector_name, vector_size in vector_sizes.items(): # use sparse vectors with 20% density if even: vectors[vector_name] = generate_random_sparse_vector(vector_size, density=0.2) else: vectors[vector_name] = generate_random_sparse_vector_uneven(vector_size, density=0.2) return vectors def generate_points( num_points: int, vector_sizes: dict[str, int] | int, with_payload: bool = False, random_ids: bool = False, skip_vectors: bool = False, sparse: bool = False, even_sparse: bool = True, multivector: bool = False, ) -> list[models.PointStruct]: if skip_vectors and isinstance(vector_sizes, int): raise ValueError("skip_vectors is not supported for single vector") points = [] for i in range(num_points): payload = None if with_payload: payload = one_random_payload_please(i) idx = i if random_ids: idx = str(uuid.uuid4()) if sparse: vectors = random_sparse_vectors(vector_sizes, even=even_sparse) elif multivector: vectors = random_multivectors(vector_sizes) else: vectors = random_vectors(vector_sizes) if skip_vectors: if random.random() > 0.8: vector_to_skip = random.choice(list(vectors.keys())) vectors.pop(vector_to_skip) points.append( construct( models.PointStruct, id=idx, vector=vectors, payload=payload, ) ) return points