Skip to content

Commit

Permalink
fix qdrant client: compatibe with open-source qdrant
Browse files Browse the repository at this point in the history
Signed-off-by: min.tian <min.tian.cn@gmail.com>
  • Loading branch information
alwayslove2013 committed Apr 22, 2024
1 parent 2966832 commit 6e0ad2e
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 13 deletions.
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

0 comments on commit 6e0ad2e

Please sign in to comment.