upstream-archive/qdrant-client
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant-client.git synced 2026-07-26 20:51:09 -05:00
Code Issues Packages Projects Releases Wiki Activity
Files
drop-python3.9
qdrant-client/tests/embed_tests/test_local_inference.py
George Panchuk baf2d0e303 new: replace optional and union with |
2025-12-02 12:08:55 +07:00

1738 lines
63 KiB
Python
Raw Permalink Blame History

from pathlib import Path
from unittest.mock import patch
from functools import lru_cache
import contextlib
import tempfile
import shutil
import numpy as np
import pytest
import qdrant_client.embed.embedder
from qdrant_client import QdrantClient, models
from qdrant_client.client_base import QdrantBase
from qdrant_client.qdrant_fastembed import IDF_EMBEDDING_MODELS
from qdrant_client.fastembed_common import (
TextEmbedding,
SparseTextEmbedding,
LateInteractionTextEmbedding,
ImageEmbedding,
)
COLLECTION_NAME = "inference_collection"
DENSE_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
DENSE_DIM = 384
SPARSE_MODEL_NAME = "Qdrant/bm42-all-minilm-l6-v2-attentions"
COLBERT_MODEL_NAME = "answerdotai/answerai-colbert-small-v1"
COLBERT_DIM = 96
DENSE_IMAGE_MODEL_NAME = "Qdrant/resnet50-onnx"
DENSE_IMAGE_DIM = 2048
TEST_IMAGE_PATH = Path(__file__).parent / "misc" / "image.jpeg"
SESSION_TMP_DIR: str = tempfile.mkdtemp()
def _get_cache_dir_name(model_name: str) -> Path:
return Path(SESSION_TMP_DIR) / model_name.lower().replace("/", "_").replace("-", "_")
@pytest.fixture(scope="module", autouse=True)
def cleanup():
yield
shutil.rmtree(SESSION_TMP_DIR)
@lru_cache
def _cached_text_embedding(model_name, *args, **kwargs):
cache_dir = kwargs.get("cache_dir") or _get_cache_dir_name(model_name)
kwargs["cache_dir"] = cache_dir
kwargs["local_files_only"] = kwargs.get("local_files_only", Path(cache_dir).exists())
return TextEmbedding(model_name=model_name, *args, **kwargs)
@lru_cache
def _cached_sparse_text_embedding(model_name, *args, **kwargs):
cache_dir = kwargs.get("cache_dir") or _get_cache_dir_name(model_name)
kwargs["cache_dir"] = cache_dir
kwargs["local_files_only"] = kwargs.get("local_files_only", Path(cache_dir).exists())
return SparseTextEmbedding(model_name=model_name, *args, **kwargs)
@lru_cache
def _cached_late_interaction_text_embedding(model_name, *args, **kwargs):
cache_dir = kwargs.get("cache_dir") or _get_cache_dir_name(model_name)
kwargs["cache_dir"] = cache_dir
kwargs["local_files_only"] = kwargs.get("local_files_only", Path(cache_dir).exists())
return LateInteractionTextEmbedding(model_name=model_name, *args, **kwargs)
@lru_cache
def _cached_image_embedding(model_name, *args, **kwargs):
cache_dir = kwargs.get("cache_dir") or _get_cache_dir_name(model_name)
kwargs["cache_dir"] = cache_dir
kwargs["local_files_only"] = kwargs.get("local_files_only", Path(cache_dir).exists())
return ImageEmbedding(model_name=model_name, *args, **kwargs)
@pytest.fixture(scope="module")
def cached_embeddings():
with contextlib.ExitStack() as stack:
stack.enter_context(
patch(
"qdrant_client.embed.embedder.TextEmbedding",
side_effect=_cached_text_embedding,
)
)
stack.enter_context(
patch(
"qdrant_client.embed.embedder.SparseTextEmbedding",
side_effect=_cached_sparse_text_embedding,
)
)
stack.enter_context(
patch(
"qdrant_client.embed.embedder.LateInteractionTextEmbedding",
side_effect=_cached_late_interaction_text_embedding,
)
)
stack.enter_context(
patch(
"qdrant_client.embed.embedder.ImageEmbedding",
side_effect=_cached_image_embedding,
)
)
yield
def arg_interceptor(func, kwarg_storage):
kwarg_storage.clear()
def wrapper(**kwargs):
kwarg_storage.update(kwargs)
return func(**kwargs)
return wrapper
def populate_dense_collection(
client: QdrantBase,
points: list[models.PointStruct],
vector_name: str | None = None,
collection_name: str = COLLECTION_NAME,
recreate: bool = True,
) -> None:
if recreate:
if client.collection_exists(collection_name):
client.delete_collection(collection_name)
vector_params = models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE)
vectors_config = {vector_name: vector_params} if vector_name else vector_params
client.create_collection(collection_name, vectors_config=vectors_config)
client.upsert(collection_name, points)
def populate_sparse_collection(
client: QdrantBase,
points: list[models.PointStruct],
vector_name: str,
collection_name: str = COLLECTION_NAME,
recreate: bool = True,
model_name: str = SPARSE_MODEL_NAME,
) -> None:
if recreate:
if client.collection_exists(collection_name):
client.delete_collection(collection_name)
sparse_vector_params = models.SparseVectorParams(
modifier=(
models.Modifier.IDF if model_name in IDF_EMBEDDING_MODELS else models.Modifier.NONE
)
)
sparse_vectors_config = {vector_name: sparse_vector_params}
client.create_collection(
collection_name, vectors_config={}, sparse_vectors_config=sparse_vectors_config
)
client.upsert(collection_name, points)
def test_upsert(cached_embeddings):
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
local_kwargs = {}
local_client._client.upsert = arg_interceptor(local_client._client.upsert, local_kwargs)
dense_doc_1 = models.Document(text="hello world", model=DENSE_MODEL_NAME)
dense_doc_2 = models.Document(text="bye world", model=DENSE_MODEL_NAME)
sparse_doc_1 = models.Document(text="hello world", model=SPARSE_MODEL_NAME)
sparse_doc_2 = models.Document(text="bye world", model=SPARSE_MODEL_NAME)
multi_doc_1 = models.Document(text="hello world", model=COLBERT_MODEL_NAME)
multi_doc_2 = models.Document(text="bye world", model=COLBERT_MODEL_NAME)
dense_image_1 = models.Image(image=TEST_IMAGE_PATH, model=DENSE_IMAGE_MODEL_NAME)
dense_image_2 = models.Image(image=TEST_IMAGE_PATH, model=DENSE_IMAGE_MODEL_NAME)
# region dense unnamed
points = [
models.PointStruct(id=1, vector=dense_doc_1),
models.PointStruct(id=2, vector=dense_doc_2),
]
populate_dense_collection(local_client, points)
vec_points = local_kwargs["points"]
assert all([isinstance(vec_point.vector, list) for vec_point in vec_points])
batch = models.Batch(ids=[1, 2], vectors=[dense_doc_1, dense_doc_2])
local_client.upsert(COLLECTION_NAME, batch)
batch = local_kwargs["points"]
assert all([isinstance(vector, list) for vector in batch.vectors])
local_client.delete_collection(COLLECTION_NAME)
# endregion
# region named vectors
vectors_config = {
"text": models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE),
"multi-text": models.VectorParams(
size=COLBERT_DIM,
distance=models.Distance.COSINE,
multivector_config=models.MultiVectorConfig(
comparator=models.MultiVectorComparator.MAX_SIM
),
),
"image": models.VectorParams(size=DENSE_IMAGE_DIM, distance=models.Distance.COSINE),
}
sparse_vectors_config = {
"sparse-text": models.SparseVectorParams(modifier=models.Modifier.IDF)
}
local_client.create_collection(
COLLECTION_NAME,
vectors_config=vectors_config,
sparse_vectors_config=sparse_vectors_config,
)
points = [
models.PointStruct(
id=1,
vector={
"text": dense_doc_1,
"multi-text": multi_doc_1,
"sparse-text": sparse_doc_1,
"image": dense_image_1,
},
),
models.PointStruct(
id=2,
vector={
"text": dense_doc_2,
"multi-text": multi_doc_2,
"sparse-text": sparse_doc_2,
"image": dense_image_2,
},
),
]
local_client.upsert(COLLECTION_NAME, points)
vec_points = local_kwargs["points"]
for vec_point in vec_points:
assert isinstance(vec_point.vector, dict)
assert isinstance(vec_point.vector["text"], list)
assert isinstance(vec_point.vector["multi-text"], list)
assert isinstance(vec_point.vector["sparse-text"], models.SparseVector)
assert isinstance(vec_point.vector["image"], list)
batch = models.Batch(
ids=[1, 2],
vectors={
"text": [dense_doc_1, dense_doc_2],
"multi-text": [multi_doc_1, multi_doc_2],
"sparse-text": [sparse_doc_1, sparse_doc_2],
"image": [dense_image_1, dense_image_2],
},
)
local_client.upsert(COLLECTION_NAME, batch)
batch = local_kwargs["points"]
vectors = batch.vectors
assert isinstance(vectors, dict)
assert all([isinstance(vector, list) for vector in vectors["text"]])
assert all([isinstance(vector, list) for vector in vectors["multi-text"]])
assert all([isinstance(vector, list) for vector in vectors["image"]])
assert all([isinstance(vector, models.SparseVector) for vector in vectors["sparse-text"]])
local_client.delete_collection(COLLECTION_NAME)
# endregion
def test_upload(cached_embeddings):
def recreate_collection(client, collection_name):
if client.collection_exists(collection_name):
client.delete_collection(collection_name)
vector_params = {
"text": models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE),
"image": models.VectorParams(size=DENSE_IMAGE_DIM, distance=models.Distance.COSINE),
}
client.create_collection(
collection_name,
vectors_config=vector_params,
sparse_vectors_config={
"sparse-text": models.SparseVectorParams(modifier=models.Modifier.IDF)
},
)
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
remote_client = QdrantClient()
dense_model_cache_dir = _get_cache_dir_name(DENSE_MODEL_NAME)
dense_options = {"cache_dir": str(dense_model_cache_dir)}
if dense_model_cache_dir.exists():
dense_options["local_files_only"] = True
dense_doc_1 = models.Document(
text="hello world", model=DENSE_MODEL_NAME, options=dense_options
)
dense_doc_2 = models.Document(text="bye world", model=DENSE_MODEL_NAME, options=dense_options)
dense_doc_3 = models.Document(
text="world world", model=DENSE_MODEL_NAME, options=dense_options
)
sparse_model_cache_dir = _get_cache_dir_name(SPARSE_MODEL_NAME)
sparse_options = {"cache_dir": str(sparse_model_cache_dir)}
if sparse_model_cache_dir.exists():
sparse_options["local_files_only"] = True
sparse_doc_1 = models.Document(
text="hello world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_2 = models.Document(
text="bye world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_3 = models.Document(
text="world world", model=SPARSE_MODEL_NAME, options=sparse_options
)
image_model_cache_dir = _get_cache_dir_name(DENSE_IMAGE_MODEL_NAME)
image_options = {"cache_dir": str(image_model_cache_dir)}
if image_model_cache_dir.exists():
image_options["local_files_only"] = True
dense_image_1 = models.Image(
image=TEST_IMAGE_PATH, model=DENSE_IMAGE_MODEL_NAME, options=image_options
)
dense_image_2 = models.Image(
image=TEST_IMAGE_PATH, model=DENSE_IMAGE_MODEL_NAME, options=image_options
)
dense_image_3 = models.Image(
image=TEST_IMAGE_PATH, model=DENSE_IMAGE_MODEL_NAME, options=image_options
)
points = [
models.PointStruct(
id=1, vector={"text": dense_doc_1, "image": dense_image_1, "sparse-text": sparse_doc_1}
),
models.PointStruct(
id=2, vector={"text": dense_doc_2, "image": dense_image_2, "sparse-text": sparse_doc_2}
),
models.PointStruct(
id=3, vector={"text": dense_doc_3, "image": dense_image_3, "sparse-text": sparse_doc_3}
),
]
recreate_collection(remote_client, COLLECTION_NAME)
remote_client.upload_points(COLLECTION_NAME, points, wait=True)
assert remote_client.count(COLLECTION_NAME).count == len(points)
assert isinstance(
remote_client.retrieve(COLLECTION_NAME, ids=[1], with_vectors=True)[0].vector["text"], list
) # assert doc has been substituted with its embedding
recreate_collection(remote_client, COLLECTION_NAME)
for text_doc in (dense_doc_1, dense_doc_2, dense_doc_3):
text_doc.options["local_files_only"] = Path(text_doc.options["cache_dir"]).exists()
for image_doc in (dense_image_1, dense_image_2, dense_image_3):
image_doc.options["local_files_only"] = Path(image_doc.options["cache_dir"]).exists()
for sparse_doc in (sparse_doc_1, sparse_doc_2, sparse_doc_3):
sparse_doc.options["local_files_only"] = Path(sparse_doc.options["cache_dir"]).exists()
vectors = [
{"text": dense_doc_1, "image": dense_image_1, "sparse-text": sparse_doc_1},
{"text": dense_doc_2, "image": dense_image_2, "sparse-text": sparse_doc_2},
{"text": dense_doc_3, "image": dense_image_3, "sparse-text": sparse_doc_3},
]
ids = list(range(len(vectors)))
remote_client.upload_collection(COLLECTION_NAME, ids=ids, vectors=vectors, wait=True)
assert remote_client.count(COLLECTION_NAME).count == len(vectors)
assert isinstance(
remote_client.retrieve(COLLECTION_NAME, ids=[1], with_vectors=True)[0].vector["text"], list
) # assert doc has been substituted with its embedding
recreate_collection(remote_client, COLLECTION_NAME)
remote_client.upload_points(COLLECTION_NAME, points, parallel=2, batch_size=2, wait=True)
assert remote_client.count(COLLECTION_NAME).count == len(points)
assert isinstance(
remote_client.retrieve(COLLECTION_NAME, ids=[1], with_vectors=True)[0].vector["text"], list
) # assert doc has been substituted with its embedding
recreate_collection(remote_client, COLLECTION_NAME)
remote_client.upload_collection(
COLLECTION_NAME, ids=ids, vectors=vectors, parallel=2, batch_size=2, wait=True
)
assert remote_client.count(COLLECTION_NAME).count == len(vectors)
assert isinstance(
remote_client.retrieve(COLLECTION_NAME, ids=[1], with_vectors=True)[0].vector["text"], list
) # assert doc has been substituted with its embedding
assert isinstance(points[0].vector["text"], models.Document)
recreate_collection(remote_client, COLLECTION_NAME)
remote_client.upload_points(COLLECTION_NAME, iter(points), parallel=2, batch_size=2, wait=True)
assert remote_client.count(COLLECTION_NAME).count == len(points)
assert isinstance(
remote_client.retrieve(COLLECTION_NAME, ids=[1], with_vectors=True)[0].vector["text"], list
) # assert doc has been substituted with its embedding
assert isinstance(vectors[0]["text"], models.Document)
recreate_collection(remote_client, COLLECTION_NAME)
remote_client.upload_collection(
COLLECTION_NAME, ids=ids, vectors=iter(vectors), parallel=2, batch_size=2, wait=True
)
assert remote_client.count(COLLECTION_NAME).count == len(vectors)
assert isinstance(
remote_client.retrieve(COLLECTION_NAME, ids=[1], with_vectors=True)[0].vector["text"], list
) # assert doc has been substituted with its embedding
def test_query_points(cached_embeddings):
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
local_kwargs = {}
local_client._client.query_points = arg_interceptor(
local_client._client.query_points, local_kwargs
)
sparse_model_cache_dir = _get_cache_dir_name(SPARSE_MODEL_NAME)
sparse_options = {"cache_dir": str(sparse_model_cache_dir)}
if Path(sparse_options["cache_dir"]).exists():
sparse_options["local_files_only"] = True
sparse_doc_1 = models.Document(
text="hello world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_2 = models.Document(
text="bye world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_3 = models.Document(
text="goodbye world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_4 = models.Document(
text="good afternoon world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_5 = models.Document(
text="good morning world", model=SPARSE_MODEL_NAME, options=sparse_options
)
points = [
models.PointStruct(id=i, vector={"sparse-text": doc}, payload={"content": doc.text})
for i, doc in enumerate(
[sparse_doc_1, sparse_doc_2, sparse_doc_3, sparse_doc_4, sparse_doc_5],
)
]
populate_sparse_collection(local_client, points, vector_name="sparse-text")
# region non-prefetch queries
local_client.query_points(COLLECTION_NAME, sparse_doc_1, using="sparse-text")
current_query = local_kwargs["query"]
assert isinstance(current_query.nearest, models.SparseVector)
retrieved_point_id_0 = local_client.retrieve(COLLECTION_NAME, ids=[0], with_vectors=True)[0]
# assert that we generate different embeddings for doc and query
nearest_query = models.NearestQuery(nearest=sparse_doc_1)
local_client.query_points(COLLECTION_NAME, nearest_query, using="sparse-text")
current_query = local_kwargs["query"]
assert isinstance(current_query.nearest, models.SparseVector)
recommend_query = models.RecommendQuery(
recommend=models.RecommendInput(
positive=[sparse_doc_1],
negative=[sparse_doc_1],
)
)
local_client.query_points(COLLECTION_NAME, recommend_query, using="sparse-text")
current_query = local_kwargs["query"]
assert all(
isinstance(vector, models.SparseVector) for vector in current_query.recommend.positive
)
assert all(
isinstance(vector, models.SparseVector) for vector in current_query.recommend.negative
)
discover_query = models.DiscoverQuery(
discover=models.DiscoverInput(
target=sparse_doc_1,
context=models.ContextPair(
positive=sparse_doc_2,
negative=sparse_doc_3,
),
)
)
local_client.query_points(COLLECTION_NAME, discover_query, using="sparse-text")
current_query = local_kwargs["query"]
assert isinstance(current_query.discover.target, models.SparseVector)
context_pair = current_query.discover.context
assert isinstance(context_pair.positive, models.SparseVector)
assert isinstance(context_pair.negative, models.SparseVector)
discover_query_list = models.DiscoverQuery(
discover=models.DiscoverInput(
target=sparse_doc_1,
context=[
models.ContextPair(
positive=sparse_doc_2,
negative=sparse_doc_3,
)
],
)
)
local_client.query_points(COLLECTION_NAME, discover_query_list, using="sparse-text")
current_query = local_kwargs["query"]
assert isinstance(current_query.discover.target, models.SparseVector)
context_pairs = current_query.discover.context
assert all(isinstance(pair.positive, models.SparseVector) for pair in context_pairs)
assert all(isinstance(pair.negative, models.SparseVector) for pair in context_pairs)
context_query = models.ContextQuery(
context=models.ContextPair(
positive=sparse_doc_1,
negative=sparse_doc_2,
)
)
local_client.query_points(COLLECTION_NAME, context_query, using="sparse-text")
current_query = local_kwargs["query"]
context = current_query.context
assert isinstance(context.positive, models.SparseVector)
assert isinstance(context.negative, models.SparseVector)
context_query_list = models.ContextQuery(
context=[
models.ContextPair(
positive=sparse_doc_1,
negative=sparse_doc_2,
),
models.ContextPair(
positive=sparse_doc_3,
negative=sparse_doc_4,
),
]
)
local_client.query_points(COLLECTION_NAME, context_query_list, using="sparse-text")
current_query = local_kwargs["query"]
contexts = current_query.context
assert all(isinstance(context.positive, models.SparseVector) for context in contexts)
assert all(isinstance(context.negative, models.SparseVector) for context in contexts)
# endregion
# region prefetch queries
prefetch = models.Prefetch(
query=nearest_query,
prefetch=models.Prefetch(
query=nearest_query,
prefetch=models.Prefetch(
query=nearest_query,
prefetch=[
models.Prefetch(query=discover_query_list, limit=5, using="sparse-text"),
models.Prefetch(query=nearest_query, using="sparse-text", limit=5),
],
using="sparse-text",
limit=4,
),
using="sparse-text",
limit=3,
),
using="sparse-text",
limit=2,
)
local_client.query_points(
COLLECTION_NAME, query=nearest_query, prefetch=prefetch, limit=1, using="sparse-text"
)
current_query = local_kwargs["query"]
current_prefetch = local_kwargs["prefetch"]
assert isinstance(current_query.nearest, models.SparseVector)
assert isinstance(current_prefetch.query.nearest, models.SparseVector)
assert isinstance(current_prefetch.prefetch.query.nearest, models.SparseVector)
assert isinstance(current_prefetch.prefetch.prefetch.query.nearest, models.SparseVector)
assert isinstance(
current_prefetch.prefetch.prefetch.prefetch[0].query.discover.target, models.SparseVector
)
context_pairs = current_prefetch.prefetch.prefetch.prefetch[0].query.discover.context
assert all(isinstance(pair.positive, models.SparseVector) for pair in context_pairs)
assert all(isinstance(pair.negative, models.SparseVector) for pair in context_pairs)
assert isinstance(
current_prefetch.prefetch.prefetch.prefetch[1].query.nearest, models.SparseVector
)
# endregion
local_client.delete_collection(COLLECTION_NAME)
def test_query_points_is_query(cached_embeddings):
# dense_model_name = "jinaai/jina-embeddings-v3"
# dense_dim = 1024
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
local_kwargs = {}
local_client._client.query_points = arg_interceptor(
local_client._client.query_points, local_kwargs
)
# dense_doc_1 = models.Document(text="hello world", model=dense_model_name) # todo: uncomment once this model is supported
sparse_model_cache_dir = _get_cache_dir_name(SPARSE_MODEL_NAME)
sparse_options = {"cache_dir": str(sparse_model_cache_dir)}
if Path(sparse_options["cache_dir"]).exists():
sparse_options["local_files_only"] = True
colbert_model_cache_dir = _get_cache_dir_name(COLBERT_MODEL_NAME)
colbert_options = {"cache_dir": str(colbert_model_cache_dir)}
if Path(colbert_options["cache_dir"]).exists():
colbert_options["local_files_only"] = True
sparse_doc_1 = models.Document(
text="hello world", model=SPARSE_MODEL_NAME, options=sparse_options
)
colbert_doc_1 = models.Document(
text="hello world", model=COLBERT_MODEL_NAME, options=colbert_options
)
vectors_config = {
# "dense-text": models.VectorParams(size=dense_dim, distance=models.Distance.COSINE),
"colbert-text": models.VectorParams(
size=COLBERT_DIM,
distance=models.Distance.COSINE,
multivector_config=models.MultiVectorConfig(
comparator=models.MultiVectorComparator.MAX_SIM
),
),
}
sparse_vectors_config = {
"sparse-text": models.SparseVectorParams(modifier=models.Modifier.IDF)
}
local_client.create_collection(
COLLECTION_NAME,
vectors_config=vectors_config,
sparse_vectors_config=sparse_vectors_config,
)
points = [
models.PointStruct(
id=0, vector={"colbert-text": colbert_doc_1, "sparse-text": sparse_doc_1}
)
]
local_client.upsert(COLLECTION_NAME, points)
retrieved_point = local_client.retrieve(COLLECTION_NAME, ids=[0], with_vectors=True)[0]
# local_client.query_points(COLLECTION_NAME, dense_doc_1, using="dense-text")
# assert isinstance(local_kwargs["query"].nearest, list)
# assert not np.allclose(retrieved_point.vector["dense-text"], local_kwargs["query"].nearest, atol=1e-3)
local_client.query_points(COLLECTION_NAME, sparse_doc_1, using="sparse-text")
assert isinstance(local_kwargs["query"].nearest, models.SparseVector)
assert not np.allclose(
retrieved_point.vector["sparse-text"].values,
local_kwargs["query"].nearest.values,
atol=1e-3,
)
local_client.query_points(COLLECTION_NAME, colbert_doc_1, using="colbert-text")
assert isinstance(local_kwargs["query"].nearest, list)
# colbert has a min number of 32 tokens for query
assert len(retrieved_point.vector["colbert-text"]) != len(local_kwargs["query"].nearest)
local_client.delete_collection(COLLECTION_NAME)
def test_query_points_groups(cached_embeddings):
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
local_kwargs = {}
local_client._client.query_points_groups = arg_interceptor(
local_client._client.query_points_groups, local_kwargs
)
sparse_model_cache_dir = _get_cache_dir_name(SPARSE_MODEL_NAME)
sparse_options = {"cache_dir": str(sparse_model_cache_dir)}
if Path(sparse_options["cache_dir"]).exists():
sparse_options["local_files_only"] = True
sparse_doc_1 = models.Document(
text="hello world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_2 = models.Document(
text="bye world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_3 = models.Document(
text="goodbye world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_4 = models.Document(
text="good afternoon world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_5 = models.Document(
text="good morning world", model=SPARSE_MODEL_NAME, options=sparse_options
)
points = [
models.PointStruct(id=i, vector={"sparse-text": doc}, payload={"content": doc.text})
for i, doc in enumerate(
[sparse_doc_1, sparse_doc_2, sparse_doc_3, sparse_doc_4, sparse_doc_5],
)
]
populate_sparse_collection(local_client, points, vector_name="sparse-text")
# region query_points_groups
local_client.query_points_groups(
COLLECTION_NAME, group_by="content", query=sparse_doc_1, using="sparse-text"
)
current_query = local_kwargs["query"]
assert isinstance(current_query.nearest, models.SparseVector)
retrieved_point_id_0 = local_client.retrieve(COLLECTION_NAME, ids=[0], with_vectors=True)[0]
# assert that we generate different embeddings for doc and query
# we are using sparse_doc_1 as a query
assert not (
np.allclose(
retrieved_point_id_0.vector["sparse-text"].values,
current_query.nearest.values,
atol=1e-3,
)
)
prefetch_1 = models.Prefetch(
query=models.NearestQuery(nearest=sparse_doc_2), using="sparse-text", limit=3
)
prefetch_2 = models.Prefetch(
query=models.NearestQuery(nearest=sparse_doc_3), using="sparse-text", limit=3
)
local_client.query_points_groups(
COLLECTION_NAME,
group_by="content",
query=sparse_doc_1,
prefetch=[prefetch_1, prefetch_2],
using="sparse-text",
)
current_query = local_kwargs["query"]
current_prefetch = local_kwargs["prefetch"]
assert isinstance(current_query.nearest, models.SparseVector)
assert isinstance(current_prefetch[0].query.nearest, models.SparseVector)
assert isinstance(current_prefetch[1].query.nearest, models.SparseVector)
assert not (
np.allclose(
retrieved_point_id_0.vector["sparse-text"].values,
current_query.nearest.values,
atol=1e-3,
)
)
retrieved_point_id_1 = local_client.retrieve(COLLECTION_NAME, ids=[1], with_vectors=True)[0]
assert not (
np.allclose(
retrieved_point_id_1.vector["sparse-text"].values,
current_prefetch[0].query.nearest.values,
atol=1e-3,
)
)
assert isinstance(prefetch_1.query.nearest, models.Document)
local_kwargs.clear()
local_client.query_points_groups(
COLLECTION_NAME,
group_by="content",
query=sparse_doc_1,
prefetch=prefetch_1,
using="sparse-text",
)
current_prefetch = local_kwargs["prefetch"]
assert isinstance(current_prefetch.query.nearest, models.SparseVector)
assert not (
np.allclose(
retrieved_point_id_1.vector["sparse-text"].values,
current_prefetch.query.nearest.values,
atol=1e-3,
)
)
# endregion
local_client.delete_collection(COLLECTION_NAME)
def test_query_batch_points(cached_embeddings):
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
local_kwargs = {}
local_client._client.query_batch_points = arg_interceptor(
local_client._client.query_batch_points, local_kwargs
)
sparse_model_cache_dir = _get_cache_dir_name(SPARSE_MODEL_NAME)
sparse_options = {"cache_dir": str(sparse_model_cache_dir)}
if Path(sparse_options["cache_dir"]).exists():
sparse_options["local_files_only"] = True
sparse_doc_1 = models.Document(
text="hello world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_2 = models.Document(
text="bye world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_3 = models.Document(
text="goodbye world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_4 = models.Document(
text="good afternoon world", model=SPARSE_MODEL_NAME, options=sparse_options
)
sparse_doc_5 = models.Document(
text="good morning world", model=SPARSE_MODEL_NAME, options=sparse_options
)
points = [
models.PointStruct(id=i, vector={"sparse-text": dense_doc})
for i, dense_doc in enumerate(
[sparse_doc_1, sparse_doc_2, sparse_doc_3, sparse_doc_4, sparse_doc_5]
)
]
populate_sparse_collection(local_client, points, vector_name="sparse-text")
prefetch_1 = models.Prefetch(
query=models.NearestQuery(nearest=sparse_doc_2), limit=3, using="sparse-text"
)
prefetch_2 = models.Prefetch(
query=models.NearestQuery(nearest=sparse_doc_3), limit=3, using="sparse-text"
)
query_requests = [
models.QueryRequest(query=models.NearestQuery(nearest=sparse_doc_1), using="sparse-text"),
models.QueryRequest(
query=models.NearestQuery(nearest=sparse_doc_2),
prefetch=[prefetch_1, prefetch_2],
using="sparse-text",
),
]
local_client.query_batch_points(COLLECTION_NAME, query_requests)
current_requests = local_kwargs["requests"]
assert all(
[isinstance(request.query.nearest, models.SparseVector) for request in current_requests]
)
assert all(
[
isinstance(prefetch.query.nearest, models.SparseVector)
for prefetch in current_requests[1].prefetch
]
)
retrieved_point = local_client.retrieve(COLLECTION_NAME, ids=[0], with_vectors=True)[0]
assert not np.allclose(
retrieved_point.vector["sparse-text"].values,
current_requests[0].query.nearest.values,
atol=1e-3,
)
local_client.delete_collection(COLLECTION_NAME)
def test_batch_update_points(cached_embeddings):
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
local_kwargs = {}
local_client._client.batch_update_points = arg_interceptor(
local_client._client.batch_update_points, local_kwargs
)
dense_model_cache_dir = _get_cache_dir_name(DENSE_MODEL_NAME)
dense_options = {"cache_dir": str(dense_model_cache_dir)}
if Path(dense_options["cache_dir"]).exists():
dense_options["local_files_only"] = True
dense_doc_1 = models.Document(
text="hello world", model=DENSE_MODEL_NAME, options=dense_options
)
dense_doc_2 = models.Document(text="bye world", model=DENSE_MODEL_NAME, options=dense_options)
# region unnamed
points = [
models.PointStruct(id=1, vector=dense_doc_1),
models.PointStruct(id=2, vector=dense_doc_2),
]
populate_dense_collection(local_client, points)
batch = models.Batch(ids=[2, 3], vectors=[dense_doc_1, dense_doc_2])
upsert_operation = models.UpsertOperation(upsert=models.PointsBatch(batch=batch))
local_client.batch_update_points(COLLECTION_NAME, [upsert_operation])
current_operation = local_kwargs["update_operations"][0]
current_batch = current_operation.upsert.batch
assert all([isinstance(vector, list) for vector in current_batch.vectors])
new_points = [
models.PointStruct(id=3, vector=dense_doc_1),
models.PointStruct(id=4, vector=dense_doc_2),
]
upsert_operation = models.UpsertOperation(upsert=models.PointsList(points=new_points))
local_client.batch_update_points(COLLECTION_NAME, [upsert_operation])
current_operation = local_kwargs["update_operations"][0]
current_batch = current_operation.upsert.points
assert all([isinstance(vector.vector, list) for vector in current_batch])
update_vectors_operation = models.UpdateVectorsOperation(
update_vectors=models.UpdateVectors(points=[models.PointVectors(id=1, vector=dense_doc_2)])
)
upsert_operation = models.UpsertOperation(
upsert=models.PointsList(points=[models.PointStruct(id=5, vector=dense_doc_2)])
)
local_client.batch_update_points(COLLECTION_NAME, [update_vectors_operation, upsert_operation])
current_update_operation = local_kwargs["update_operations"][0]
current_upsert_operation = local_kwargs["update_operations"][1]
assert all(
[
isinstance(vector.vector, list)
for vector in current_update_operation.update_vectors.points
]
)
assert all(
[isinstance(vector.vector, list) for vector in current_upsert_operation.upsert.points]
)
local_client.delete_collection(COLLECTION_NAME)
# endregion
# region named
points = [
models.PointStruct(id=1, vector={"text": dense_doc_1}),
models.PointStruct(id=2, vector={"text": dense_doc_2}),
]
populate_dense_collection(local_client, points, vector_name="text")
batch = models.Batch(ids=[2, 3], vectors={"text": [dense_doc_1, dense_doc_2]})
upsert_operation = models.UpsertOperation(upsert=models.PointsBatch(batch=batch))
local_client.batch_update_points(COLLECTION_NAME, [upsert_operation])
current_operation = local_kwargs["update_operations"][0]
current_batch = current_operation.upsert.batch
assert all([isinstance(vector, list) for vector in current_batch.vectors.values()])
new_points = [
models.PointStruct(id=3, vector={"text": dense_doc_1}),
models.PointStruct(id=4, vector={"text": dense_doc_2}),
]
upsert_operation = models.UpsertOperation(upsert=models.PointsList(points=new_points))
local_client.batch_update_points(COLLECTION_NAME, [upsert_operation])
current_operation = local_kwargs["update_operations"][0]
current_batch = current_operation.upsert.points
assert all([isinstance(vector.vector["text"], list) for vector in current_batch])
update_vectors_operation = models.UpdateVectorsOperation(
update_vectors=models.UpdateVectors(
points=[models.PointVectors(id=1, vector={"text": dense_doc_2})]
)
)
upsert_operation = models.UpsertOperation(
upsert=models.PointsList(points=[models.PointStruct(id=5, vector={"text": dense_doc_2})])
)
local_client.batch_update_points(COLLECTION_NAME, [update_vectors_operation, upsert_operation])
current_update_operation = local_kwargs["update_operations"][0]
current_upsert_operation = local_kwargs["update_operations"][1]
assert all(
[
isinstance(vector.vector["text"], list)
for vector in current_update_operation.update_vectors.points
]
)
assert all(
[
isinstance(vector.vector["text"], list)
for vector in current_upsert_operation.upsert.points
]
)
local_client.delete_collection(COLLECTION_NAME)
# endregion
def test_update_vectors(cached_embeddings):
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
local_kwargs = {}
local_client._client.update_vectors = arg_interceptor(
local_client._client.update_vectors, local_kwargs
)
dense_model_cache_dir = _get_cache_dir_name(DENSE_MODEL_NAME)
dense_options = {"cache_dir": str(dense_model_cache_dir)}
if Path(dense_options["cache_dir"]).exists():
dense_options["local_files_only"] = True
dense_doc_1 = models.Document(
text="hello world",
model=DENSE_MODEL_NAME,
options=dense_options,
)
dense_doc_2 = models.Document(
text="bye world",
model=DENSE_MODEL_NAME,
options=dense_options,
)
dense_doc_3 = models.Document(
text="goodbye world",
model=DENSE_MODEL_NAME,
options=dense_options,
)
# region unnamed
points = [
models.PointStruct(id=1, vector=dense_doc_1),
models.PointStruct(id=2, vector=dense_doc_2),
]
populate_dense_collection(local_client, points)
point_vectors = [
models.PointVectors(id=1, vector=dense_doc_2),
models.PointVectors(id=2, vector=dense_doc_3),
]
local_client.update_vectors(COLLECTION_NAME, point_vectors)
current_vectors = local_kwargs["points"]
assert all([isinstance(vector.vector, list) for vector in current_vectors])
local_client.delete_collection(COLLECTION_NAME)
# endregion
# region named
points = [
models.PointStruct(id=1, vector={"text": dense_doc_1}),
models.PointStruct(id=2, vector={"text": dense_doc_2}),
]
populate_dense_collection(local_client, points, vector_name="text")
point_vectors = [
models.PointVectors(id=1, vector={"text": dense_doc_2}),
models.PointVectors(id=2, vector={"text": dense_doc_3}),
]
local_client.update_vectors(COLLECTION_NAME, point_vectors)
current_vectors = local_kwargs["points"]
assert all([isinstance(vector.vector["text"], list) for vector in current_vectors])
local_client.delete_collection(COLLECTION_NAME)
# endregion
def test_propagate_options(cached_embeddings):
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
dense_doc_1 = models.Document(
text="hello world", model=DENSE_MODEL_NAME, options={"lazy_load": True}
)
sparse_doc_1 = models.Document(
text="hello world", model=SPARSE_MODEL_NAME, options={"lazy_load": True}
)
multi_doc_1 = models.Document(
text="hello world", model=COLBERT_MODEL_NAME, options={"lazy_load": True}
)
dense_image_1 = models.Image(
image=TEST_IMAGE_PATH, model=DENSE_IMAGE_MODEL_NAME, options={"lazy_load": True}
)
points = [
models.PointStruct(
id=1,
vector={
"text": dense_doc_1,
"multi-text": multi_doc_1,
"sparse-text": sparse_doc_1,
"image": dense_image_1,
},
)
]
vectors_config = {
"text": models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE),
"multi-text": models.VectorParams(
size=COLBERT_DIM,
distance=models.Distance.COSINE,
multivector_config=models.MultiVectorConfig(
comparator=models.MultiVectorComparator.MAX_SIM
),
),
"image": models.VectorParams(size=DENSE_IMAGE_DIM, distance=models.Distance.COSINE),
}
sparse_vectors_config = {
"sparse-text": models.SparseVectorParams(modifier=models.Modifier.IDF)
}
local_client.create_collection(
COLLECTION_NAME,
vectors_config=vectors_config,
sparse_vectors_config=sparse_vectors_config,
)
local_client.upsert(COLLECTION_NAME, points)
assert local_client._model_embedder.embedder.embedding_models[DENSE_MODEL_NAME][
0
].model.model.lazy_load
assert local_client._model_embedder.embedder.sparse_embedding_models[SPARSE_MODEL_NAME][
0
].model.model.lazy_load
assert local_client._model_embedder.embedder.late_interaction_embedding_models[
COLBERT_MODEL_NAME
][0].model.model.lazy_load
assert local_client._model_embedder.embedder.image_embedding_models[DENSE_IMAGE_MODEL_NAME][
0
].model.model.lazy_load
local_client._model_embedder.embedder.embedding_models.clear()
local_client._model_embedder.embedder.sparse_embedding_models.clear()
local_client._model_embedder.embedder.late_interaction_embedding_models.clear()
local_client._model_embedder.embedder.image_embedding_models.clear()
inference_object_dense_doc_1 = models.InferenceObject(
object="hello world",
model=DENSE_MODEL_NAME,
options={"lazy_load": True},
)
inference_object_sparse_doc_1 = models.InferenceObject(
object="hello world",
model=SPARSE_MODEL_NAME,
options={"lazy_load": True},
)
inference_object_multi_doc_1 = models.InferenceObject(
object="hello world",
model=COLBERT_MODEL_NAME,
options={"lazy_load": True},
)
inference_object_dense_image_1 = models.InferenceObject(
object=TEST_IMAGE_PATH,
model=DENSE_IMAGE_MODEL_NAME,
options={"lazy_load": True},
)
points = [
models.PointStruct(
id=2,
vector={
"text": inference_object_dense_doc_1,
"multi-text": inference_object_multi_doc_1,
"sparse-text": inference_object_sparse_doc_1,
"image": inference_object_dense_image_1,
},
)
]
local_client.upsert(COLLECTION_NAME, points)
assert local_client._model_embedder.embedder.embedding_models[DENSE_MODEL_NAME][
0
].model.model.lazy_load
assert local_client._model_embedder.embedder.sparse_embedding_models[SPARSE_MODEL_NAME][
0
].model.model.lazy_load
assert local_client._model_embedder.embedder.late_interaction_embedding_models[
COLBERT_MODEL_NAME
][0].model.model.lazy_load
assert local_client._model_embedder.embedder.image_embedding_models[DENSE_IMAGE_MODEL_NAME][
0
].model.model.lazy_load
def test_inference_object(cached_embeddings):
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
local_kwargs = {}
local_client._client.upsert = arg_interceptor(local_client._client.upsert, local_kwargs)
inference_object_dense_doc_1 = models.InferenceObject(
object="hello world",
model=DENSE_MODEL_NAME,
options={"lazy_load": True},
)
inference_object_sparse_doc_1 = models.InferenceObject(
object="hello world",
model=SPARSE_MODEL_NAME,
options={"lazy_load": True},
)
inference_object_multi_doc_1 = models.InferenceObject(
object="hello world",
model=COLBERT_MODEL_NAME,
options={"lazy_load": True},
)
inference_object_dense_image_1 = models.InferenceObject(
object=TEST_IMAGE_PATH,
model=DENSE_IMAGE_MODEL_NAME,
options={"lazy_load": True},
)
points = [
models.PointStruct(
id=1,
vector={
"text": inference_object_dense_doc_1,
"multi-text": inference_object_multi_doc_1,
"sparse-text": inference_object_sparse_doc_1,
"image": inference_object_dense_image_1,
},
)
]
vectors_config = {
"text": models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE),
"multi-text": models.VectorParams(
size=COLBERT_DIM,
distance=models.Distance.COSINE,
multivector_config=models.MultiVectorConfig(
comparator=models.MultiVectorComparator.MAX_SIM
),
),
"image": models.VectorParams(size=DENSE_IMAGE_DIM, distance=models.Distance.COSINE),
}
sparse_vectors_config = {
"sparse-text": models.SparseVectorParams(modifier=models.Modifier.IDF)
}
local_client.create_collection(
COLLECTION_NAME,
vectors_config=vectors_config,
sparse_vectors_config=sparse_vectors_config,
)
local_client.upsert(COLLECTION_NAME, points)
vec_points = local_kwargs["points"]
vector = vec_points[0].vector
assert isinstance(vector["text"], list)
assert isinstance(vector["multi-text"], list)
assert isinstance(vector["sparse-text"], models.SparseVector)
assert isinstance(vector["image"], list)
assert local_client.scroll(COLLECTION_NAME, limit=1, with_vectors=True)[0]
local_client.query_points(COLLECTION_NAME, inference_object_dense_doc_1, using="text")
local_client.query_points(COLLECTION_NAME, inference_object_sparse_doc_1, using="sparse-text")
local_client.query_points(COLLECTION_NAME, inference_object_multi_doc_1, using="multi-text")
local_client.query_points(COLLECTION_NAME, inference_object_dense_image_1, using="image")
local_client.delete_collection(COLLECTION_NAME)
@pytest.mark.parametrize("parallel", [1, 2])
def test_upload_mixed_batches_upload_points(parallel, cached_embeddings):
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
remote_client = QdrantClient()
half_dense_dim = DENSE_DIM // 2
batch_size = 2
ref_vector = [0.0, 0.2] * half_dense_dim
norm_ref_vector = (np.array(ref_vector) / np.linalg.norm(ref_vector)).tolist()
# region separate plain batches
dense_model_cache_dir = _get_cache_dir_name(DENSE_MODEL_NAME)
dense_options = {"cache_dir": str(dense_model_cache_dir)}
if Path(dense_options["cache_dir"]).exists():
dense_options["local_files_only"] = True
points = [
models.PointStruct(
id=1,
vector=models.Document(
text="hello world", model=DENSE_MODEL_NAME, options=dense_options
),
),
models.PointStruct(
id=2,
vector=models.Document(
text="bye world", model=DENSE_MODEL_NAME, options=dense_options
),
),
models.PointStruct(id=3, vector=ref_vector),
models.PointStruct(id=4, vector=[0.1, 0.2] * half_dense_dim),
]
vectors_config = models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE)
if remote_client.collection_exists(COLLECTION_NAME):
remote_client.delete_collection(COLLECTION_NAME)
remote_client.create_collection(COLLECTION_NAME, vectors_config=vectors_config)
remote_client.upload_points(
COLLECTION_NAME, points, batch_size=batch_size, wait=True, parallel=parallel
)
assert remote_client.count(COLLECTION_NAME).count == len(points)
assert np.allclose(
remote_client.retrieve(COLLECTION_NAME, ids=[3], with_vectors=True)[0].vector,
norm_ref_vector,
)
remote_client.delete_collection(COLLECTION_NAME)
# endregion
# region mixed plain batches
if Path(dense_options["cache_dir"]).exists():
dense_options["local_files_only"] = True
points = [
models.PointStruct(
id=1,
vector=models.Document(
text="hello world", model=DENSE_MODEL_NAME, options=dense_options
),
),
models.PointStruct(id=2, vector=ref_vector),
models.PointStruct(
id=3,
vector=models.Document(
text="bye world", model=DENSE_MODEL_NAME, options=dense_options
),
),
models.PointStruct(id=4, vector=[0.1, 0.2] * half_dense_dim),
]
remote_client.create_collection(COLLECTION_NAME, vectors_config=vectors_config)
remote_client.upload_points(
COLLECTION_NAME, points, batch_size=batch_size, wait=True, parallel=parallel
)
assert remote_client.count(COLLECTION_NAME).count == len(points)
assert np.allclose(
remote_client.retrieve(COLLECTION_NAME, ids=[2], with_vectors=True)[0].vector,
norm_ref_vector,
)
remote_client.delete_collection(COLLECTION_NAME)
# endregion
# region mixed named batches
vectors_config = {
"text": models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE),
"plain": models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE),
}
points = [
models.PointStruct(
id=1,
vector={
"text": models.Document(
text="hello world", model=DENSE_MODEL_NAME, options=dense_options
),
"plain": [0.1, 0.2] * half_dense_dim,
},
),
models.PointStruct(
id=2,
vector={
"plain": ref_vector,
"text": models.Document(
text="bye world", model=DENSE_MODEL_NAME, options=dense_options
),
},
),
models.PointStruct(
id=3,
vector={"plain": [0.3, 0.2] * half_dense_dim},
),
models.PointStruct(
id=4,
vector={
"text": models.Document(
text="bye world", model=DENSE_MODEL_NAME, options=dense_options
)
},
),
]
remote_client.create_collection(COLLECTION_NAME, vectors_config=vectors_config)
remote_client.upload_points(
COLLECTION_NAME, points, batch_size=batch_size, wait=True, parallel=parallel
)
assert remote_client.count(COLLECTION_NAME).count == len(points)
assert np.allclose(
remote_client.retrieve(COLLECTION_NAME, ids=[2], with_vectors=True)[0].vector["plain"],
norm_ref_vector,
)
remote_client.delete_collection(COLLECTION_NAME)
# endregion
@pytest.mark.parametrize("parallel", [1, 2])
def test_upload_mixed_batches_upload_collection(parallel, cached_embeddings):
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
remote_client = QdrantClient()
half_dense_dim = DENSE_DIM // 2
batch_size = 2
ref_vector = [0.0, 0.2] * half_dense_dim
norm_ref_vector = (np.array(ref_vector) / np.linalg.norm(ref_vector)).tolist()
# region separate plain batches
ids = [0, 1, 2, 3]
dense_model_cache_dir = _get_cache_dir_name(DENSE_MODEL_NAME)
dense_options = {"cache_dir": str(dense_model_cache_dir)}
if Path(dense_options["cache_dir"]).exists():
dense_options["local_files_only"] = True
vectors = [
models.Document(text="hello world", model=DENSE_MODEL_NAME, options=dense_options),
models.Document(text="bye world", model=DENSE_MODEL_NAME, options=dense_options),
ref_vector,
[0.1, 0.2] * half_dense_dim,
]
vectors_config = models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE)
if remote_client.collection_exists(COLLECTION_NAME):
remote_client.delete_collection(COLLECTION_NAME)
remote_client.create_collection(COLLECTION_NAME, vectors_config=vectors_config)
remote_client.upload_collection(
COLLECTION_NAME,
ids=ids,
vectors=vectors,
batch_size=batch_size,
wait=True,
parallel=parallel,
)
assert remote_client.count(COLLECTION_NAME).count == len(vectors)
assert np.allclose(
remote_client.retrieve(COLLECTION_NAME, ids=[2], with_vectors=True)[0].vector,
norm_ref_vector,
)
remote_client.delete_collection(COLLECTION_NAME)
# endregion
# region mixed plain batches
if Path(dense_options["cache_dir"]).exists():
dense_options["local_files_only"] = True
vectors = [
models.Document(text="hello world", model=DENSE_MODEL_NAME, options=dense_options),
ref_vector,
models.Document(text="bye world", model=DENSE_MODEL_NAME, options=dense_options),
[0.1, 0.2] * half_dense_dim,
]
remote_client.create_collection(COLLECTION_NAME, vectors_config=vectors_config)
remote_client.upload_collection(
COLLECTION_NAME,
ids=ids,
vectors=vectors,
batch_size=batch_size,
wait=True,
parallel=parallel,
)
assert remote_client.count(COLLECTION_NAME).count == len(vectors)
assert np.allclose(
remote_client.retrieve(COLLECTION_NAME, ids=[1], with_vectors=True)[0].vector,
norm_ref_vector,
)
remote_client.delete_collection(COLLECTION_NAME)
# endregion
# region mixed named batches
vectors_config = {
"text": models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE),
"plain": models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE),
}
vectors = [
{
"text": models.Document(
text="hello world", model=DENSE_MODEL_NAME, options=dense_options
),
"plain": [0.0, 0.2] * half_dense_dim,
},
{
"plain": ref_vector,
"text": models.Document(
text="bye world", model=DENSE_MODEL_NAME, options=dense_options
),
},
{"plain": [0.3, 0.2] * half_dense_dim},
{"text": models.Document(text="bye world", model=DENSE_MODEL_NAME, options=dense_options)},
]
remote_client.create_collection(COLLECTION_NAME, vectors_config=vectors_config)
remote_client.upload_collection(
COLLECTION_NAME,
ids=ids,
vectors=vectors,
batch_size=batch_size,
wait=True,
parallel=parallel,
)
assert remote_client.count(COLLECTION_NAME).count == len(vectors)
assert np.allclose(
remote_client.retrieve(COLLECTION_NAME, ids=[1], with_vectors=True)[0].vector["plain"],
norm_ref_vector,
)
remote_client.delete_collection(COLLECTION_NAME)
# endregion
@pytest.mark.skip
def test_upsert_batch_with_different_options():
bm25_name = "Qdrant/bm25"
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
download_options = dict()
download_options["cache_dir"] = _get_cache_dir_name("Qdrant/bm25")
if download_options["cache_dir"].exists():
download_options["local_files_only"] = True
sparse_doc_1 = models.Document(
text="running run", model=bm25_name, options={"language": "english", **download_options}
)
sparse_doc_2 = models.Document(
text="running run", model=bm25_name, options={"language": "german", **download_options}
)
sparse_doc_3 = models.Document(
text="running run", model=bm25_name, options={"language": "english", **download_options}
)
sparse_doc_4 = models.Document(
text="running run", model=bm25_name, options={"language": "german", **download_options}
)
sparse_vectors_config = {
"sparse-text-en": models.SparseVectorParams(modifier=models.Modifier.IDF),
"sparse-text-de": models.SparseVectorParams(modifier=models.Modifier.IDF),
}
local_client.create_collection(
COLLECTION_NAME, vectors_config={}, sparse_vectors_config=sparse_vectors_config
)
points = [
models.PointStruct(
id=0,
vector={
"sparse-text-en": sparse_doc_1,
"sparse-text-de": sparse_doc_2,
**download_options,
},
),
models.PointStruct(id=1, vector={"sparse-text-en": sparse_doc_3}),
models.PointStruct(id=2, vector={"sparse-text-de": sparse_doc_4}),
]
local_client.upsert(COLLECTION_NAME, points)
read_points, _ = local_client.scroll(COLLECTION_NAME, limit=4, with_vectors=True)
assert len(read_points) == 3
assert (
read_points[0].vector["sparse-text-en"].indices
!= read_points[0].vector["sparse-text-de"].indices
)
assert (
read_points[0].vector["sparse-text-en"].indices
== read_points[1].vector["sparse-text-en"].indices
)
assert (
read_points[0].vector["sparse-text-de"].indices
== read_points[2].vector["sparse-text-de"].indices
)
@pytest.mark.skip
def test_batch_size_propagation():
def mock(func, kw_param_storage):
def decorated(*args, **kwargs):
for k in kwargs:
kw_param_storage[k] = kwargs[k]
return func(*args, **kwargs)
return decorated
param_storage = {}
download_options = dict()
download_options["cache_dir"] = _get_cache_dir_name("Qdrant/bm25")
if download_options["cache_dir"].exists():
download_options["local_files_only"] = True
bm25_name = "Qdrant/bm25"
inference_batch_size = 2
local_client = QdrantClient(":memory:", local_inference_batch_size=inference_batch_size)
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
local_client._model_embedder.embedder.embed = mock(
local_client._model_embedder.embedder.embed, param_storage
)
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
sparse_doc_1 = models.Document(text="a quick", model=bm25_name, options=download_options)
sparse_doc_2 = models.Document(text="brown fox", model=bm25_name, options=download_options)
sparse_doc_3 = models.Document(text="jumps over", model=bm25_name, options=download_options)
sparse_doc_4 = models.Document(text="a lazy dog", model=bm25_name, options=download_options)
local_client.create_collection(
COLLECTION_NAME, sparse_vectors_config={"sparse": models.SparseVectorParams()}
)
points = [
models.PointStruct(id=0, vector={"sparse": sparse_doc_1}),
models.PointStruct(id=1, vector={"sparse": sparse_doc_2}),
models.PointStruct(id=2, vector={"sparse": sparse_doc_3}),
models.PointStruct(id=3, vector={"sparse": sparse_doc_4}),
]
local_client.upsert( # uses _embed_models
COLLECTION_NAME, points
)
assert param_storage["batch_size"] == inference_batch_size
param_storage.clear()
local_client.upload_points(
COLLECTION_NAME, points=points, batch_size=3
) # uses _embed_models_strict
assert param_storage["batch_size"] == inference_batch_size
param_storage.clear()
@pytest.fixture
def mock_late_interaction_multimodal_embedding():
try:
import fastembed
except ImportError:
pytest.skip("FastEmbed is not installed, skipping")
class LateInteractionMultimodalEmbeddingMock:
text_calls = 0
image_calls = 0
num_image_columns = 10
num_text_columns = 20
text_embedding = np.array([[0.1, 0.2, 0.3]] * num_text_columns)
image_embedding = np.array([[0.1, 0.2, 0.3]] * num_image_columns)
def __init__(self, *args, **kwargs):
pass
def embed_text(self, documents, *args, **kwargs):
LateInteractionMultimodalEmbeddingMock.text_calls += 1
for _ in documents:
yield self.text_embedding
def embed_image(self, images, *args, **kwargs):
LateInteractionMultimodalEmbeddingMock.image_calls += 1
for _ in images:
yield self.image_embedding
original_class = qdrant_client.embed.embedder.LateInteractionMultimodalEmbedding
qdrant_client.embed.embedder.LateInteractionMultimodalEmbedding = (
LateInteractionMultimodalEmbeddingMock
)
yield LateInteractionMultimodalEmbeddingMock
qdrant_client.embed.embedder.LateInteractionMultimodalEmbedding = original_class
@pytest.mark.skip
def test_embed_multimodal(mock_late_interaction_multimodal_embedding):
mock_cls = mock_late_interaction_multimodal_embedding
local_client = QdrantClient(":memory:")
if not local_client._FASTEMBED_INSTALLED:
pytest.skip("FastEmbed is not installed, skipping")
point_1 = models.PointStruct(
id=1,
vector={
"text": models.Document(text="a quick brown fox", model="Qdrant/colpali-v1.3-fp16"),
"image": models.Image(image=TEST_IMAGE_PATH, model="Qdrant/colpali-v1.3-fp16"),
},
)
point_2 = models.PointStruct(
id=2,
vector={
"text": models.Document(
text="jumped over a lazy dog", model="Qdrant/colpali-v1.3-fp16"
),
},
)
dim = 3
local_client.create_collection(
COLLECTION_NAME,
vectors_config={
"text": models.VectorParams(
size=dim,
distance=models.Distance.MANHATTAN,
multivector_config=models.MultiVectorConfig(
comparator=models.MultiVectorComparator.MAX_SIM
),
),
"image": models.VectorParams(
size=dim,
distance=models.Distance.MANHATTAN,
multivector_config=models.MultiVectorConfig(
comparator=models.MultiVectorComparator.MAX_SIM
),
),
},
)
local_client.upsert(COLLECTION_NAME, [point_1, point_2])
assert mock_cls.text_calls == 1 # assert that text records were assembled into batches
assert mock_cls.image_calls == 1
records, _ = local_client.scroll(COLLECTION_NAME, limit=2, with_vectors=True)
assert len(records) == 2
assert len(records[1].vector) == 1
np_text_vectors = np.array([records[0].vector["text"], records[1].vector["text"]])
np_image_vector = np.array(records[0].vector["image"])
assert all(vector.shape == (mock_cls.num_text_columns, dim) for vector in np_text_vectors)
assert np_image_vector.shape == (mock_cls.num_image_columns, dim)
# check that embeddings order was right
assert np.allclose(np_text_vectors[0], mock_cls.text_embedding)
assert np.allclose(np_text_vectors[1], mock_cls.text_embedding)
assert np.allclose(np_image_vector, mock_cls.image_embedding)
Reference in New Issue View Git Blame Copy Permalink