Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix qdrant client: compatibe with open-source qdrant #308

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 19 additions & 6 deletions vectordb_bench/backend/clients/qdrant_cloud/config.py
Original file line number Diff line number Diff line change
@@ -1,18 +1,31 @@
from pydantic import BaseModel, SecretStr

from ..api import DBConfig, DBCaseConfig, MetricType
from pydantic import validator


# Allowing `api_key` to be left empty, to ensure compatibility with the open-source Qdrant.
class QdrantConfig(DBConfig):
url: SecretStr
api_key: SecretStr

def to_dict(self) -> dict:
return {
"url": self.url.get_secret_value(),
"api_key": self.api_key.get_secret_value(),
"prefer_grpc": True,
}
api_key = self.api_key.get_secret_value()
if len(api_key) > 0:
return {
"url": self.url.get_secret_value(),
"api_key": self.api_key.get_secret_value(),
"prefer_grpc": True,
}
else:
return {"url": self.url.get_secret_value(),}

@validator("*")
def not_empty_field(cls, v, field):
if field.name in ["api_key", "db_label"]:
return v
if isinstance(v, (str, SecretStr)) and len(v) == 0:
raise ValueError("Empty string!")
return v

class QdrantIndexConfig(BaseModel, DBCaseConfig):
metric_type: MetricType | None = None
Expand Down
18 changes: 11 additions & 7 deletions vectordb_bench/backend/clients/qdrant_cloud/qdrant_cloud.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,7 @@ def __init__(
if drop_old:
log.info(f"QdrantCloud client drop_old collection: {self.collection_name}")
tmp_client.delete_collection(self.collection_name)

self._create_collection(dim, tmp_client)
self._create_collection(dim, tmp_client)
tmp_client = None

@contextmanager
Expand Down Expand Up @@ -110,13 +109,18 @@ def insert_embeddings(
) -> (int, Exception):
"""Insert embeddings into Milvus. should call self.init() first"""
assert self.qdrant_client is not None
QDRANT_BATCH_SIZE = 500
try:
# TODO: counts
_ = self.qdrant_client.upsert(
collection_name=self.collection_name,
wait=True,
points=Batch(ids=metadata, payloads=[{self._primary_field: v} for v in metadata], vectors=embeddings)
)
for offset in range(0, len(embeddings), QDRANT_BATCH_SIZE):
vectors = embeddings[offset: offset + QDRANT_BATCH_SIZE]
ids = metadata[offset: offset + QDRANT_BATCH_SIZE]
payloads=[{self._primary_field: v} for v in ids]
_ = self.qdrant_client.upsert(
collection_name=self.collection_name,
wait=True,
points=Batch(ids=ids, payloads=payloads, vectors=vectors),
)
except Exception as e:
log.info(f"Failed to insert data, {e}")
return 0, e
Expand Down