semantic-kernel 1.8.3__zip

semantic_kernel-1.8.3/semantic_kernel/connectors/memory/qdrant/qdrant_collection.py

@@ -0,0 +1,257 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+import sys
+from collections.abc import Mapping, Sequence
+from typing import Any, ClassVar, TypeVar
+
+if sys.version_info >= (3, 12):
+    from typing import override  # pragma: no cover
+else:
+    from typing_extensions import override  # pragma: no cover
+
+from pydantic import ValidationError
+from qdrant_client.async_qdrant_client import AsyncQdrantClient
+from qdrant_client.models import PointStruct, VectorParams
+
+from semantic_kernel.connectors.memory.qdrant.const import DISTANCE_FUNCTION_MAP, TYPE_MAPPER_VECTOR
+from semantic_kernel.connectors.memory.qdrant.utils import AsyncQdrantClientWrapper
+from semantic_kernel.data.vector_store_model_definition import VectorStoreRecordDefinition
+from semantic_kernel.data.vector_store_record_collection import VectorStoreRecordCollection
+from semantic_kernel.data.vector_store_record_fields import VectorStoreRecordVectorField
+from semantic_kernel.exceptions import (
+    MemoryConnectorInitializationError,
+    VectorStoreModelValidationError,
+)
+from semantic_kernel.exceptions.memory_connector_exceptions import MemoryConnectorException
+from semantic_kernel.kernel_types import OneOrMany
+from semantic_kernel.utils.experimental_decorator import experimental_class
+from semantic_kernel.utils.telemetry.user_agent import APP_INFO, prepend_semantic_kernel_to_user_agent
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+TModel = TypeVar("TModel")
+TKey = TypeVar("TKey", str, int)
+
+
+@experimental_class
+class QdrantCollection(VectorStoreRecordCollection[str | int, TModel]):
+    """A QdrantCollection is a memory collection that uses Qdrant as the backend."""
+
+    qdrant_client: AsyncQdrantClient
+    named_vectors: bool
+    supported_key_types: ClassVar[list[str] | None] = ["str", "int"]
+    supported_vector_types: ClassVar[list[str] | None] = ["float", "int"]
+
+    def __init__(
+        self,
+        data_model_type: type[TModel],
+        data_model_definition: VectorStoreRecordDefinition | None = None,
+        collection_name: str | 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:
+            data_model_type (type[TModel]): The type of the data model.
+            data_model_definition (VectorStoreRecordDefinition): The model fields, optional.
+            collection_name (str): The name of the collection, 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__(
+                data_model_type=data_model_type,
+                data_model_definition=data_model_definition,
+                collection_name=collection_name,
+                qdrant_client=client,  # type: ignore
+                named_vectors=named_vectors,  # type: ignore
+            )
+            return
+
+        from semantic_kernel.connectors.memory.qdrant.qdrant_settings import QdrantSettings
+
+        try:
+            settings = QdrantSettings.create(
+                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 MemoryConnectorInitializationError("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 = AsyncQdrantClientWrapper(**settings.model_dump(exclude_none=True), **kwargs)
+        except ValueError as ex:
+            raise MemoryConnectorInitializationError("Failed to create Qdrant client.", ex) from ex
+        super().__init__(
+            data_model_type=data_model_type,
+            data_model_definition=data_model_definition,
+            collection_name=collection_name,
+            qdrant_client=client,
+            named_vectors=named_vectors,
+        )
+
+    @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]
+
+    @override
+    async def _inner_get(self, keys: Sequence[TKey], **kwargs: Any) -> OneOrMany[Any] | 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
+    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.data_model_definition.vector_field_names[0])
+                if not self.named_vectors
+                else {field: record.pop(field) for field in self.data_model_definition.vector_field_names},
+                payload=record,
+            )
+            for record in records
+        ]
+
+    @override
+    def _deserialize_store_models_to_dicts(
+        self,
+        records: Sequence[PointStruct],
+        **kwargs: Any,
+    ) -> Sequence[dict[str, Any]]:
+        return [
+            {
+                self._key_field_name: record.id,
+                **(record.payload if record.payload else {}),
+                **(
+                    record.vector
+                    if isinstance(record.vector, dict)
+                    else {self.data_model_definition.vector_field_names[0]: record.vector}
+                ),
+            }
+            for record in records
+        ]
+
+    @override
+    async def create_collection(self, **kwargs) -> None:
+        """Create a new collection in Azure AI Search.
+
+        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:
+            vectors_config: VectorParams | Mapping[str, VectorParams] = {}
+            if self.named_vectors:
+                for field in self.data_model_definition.vector_field_names:
+                    vector = self.data_model_definition.fields[field]
+                    assert isinstance(vector, VectorStoreRecordVectorField)  # nosec
+                    if not vector.dimensions:
+                        raise MemoryConnectorException("Vector field must have dimensions.")
+                    vectors_config[field] = VectorParams(
+                        size=vector.dimensions,
+                        distance=DISTANCE_FUNCTION_MAP[vector.distance_function or "default"],
+                        datatype=TYPE_MAPPER_VECTOR[vector.property_type or "default"],
+                    )
+            else:
+                vector = self.data_model_definition.fields[self.data_model_definition.vector_field_names[0]]
+                assert isinstance(vector, VectorStoreRecordVectorField)  # nosec
+                if not vector.dimensions:
+                    raise MemoryConnectorException("Vector field must have dimensions.")
+                vectors_config = VectorParams(
+                    size=vector.dimensions,
+                    distance=DISTANCE_FUNCTION_MAP[vector.distance_function or "default"],
+                    datatype=TYPE_MAPPER_VECTOR[vector.property_type or "default"],
+                )
+            kwargs["vectors_config"] = vectors_config
+        if "collection_name" not in kwargs:
+            kwargs["collection_name"] = self.collection_name
+        await self.qdrant_client.create_collection(**kwargs)
+
+    @override
+    async def does_collection_exist(self, **kwargs) -> bool:
+        return await self.qdrant_client.collection_exists(self.collection_name, **kwargs)
+
+    @override
+    async def delete_collection(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.data_model_definition.vector_field_names) > 1 and not self.named_vectors:
+            raise VectorStoreModelValidationError("Only one vector field is allowed when not using named vectors.")

semantic_kernel-1.8.3/semantic_kernel/connectors/memory/qdrant/qdrant_memory_store.py

@@ -0,0 +1,303 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import sys
+import uuid
+
+from numpy import ndarray
+from qdrant_client import QdrantClient
+from qdrant_client import models as qdrant_models
+
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
+from semantic_kernel.exceptions import ServiceResponseException
+from semantic_kernel.memory.memory_record import MemoryRecord
+from semantic_kernel.memory.memory_store_base import MemoryStoreBase
+from semantic_kernel.utils.experimental_decorator import experimental_class
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+@experimental_class
+class QdrantMemoryStore(MemoryStoreBase):
+    """QdrantMemoryStore."""
+
+    _qdrantclient: QdrantClient
+
+    def __init__(
+        self,
+        vector_size: int,
+        url: str | None = None,
+        port: int | None = 6333,
+        local: bool | None = False,
+        **kwargs,
+    ) -> None:
+        """Initializes a new instance of the QdrantMemoryStore class."""
+        if local:
+            if url:
+                self._qdrantclient = QdrantClient(location=url)
+            else:
+                self._qdrantclient = QdrantClient(location=":memory:")
+        else:
+            self._qdrantclient = QdrantClient(url=url, port=port)
+
+        self._default_vector_size = vector_size
+
+    @override
+    async def create_collection(self, collection_name: str) -> None:
+        self._qdrantclient.recreate_collection(
+            collection_name=collection_name,
+            vectors_config=qdrant_models.VectorParams(
+                size=self._default_vector_size, distance=qdrant_models.Distance.COSINE
+            ),
+        )
+
+    @override
+    async def get_collections(
+        self,
+    ) -> list[str]:
+        collection_info = self._qdrantclient.get_collections()
+        return [collection.name for collection in collection_info.collections]
+
+    async def get_collection(self, collection_name: str) -> qdrant_models.CollectionInfo:
+        """Gets the collection based upon collection name.
+
+        Returns:
+            CollectionInfo -- Collection Information from Qdrant about collection.
+        """
+        return self._qdrantclient.get_collection(collection_name=collection_name)
+
+    @override
+    async def delete_collection(self, collection_name: str) -> None:
+        self._qdrantclient.delete_collection(collection_name=collection_name)
+
+    @override
+    async def does_collection_exist(self, collection_name: str) -> bool:
+        return self._qdrantclient.collection_exists(collection_name=collection_name)
+
+    @override
+    async def upsert(self, collection_name: str, record: MemoryRecord) -> str:
+        data_to_upsert = await self._convert_from_memory_record(
+            collection_name=collection_name,
+            record=record,
+        )
+
+        result = self._qdrantclient.upsert(
+            collection_name=collection_name,
+            points=[data_to_upsert],
+        )
+
+        if result.status == qdrant_models.UpdateStatus.COMPLETED:
+            return data_to_upsert.id
+        raise ServiceResponseException("Upsert failed")
+
+    @override
+    async def upsert_batch(self, collection_name: str, records: list[MemoryRecord]) -> list[str]:
+        tasks = []
+        for record in records:
+            tasks.append(
+                self._convert_from_memory_record(
+                    collection_name=collection_name,
+                    record=record,
+                )
+            )
+
+        data_to_upsert = await asyncio.gather(*tasks)
+
+        result = self._qdrantclient.upsert(
+            collection_name=collection_name,
+            points=data_to_upsert,
+        )
+
+        if result.status == qdrant_models.UpdateStatus.COMPLETED:
+            return [data.id for data in data_to_upsert]
+        raise ServiceResponseException("Batch upsert failed")
+
+    @override
+    async def get(self, collection_name: str, key: str, with_embedding: bool = False) -> MemoryRecord | None:
+        result = await self._get_existing_record_by_payload_id(
+            collection_name=collection_name,
+            payload_id=key,
+            with_embedding=with_embedding,
+        )
+
+        if result:
+            return MemoryRecord(
+                is_reference=result.payload["_is_reference"],
+                external_source_name=result.payload["_external_source_name"],
+                id=result.payload["_id"],
+                description=result.payload["_description"],
+                text=result.payload["_text"],
+                additional_metadata=result.payload["_additional_metadata"],
+                embedding=result.vector,
+                key=result.id,
+                timestamp=result.payload["_timestamp"],
+            )
+        return None
+
+    @override
+    async def get_batch(
+        self,
+        collection_name: str,
+        keys: list[str],
+        with_embeddings: bool = False,
+    ) -> list[MemoryRecord]:
+        tasks = []
+        for key in keys:
+            tasks.append(
+                self.get(
+                    collection_name=collection_name,
+                    key=key,
+                    with_embedding=with_embeddings,
+                )
+            )
+        return await asyncio.gather(*tasks)
+
+    @override
+    async def remove(self, collection_name: str, key: str) -> None:
+        existing_record = await self._get_existing_record_by_payload_id(
+            collection_name=collection_name,
+            payload_id=key,
+            with_embedding=False,
+        )
+
+        if existing_record:
+            pointId = existing_record.id
+            result = self._qdrantclient.delete(collection_name=collection_name, points_selector=[pointId])
+            if result.status != qdrant_models.UpdateStatus.COMPLETED:
+                raise ServiceResponseException("Delete failed")
+
+    @override
+    async def remove_batch(self, collection_name: str, keys: list[str]) -> None:
+        tasks = []
+        for key in keys:
+            tasks.append(
+                self._get_existing_record_by_payload_id(
+                    collection_name=collection_name,
+                    payload_id=key,
+                    with_embedding=False,
+                )
+            )
+
+        existing_records = await asyncio.gather(*tasks)
+
+        if len(existing_records) > 0:
+            result = self._qdrantclient.delete(
+                collection_name=collection_name,
+                points_selector=[record.id for record in existing_records],
+            )
+            if result.status != qdrant_models.UpdateStatus.COMPLETED:
+                raise ServiceResponseException("Delete failed")
+
+    @override
+    async def get_nearest_matches(
+        self,
+        collection_name: str,
+        embedding: ndarray,
+        limit: int,
+        min_relevance_score: float,
+        with_embeddings: bool = False,
+    ) -> list[tuple[MemoryRecord, float]]:
+        match_results = self._qdrantclient.search(
+            collection_name=collection_name,
+            query_vector=embedding,
+            limit=limit,
+            score_threshold=min_relevance_score,
+            with_vectors=with_embeddings,
+        )
+
+        return [
+            (
+                MemoryRecord(
+                    is_reference=result.payload["_is_reference"],
+                    external_source_name=result.payload["_external_source_name"],
+                    id=result.payload["_id"],
+                    description=result.payload["_description"],
+                    text=result.payload["_text"],
+                    additional_metadata=result.payload["_additional_metadata"],
+                    embedding=result.vector,
+                    key=result.id,
+                    timestamp=result.payload["_timestamp"],
+                ),
+                result.score,
+            )
+            for result in match_results
+        ]
+
+    @override
+    async def get_nearest_match(
+        self,
+        collection_name: str,
+        embedding: ndarray,
+        min_relevance_score: float,
+        with_embedding: bool = False,
+    ) -> tuple[MemoryRecord, float]:
+        result = await self.get_nearest_matches(
+            collection_name=collection_name,
+            embedding=embedding,
+            limit=1,
+            min_relevance_score=min_relevance_score,
+            with_embeddings=with_embedding,
+        )
+        return result[0] if result else None
+
+    async def _get_existing_record_by_payload_id(
+        self,
+        collection_name: str,
+        payload_id: str,
+        with_embedding: bool = False,
+    ) -> qdrant_models.ScoredPoint | None:
+        """Gets an existing record based upon payload id.
+
+        Args:
+            collection_name (str): The name of the collection.
+            payload_id (str): The payload id to search for.
+            with_embedding (bool): If true, the embedding will be returned in the memory records.
+
+        Returns:
+            Optional[ScoredPoint]: The existing record if found; otherwise, None.
+        """
+        filter = qdrant_models.Filter(
+            must=[
+                qdrant_models.FieldCondition(
+                    key="_id",
+                    match=qdrant_models.MatchText(text=payload_id),
+                )
+            ]
+        )
+
+        existing_record = self._qdrantclient.search(
+            collection_name=collection_name,
+            query_vector=[0.0] * self._default_vector_size,
+            limit=1,
+            query_filter=filter,
+            with_payload=True,
+            with_vectors=with_embedding,
+        )
+
+        if existing_record:
+            return existing_record[0]
+        return None
+
+    async def _convert_from_memory_record(
+        self, collection_name: str, record: MemoryRecord
+    ) -> qdrant_models.PointStruct:
+        if record._key is not None and record._key != "":
+            pointId = record._key
+
+        else:
+            existing_record = await self._get_existing_record_by_payload_id(
+                collection_name=collection_name,
+                payload_id=record._id,
+            )
+
+            pointId = str(existing_record.id) if existing_record else str(uuid.uuid4())
+
+        payload = record.__dict__.copy()
+        embedding = payload.pop("_embedding")
+
+        return qdrant_models.PointStruct(id=pointId, vector=embedding.tolist(), payload=payload)

semantic_kernel-1.8.3/semantic_kernel/connectors/memory/qdrant/qdrant_settings.py

@@ -0,0 +1,42 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import ClassVar
+
+from pydantic import HttpUrl, SecretStr, model_validator
+
+from semantic_kernel.kernel_pydantic import KernelBaseSettings
+from semantic_kernel.utils.experimental_decorator import experimental_class
+
+IN_MEMORY_STRING = ":memory:"
+
+
+@experimental_class
+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):
+        """Validate the settings."""
+        if "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