chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,653 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import ast
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import MutableMapping, Sequence
|
||||
from copy import deepcopy
|
||||
from typing import Any, ClassVar, Final, Generic, TypeVar
|
||||
|
||||
from pydantic import HttpUrl, SecretStr, ValidationError, model_validator
|
||||
from qdrant_client.async_qdrant_client import AsyncQdrantClient
|
||||
from qdrant_client.models import (
|
||||
Datatype,
|
||||
Distance,
|
||||
FieldCondition,
|
||||
Filter,
|
||||
Fusion,
|
||||
FusionQuery,
|
||||
MatchAny,
|
||||
MatchValue,
|
||||
PointStruct,
|
||||
Prefetch,
|
||||
QueryResponse,
|
||||
Range,
|
||||
ScoredPoint,
|
||||
VectorParams,
|
||||
)
|
||||
from typing_extensions import override
|
||||
|
||||
from semantic_kernel.connectors.ai.embedding_generator_base import EmbeddingGeneratorBase
|
||||
from semantic_kernel.data.vector import (
|
||||
DistanceFunction,
|
||||
GetFilteredRecordOptions,
|
||||
IndexKind,
|
||||
KernelSearchResults,
|
||||
SearchType,
|
||||
TModel,
|
||||
VectorSearch,
|
||||
VectorSearchOptions,
|
||||
VectorSearchResult,
|
||||
VectorStore,
|
||||
VectorStoreCollection,
|
||||
VectorStoreCollectionDefinition,
|
||||
)
|
||||
from semantic_kernel.exceptions import (
|
||||
VectorSearchExecutionException,
|
||||
VectorStoreInitializationException,
|
||||
VectorStoreModelValidationError,
|
||||
VectorStoreOperationException,
|
||||
)
|
||||
from semantic_kernel.kernel_pydantic import KernelBaseSettings
|
||||
from semantic_kernel.kernel_types import OneOrMany
|
||||
from semantic_kernel.utils.feature_stage_decorator import release_candidate
|
||||
from semantic_kernel.utils.telemetry.user_agent import APP_INFO, prepend_semantic_kernel_to_user_agent
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import override # pragma: no cover
|
||||
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
TKey = TypeVar("TKey", bound=str | int)
|
||||
|
||||
DISTANCE_FUNCTION_MAP: Final[dict[DistanceFunction, Distance]] = {
|
||||
DistanceFunction.COSINE_SIMILARITY: Distance.COSINE,
|
||||
DistanceFunction.DOT_PROD: Distance.DOT,
|
||||
DistanceFunction.EUCLIDEAN_DISTANCE: Distance.EUCLID,
|
||||
DistanceFunction.MANHATTAN: Distance.MANHATTAN,
|
||||
DistanceFunction.DEFAULT: Distance.COSINE,
|
||||
}
|
||||
INDEX_KIND_MAP: Final[dict[IndexKind, str]] = {
|
||||
IndexKind.HNSW: "hnsw",
|
||||
IndexKind.DEFAULT: "hnsw",
|
||||
}
|
||||
TYPE_MAPPER_VECTOR: Final[dict[str, Datatype]] = {
|
||||
"float": Datatype.FLOAT32,
|
||||
"int": Datatype.UINT8,
|
||||
"binary": Datatype.UINT8,
|
||||
"default": Datatype.FLOAT32,
|
||||
}
|
||||
IN_MEMORY_STRING: Final[str] = ":memory:"
|
||||
|
||||
|
||||
@release_candidate
|
||||
class QdrantSettings(KernelBaseSettings):
|
||||
"""Qdrant settings currently used by the Qdrant Vector Record Store."""
|
||||
|
||||
env_prefix: ClassVar[str] = "QDRANT_"
|
||||
|
||||
url: HttpUrl | None = None
|
||||
api_key: SecretStr | None = None
|
||||
host: str | None = None
|
||||
port: int | None = None
|
||||
grpc_port: int | None = None
|
||||
path: str | None = None
|
||||
location: str | None = None
|
||||
prefer_grpc: bool = False
|
||||
|
||||
@model_validator(mode="before")
|
||||
def validate_settings(cls, values: dict):
|
||||
"""Validate the settings."""
|
||||
if (
|
||||
isinstance(values, dict)
|
||||
and "url" not in values
|
||||
and "host" not in values
|
||||
and "path" not in values
|
||||
and "location" not in values
|
||||
):
|
||||
values["location"] = IN_MEMORY_STRING
|
||||
return values
|
||||
|
||||
def model_dump(self, **kwargs):
|
||||
"""Dump the model."""
|
||||
dump = super().model_dump(**kwargs)
|
||||
if "api_key" in dump:
|
||||
dump["api_key"] = dump["api_key"].get_secret_value()
|
||||
if "url" in dump:
|
||||
dump["url"] = str(dump["url"])
|
||||
return dump
|
||||
|
||||
|
||||
@release_candidate
|
||||
class QdrantCollection(
|
||||
VectorStoreCollection[TKey, TModel],
|
||||
VectorSearch[TKey, TModel],
|
||||
Generic[TKey, TModel],
|
||||
):
|
||||
"""A QdrantCollection is a memory collection that uses Qdrant as the backend."""
|
||||
|
||||
qdrant_client: AsyncQdrantClient
|
||||
named_vectors: bool
|
||||
supported_key_types: ClassVar[set[str] | None] = {"str", "int"}
|
||||
supported_vector_types: ClassVar[set[str] | None] = {"float", "int"}
|
||||
supported_search_types: ClassVar[set[SearchType]] = {SearchType.VECTOR, SearchType.KEYWORD_HYBRID}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
record_type: type[TModel],
|
||||
definition: VectorStoreCollectionDefinition | None = None,
|
||||
collection_name: str | None = None,
|
||||
embedding_generator: EmbeddingGeneratorBase | None = None,
|
||||
named_vectors: bool = True,
|
||||
url: str | None = None,
|
||||
api_key: str | None = None,
|
||||
host: str | None = None,
|
||||
port: int | None = None,
|
||||
grpc_port: int | None = None,
|
||||
path: str | None = None,
|
||||
location: str | None = None,
|
||||
prefer_grpc: bool | None = None,
|
||||
client: AsyncQdrantClient | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initializes a new instance of the QdrantVectorRecordStore.
|
||||
|
||||
When using qdrant client, make sure to supply url and api_key.
|
||||
When using qdrant server, make sure to supply url or host and optionally port.
|
||||
When using qdrant local, either supply path to use a persisted qdrant instance
|
||||
or set location to ":memory:" to use an in-memory qdrant instance.
|
||||
When nothing is supplied, it defaults to an in-memory qdrant instance.
|
||||
You can also supply a async qdrant client directly.
|
||||
|
||||
Args:
|
||||
record_type (type[TModel]): The type of the data model.
|
||||
definition (VectorStoreRecordDefinition): The model fields, optional.
|
||||
collection_name (str): The name of the collection, optional.
|
||||
embedding_generator (EmbeddingGeneratorBase): The embedding generator to use, optional.
|
||||
named_vectors (bool): If true, vectors are stored with name (default: True).
|
||||
url (str): The URL of the Qdrant server (default: {None}).
|
||||
api_key (str): The API key for the Qdrant server (default: {None}).
|
||||
host (str): The host of the Qdrant server (default: {None}).
|
||||
port (int): The port of the Qdrant server (default: {None}).
|
||||
grpc_port (int): The gRPC port of the Qdrant server (default: {None}).
|
||||
path (str): The path of the Qdrant server (default: {None}).
|
||||
location (str): The location of the Qdrant server (default: {None}).
|
||||
prefer_grpc (bool): If true, gRPC will be preferred (default: {None}).
|
||||
client (AsyncQdrantClient): The Qdrant client to use (default: {None}).
|
||||
env_file_path (str): Use the environment settings file as a fallback to environment variables.
|
||||
env_file_encoding (str): The encoding of the environment settings file.
|
||||
**kwargs: Additional keyword arguments passed to the client constructor.
|
||||
|
||||
"""
|
||||
if client:
|
||||
super().__init__(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
qdrant_client=client, # type: ignore
|
||||
named_vectors=named_vectors, # type: ignore
|
||||
managed_client=False,
|
||||
embedding_generator=embedding_generator,
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
settings = QdrantSettings(
|
||||
url=url,
|
||||
api_key=api_key,
|
||||
host=host,
|
||||
port=port,
|
||||
grpc_port=grpc_port,
|
||||
path=path,
|
||||
location=location,
|
||||
prefer_grpc=prefer_grpc,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
except ValidationError as ex:
|
||||
raise VectorStoreInitializationException("Failed to create Qdrant settings.", ex) from ex
|
||||
if APP_INFO:
|
||||
kwargs.setdefault("metadata", {})
|
||||
kwargs["metadata"] = prepend_semantic_kernel_to_user_agent(kwargs["metadata"])
|
||||
try:
|
||||
client = AsyncQdrantClient(**settings.model_dump(exclude_none=True), **kwargs)
|
||||
except ValueError as ex:
|
||||
raise VectorStoreInitializationException("Failed to create Qdrant client.", ex) from ex
|
||||
super().__init__(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
qdrant_client=client,
|
||||
named_vectors=named_vectors,
|
||||
embedding_generator=embedding_generator,
|
||||
)
|
||||
|
||||
@override
|
||||
async def _inner_upsert(
|
||||
self,
|
||||
records: Sequence[PointStruct],
|
||||
**kwargs: Any,
|
||||
) -> Sequence[TKey]:
|
||||
await self.qdrant_client.upsert(
|
||||
collection_name=self.collection_name,
|
||||
points=records,
|
||||
**kwargs,
|
||||
)
|
||||
return [record.id for record in records] # type: ignore
|
||||
|
||||
@override
|
||||
async def _inner_get(
|
||||
self,
|
||||
keys: Sequence[TKey] | None = None,
|
||||
options: GetFilteredRecordOptions | None = None,
|
||||
**kwargs: Any,
|
||||
) -> OneOrMany[Any] | None:
|
||||
if not keys:
|
||||
if options is not None:
|
||||
raise NotImplementedError("Get without keys is not yet implemented.")
|
||||
return None
|
||||
if "with_vectors" not in kwargs:
|
||||
kwargs["with_vectors"] = kwargs.pop("include_vectors", True)
|
||||
return await self.qdrant_client.retrieve(
|
||||
collection_name=self.collection_name,
|
||||
ids=keys,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@override
|
||||
async def _inner_delete(self, keys: Sequence[TKey], **kwargs: Any) -> None:
|
||||
await self.qdrant_client.delete(
|
||||
collection_name=self.collection_name,
|
||||
points_selector=keys,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@override
|
||||
async def _inner_search(
|
||||
self,
|
||||
search_type: SearchType,
|
||||
options: VectorSearchOptions,
|
||||
values: Any | None = None,
|
||||
vector: Sequence[float | int] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> KernelSearchResults[VectorSearchResult[TModel]]:
|
||||
query_vector: tuple[str, Sequence[float | int]] | Sequence[float | int] | None = None
|
||||
|
||||
if not vector:
|
||||
vector = await self._generate_vector_from_values(values, options)
|
||||
|
||||
if not vector:
|
||||
raise VectorSearchExecutionException("Search requires a vector.")
|
||||
|
||||
vector_field = self.definition.try_get_vector_field(options.vector_property_name)
|
||||
if not vector_field:
|
||||
raise VectorStoreOperationException(
|
||||
f"Vector field {options.vector_property_name} not found in data model definition."
|
||||
)
|
||||
query_vector = (vector_field.storage_name or vector_field.name, vector) if self.named_vectors else vector
|
||||
filters: Filter | list[Filter] | None = self._build_filter(options.filter) # type: ignore
|
||||
filter: Filter | None = Filter(must=filters) if filters and isinstance(filters, list) else filters # type: ignore
|
||||
if search_type == SearchType.VECTOR:
|
||||
results = await self.qdrant_client.search(
|
||||
collection_name=self.collection_name,
|
||||
query_vector=query_vector, # type: ignore
|
||||
query_filter=filter,
|
||||
with_vectors=options.include_vectors,
|
||||
limit=options.top,
|
||||
offset=options.skip,
|
||||
**kwargs,
|
||||
)
|
||||
else:
|
||||
# Hybrid search: vector + keywords (RRF fusion)
|
||||
# 1. Get keywords and text field
|
||||
if not values:
|
||||
raise VectorSearchExecutionException("Hybrid search requires non-empty keywords in values.")
|
||||
if not options.additional_property_name:
|
||||
raise VectorSearchExecutionException("Hybrid search requires a keyword field name.")
|
||||
text_field = next(
|
||||
field
|
||||
for field in self.definition.fields
|
||||
if field.name == options.additional_property_name
|
||||
or field.storage_name == options.additional_property_name
|
||||
)
|
||||
if not text_field:
|
||||
raise VectorStoreOperationException(
|
||||
f"Keyword field {options.additional_property_name} not found in data model definition."
|
||||
)
|
||||
keyword_filter = deepcopy(filter) if filter else Filter()
|
||||
keyword_sub_filter = Filter(
|
||||
should=[
|
||||
FieldCondition(key=text_field.storage_name or text_field.name, match=MatchAny(any=[kw]))
|
||||
for kw in values
|
||||
]
|
||||
)
|
||||
if isinstance(keyword_filter.must, list):
|
||||
keyword_filter.must.append(keyword_sub_filter)
|
||||
elif isinstance(keyword_filter.must, Filter):
|
||||
keyword_filter.must = Filter(must=[keyword_filter.must, keyword_sub_filter])
|
||||
else:
|
||||
keyword_filter.must = keyword_sub_filter
|
||||
|
||||
points = await self.qdrant_client.query_points(
|
||||
collection_name=self.collection_name,
|
||||
prefetch=[
|
||||
Prefetch(
|
||||
query=vector, # type: ignore
|
||||
using=vector_field.storage_name or vector_field.name,
|
||||
filter=filter,
|
||||
limit=options.top,
|
||||
),
|
||||
Prefetch(filter=keyword_filter),
|
||||
],
|
||||
query=FusionQuery(fusion=Fusion.RRF),
|
||||
limit=options.top,
|
||||
offset=options.skip,
|
||||
with_vectors=options.include_vectors,
|
||||
**kwargs,
|
||||
)
|
||||
results = points.points
|
||||
|
||||
return KernelSearchResults(
|
||||
results=self._get_vector_search_results_from_results(results, options),
|
||||
total_count=len(results) if options.include_total_count else None,
|
||||
)
|
||||
|
||||
@override
|
||||
def _get_record_from_result(self, result: ScoredPoint | QueryResponse) -> Any:
|
||||
return result
|
||||
|
||||
@override
|
||||
def _get_score_from_result(self, result: ScoredPoint) -> float:
|
||||
return result.score
|
||||
|
||||
@override
|
||||
def _lambda_parser(self, node: ast.AST) -> Any:
|
||||
# Qdrant filter translation: output a qdrant_client.models.Filter or FieldCondition tree
|
||||
# Use correct Match subtypes: MatchAny, MatchValue, etc.
|
||||
# See: https://python-client.qdrant.tech/qdrant_client.http.models.models#qdrant_client.http.models.models.Filter
|
||||
match node:
|
||||
case ast.Compare():
|
||||
if len(node.ops) > 1:
|
||||
# Chain comparisons (e.g., 1 < x < 3) become AND of each comparison
|
||||
conditions = []
|
||||
for idx in range(len(node.ops)):
|
||||
left = node.left if idx == 0 else node.comparators[idx - 1]
|
||||
right = node.comparators[idx]
|
||||
op = node.ops[idx]
|
||||
conditions.append(self._lambda_parser(ast.Compare(left=left, ops=[op], comparators=[right])))
|
||||
return Filter(must=conditions)
|
||||
left = self._lambda_parser(node.left)
|
||||
right = self._lambda_parser(node.comparators[0])
|
||||
op = node.ops[0]
|
||||
match op:
|
||||
case ast.In():
|
||||
# IN: left in right (right is a list)
|
||||
return FieldCondition(key=left, match=MatchAny(any=right))
|
||||
case ast.NotIn():
|
||||
# NOT IN: left not in right
|
||||
return Filter(must_not=[FieldCondition(key=left, match=MatchAny(any=right))])
|
||||
case ast.Eq():
|
||||
return FieldCondition(key=left, match=MatchValue(value=right))
|
||||
case ast.NotEq():
|
||||
return Filter(must_not=[FieldCondition(key=left, match=MatchValue(value=right))])
|
||||
case ast.Gt():
|
||||
return FieldCondition(key=left, range=Range(gt=right))
|
||||
case ast.GtE():
|
||||
return FieldCondition(key=left, range=Range(gte=right))
|
||||
case ast.Lt():
|
||||
return FieldCondition(key=left, range=Range(lt=right))
|
||||
case ast.LtE():
|
||||
return FieldCondition(key=left, range=Range(lte=right))
|
||||
raise NotImplementedError(f"Unsupported operator: {type(op)}")
|
||||
case ast.BoolOp():
|
||||
op = node.op # type: ignore
|
||||
values = [self._lambda_parser(v) for v in node.values]
|
||||
if isinstance(op, ast.And):
|
||||
return Filter(must=values)
|
||||
if isinstance(op, ast.Or):
|
||||
return Filter(should=values)
|
||||
raise NotImplementedError(f"Unsupported BoolOp: {type(op)}")
|
||||
case ast.UnaryOp():
|
||||
match node.op:
|
||||
case ast.Not():
|
||||
operand = self._lambda_parser(node.operand)
|
||||
return Filter(must_not=[operand])
|
||||
case ast.UAdd() | ast.USub() | ast.Invert():
|
||||
raise NotImplementedError("Unary +, -, ~ are not supported in Qdrant filters.")
|
||||
case ast.Attribute():
|
||||
# Only allow attributes that are in the data model
|
||||
if node.attr not in self.definition.storage_names:
|
||||
raise VectorStoreOperationException(
|
||||
f"Field '{node.attr}' not in data model (storage property names are used)."
|
||||
)
|
||||
return node.attr
|
||||
case ast.Name():
|
||||
# Only allow names that are in the data model
|
||||
if node.id not in self.definition.storage_names:
|
||||
raise VectorStoreOperationException(
|
||||
f"Field '{node.id}' not in data model (storage property names are used)."
|
||||
)
|
||||
return node.id
|
||||
case ast.Constant():
|
||||
return node.value
|
||||
case ast.List():
|
||||
return [self._lambda_parser(elt) for elt in node.elts]
|
||||
raise NotImplementedError(f"Unsupported AST node: {type(node)}")
|
||||
|
||||
@override
|
||||
def _serialize_dicts_to_store_models(
|
||||
self,
|
||||
records: Sequence[dict[str, Any]],
|
||||
**kwargs: Any,
|
||||
) -> Sequence[PointStruct]:
|
||||
return [
|
||||
PointStruct(
|
||||
id=record.pop(self._key_field_name),
|
||||
vector=record.pop(self.definition.vector_field_names[0])
|
||||
if not self.named_vectors
|
||||
else {
|
||||
field.storage_name or field.name: record.pop(field.name) for field in self.definition.vector_fields
|
||||
},
|
||||
payload=record,
|
||||
)
|
||||
for record in records
|
||||
]
|
||||
|
||||
@override
|
||||
def _deserialize_store_models_to_dicts(
|
||||
self,
|
||||
records: Sequence[PointStruct] | Sequence[ScoredPoint],
|
||||
**kwargs: Any,
|
||||
) -> Sequence[dict[str, Any]]:
|
||||
return [
|
||||
{
|
||||
self._key_field_name: record.id,
|
||||
**(record.payload if record.payload else {}),
|
||||
**(
|
||||
{}
|
||||
if not record.vector
|
||||
else record.vector
|
||||
if isinstance(record.vector, dict)
|
||||
else {self.definition.vector_field_names[0]: record.vector}
|
||||
),
|
||||
}
|
||||
for record in records
|
||||
]
|
||||
|
||||
@override
|
||||
async def ensure_collection_exists(self, **kwargs) -> None:
|
||||
"""Create a new collection in Qdrant.
|
||||
|
||||
Args:
|
||||
**kwargs: Additional keyword arguments.
|
||||
You can supply all keyword arguments supported by the QdrantClient.create_collection method.
|
||||
This method creates the vectors_config automatically when not supplied, other params are not set.
|
||||
Collection name will be set to the collection_name property, cannot be overridden.
|
||||
"""
|
||||
if "vectors_config" not in kwargs:
|
||||
if self.named_vectors:
|
||||
vectors_config: MutableMapping[str, VectorParams] = {}
|
||||
for field in self.definition.vector_fields:
|
||||
if field.index_kind not in INDEX_KIND_MAP:
|
||||
raise VectorStoreOperationException(f"Index kind {field.index_kind} is not supported.")
|
||||
if field.distance_function not in DISTANCE_FUNCTION_MAP:
|
||||
raise VectorStoreOperationException(
|
||||
f"Distance function {field.distance_function} is not supported."
|
||||
)
|
||||
vectors_config[field.storage_name or field.name] = VectorParams(
|
||||
size=field.dimensions,
|
||||
distance=DISTANCE_FUNCTION_MAP[field.distance_function],
|
||||
datatype=TYPE_MAPPER_VECTOR[field.type_ or "default"],
|
||||
)
|
||||
kwargs["vectors_config"] = vectors_config
|
||||
else:
|
||||
vector = self.definition.try_get_vector_field(None)
|
||||
if not vector:
|
||||
raise VectorStoreOperationException("Vector field not found in data model definition.")
|
||||
if vector.distance_function not in DISTANCE_FUNCTION_MAP:
|
||||
raise VectorStoreOperationException(
|
||||
f"Distance function {vector.distance_function} is not supported."
|
||||
)
|
||||
kwargs["vectors_config"] = VectorParams(
|
||||
size=vector.dimensions,
|
||||
distance=DISTANCE_FUNCTION_MAP[vector.distance_function],
|
||||
datatype=TYPE_MAPPER_VECTOR[vector.type_ or "default"],
|
||||
)
|
||||
if "collection_name" not in kwargs:
|
||||
kwargs["collection_name"] = self.collection_name
|
||||
await self.qdrant_client.create_collection(**kwargs)
|
||||
|
||||
@override
|
||||
async def collection_exists(self, **kwargs) -> bool:
|
||||
return await self.qdrant_client.collection_exists(self.collection_name, **kwargs)
|
||||
|
||||
@override
|
||||
async def ensure_collection_deleted(self, **kwargs) -> None:
|
||||
await self.qdrant_client.delete_collection(self.collection_name, **kwargs)
|
||||
|
||||
def _validate_data_model(self):
|
||||
"""Internal function that should be overloaded by child classes to validate datatypes, etc.
|
||||
|
||||
This should take the VectorStoreRecordDefinition from the item_type and validate it against the store.
|
||||
|
||||
Checks should include, allowed naming of parameters, allowed data types, allowed vector dimensions.
|
||||
"""
|
||||
super()._validate_data_model()
|
||||
if len(self.definition.vector_field_names) > 1 and not self.named_vectors:
|
||||
raise VectorStoreModelValidationError("Only one vector field is allowed when not using named vectors.")
|
||||
|
||||
@override
|
||||
async def __aexit__(self, exc_type, exc_value, traceback) -> None:
|
||||
"""Exit the context manager."""
|
||||
if self.managed_client:
|
||||
await self.qdrant_client.close()
|
||||
|
||||
|
||||
@release_candidate
|
||||
class QdrantStore(VectorStore):
|
||||
"""A QdrantStore is a memory store that uses Qdrant as the backend."""
|
||||
|
||||
qdrant_client: AsyncQdrantClient
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str | None = None,
|
||||
api_key: str | None = None,
|
||||
host: str | None = None,
|
||||
port: int | None = None,
|
||||
grpc_port: int | None = None,
|
||||
path: str | None = None,
|
||||
location: str | None = None,
|
||||
prefer_grpc: bool | None = None,
|
||||
client: AsyncQdrantClient | None = None,
|
||||
embedding_generator: EmbeddingGeneratorBase | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initializes a new instance of the QdrantVectorRecordStore.
|
||||
|
||||
When using qdrant client, make sure to supply url and api_key.
|
||||
When using qdrant server, make sure to supply url or host and optionally port.
|
||||
When using qdrant local, either supply path to use a persisted qdrant instance
|
||||
or set location to ":memory:" to use an in-memory qdrant instance.
|
||||
When nothing is supplied, it defaults to an in-memory qdrant instance.
|
||||
You can also supply a async qdrant client directly.
|
||||
|
||||
Args:
|
||||
url: The URL of the Qdrant server (default: {None}).
|
||||
api_key: The API key for the Qdrant server (default: {None}).
|
||||
host: The host of the Qdrant server (default: {None}).
|
||||
port: The port of the Qdrant server (default: {None}).
|
||||
grpc_port: The gRPC port of the Qdrant server (default: {None}).
|
||||
path: The path of the Qdrant server (default: {None}).
|
||||
location: The location of the Qdrant server (default: {None}).
|
||||
prefer_grpc: If true, gRPC will be preferred (default: {None}).
|
||||
client: The Qdrant client to use (default: {None}).
|
||||
embedding_generator: The embedding generator to use (default: {None}).
|
||||
env_file_path: Use the environment settings file as a fallback to environment variables.
|
||||
env_file_encoding: The encoding of the environment settings file.
|
||||
**kwargs: Additional keyword arguments passed to the client constructor.
|
||||
|
||||
"""
|
||||
if client:
|
||||
super().__init__(
|
||||
qdrant_client=client, managed_client=False, embedding_generator=embedding_generator, **kwargs
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
settings = QdrantSettings(
|
||||
url=url,
|
||||
api_key=api_key,
|
||||
host=host,
|
||||
port=port,
|
||||
grpc_port=grpc_port,
|
||||
path=path,
|
||||
location=location,
|
||||
prefer_grpc=prefer_grpc,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
except ValidationError as ex:
|
||||
raise VectorStoreInitializationException("Failed to create Qdrant settings.", ex) from ex
|
||||
if APP_INFO:
|
||||
kwargs.setdefault("metadata", {})
|
||||
kwargs["metadata"] = prepend_semantic_kernel_to_user_agent(kwargs["metadata"])
|
||||
try:
|
||||
client = AsyncQdrantClient(**settings.model_dump(exclude_none=True), **kwargs)
|
||||
except ValueError as ex:
|
||||
raise VectorStoreInitializationException("Failed to create Qdrant client.", ex) from ex
|
||||
super().__init__(qdrant_client=client, embedding_generator=embedding_generator, **kwargs)
|
||||
|
||||
@override
|
||||
def get_collection(
|
||||
self,
|
||||
record_type: type[TModel],
|
||||
*,
|
||||
definition: VectorStoreCollectionDefinition | None = None,
|
||||
collection_name: str | None = None,
|
||||
embedding_generator: EmbeddingGeneratorBase | None = None,
|
||||
**kwargs: Any,
|
||||
) -> QdrantCollection:
|
||||
return QdrantCollection(
|
||||
record_type=record_type,
|
||||
definition=definition,
|
||||
collection_name=collection_name,
|
||||
client=self.qdrant_client,
|
||||
embedding_generator=embedding_generator or self.embedding_generator,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@override
|
||||
async def list_collection_names(self, **kwargs: Any) -> Sequence[str]:
|
||||
collections = await self.qdrant_client.get_collections()
|
||||
return [collection.name for collection in collections.collections]
|
||||
|
||||
@override
|
||||
async def __aexit__(self, exc_type, exc_value, traceback) -> None:
|
||||
if self.managed_client:
|
||||
await self.qdrant_client.close()
|
||||
Reference in New Issue
Block a user