semantic-kernel 1.8.3__zip

semantic_kernel-1.8.3/semantic_kernel/connectors/memory/astradb/astradb_memory_store.py

@@ -0,0 +1,329 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+
+import aiohttp
+from numpy import ndarray
+from pydantic import ValidationError
+
+from semantic_kernel.connectors.memory.astradb.astra_client import AstraClient
+from semantic_kernel.connectors.memory.astradb.astradb_settings import AstraDBSettings
+from semantic_kernel.connectors.memory.astradb.utils import build_payload, parse_payload
+from semantic_kernel.exceptions import MemoryConnectorInitializationError
+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
+
+MAX_DIMENSIONALITY = 20000
+MAX_UPSERT_BATCH_SIZE = 100
+MAX_QUERY_WITHOUT_METADATA_BATCH_SIZE = 10000
+MAX_QUERY_WITH_METADATA_BATCH_SIZE = 1000
+MAX_FETCH_BATCH_SIZE = 1000
+MAX_DELETE_BATCH_SIZE = 1000
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+@experimental_class
+class AstraDBMemoryStore(MemoryStoreBase):
+    """A memory store that uses Astra database as the backend."""
+
+    def __init__(
+        self,
+        astra_application_token: str,
+        astra_id: str,
+        astra_region: str,
+        keyspace_name: str,
+        embedding_dim: int,
+        similarity: str,
+        session: aiohttp.ClientSession | None = None,
+        env_file_path: str | None = None,
+        env_file_encoding: str | None = None,
+    ) -> None:
+        """Initializes a new instance of the AstraDBMemoryStore class.
+
+        Args:
+            astra_application_token (str): The Astra application token.
+            astra_id (str): The Astra id of database.
+            astra_region (str): The Astra region
+            keyspace_name (str): The Astra keyspace
+            embedding_dim (int): The dimensionality to use for new collections.
+            similarity (str): TODO
+            session: Optional session parameter
+            env_file_path (str | None): Use the environment settings file as a
+                fallback to environment variables. (Optional)
+            env_file_encoding (str | None): The encoding of the environment settings file. (Optional)
+        """
+        try:
+            astradb_settings = AstraDBSettings.create(
+                app_token=astra_application_token,
+                db_id=astra_id,
+                region=astra_region,
+                keyspace=keyspace_name,
+                env_file_path=env_file_path,
+                env_file_encoding=env_file_encoding,
+            )
+        except ValidationError as ex:
+            raise MemoryConnectorInitializationError("Failed to create AstraDB settings.", ex) from ex
+
+        self._embedding_dim = embedding_dim
+        self._similarity = similarity
+        self._session = session
+
+        if self._embedding_dim > MAX_DIMENSIONALITY:
+            raise MemoryConnectorInitializationError(
+                f"Dimensionality of {self._embedding_dim} exceeds "
+                f"the maximum allowed value of {MAX_DIMENSIONALITY}."
+            )
+
+        self._client = AstraClient(
+            astra_id=astradb_settings.db_id,
+            astra_region=astradb_settings.region,
+            astra_application_token=(
+                astradb_settings.app_token.get_secret_value() if astradb_settings.app_token else None
+            ),
+            keyspace_name=astradb_settings.keyspace,
+            embedding_dim=embedding_dim,
+            similarity_function=similarity,
+            session=self._session,
+        )
+
+    async def get_collections(self) -> list[str]:
+        """Gets the list of collections.
+
+        Returns:
+            List[str]: The list of collections.
+        """
+        return await self._client.find_collections(False)
+
+    async def create_collection(
+        self,
+        collection_name: str,
+        dimension_num: int | None = None,
+        distance_type: str | None = "cosine",
+    ) -> None:
+        """Creates a new collection in Astra if it does not exist.
+
+        Args:
+            collection_name (str): The name of the collection to create.
+            dimension_num (int): The dimension of the vectors to be stored in this collection.
+            distance_type (str): Specifies the similarity metric to be used when querying or comparing vectors within
+            this collection. The available options are dot_product, euclidean, and cosine.
+
+        Returns:
+            None
+        """
+        dimension_num = dimension_num if dimension_num is not None else self._embedding_dim
+        distance_type = distance_type if distance_type is not None else self._similarity
+
+        if dimension_num > MAX_DIMENSIONALITY:
+            raise ValueError(
+                f"Dimensionality of {dimension_num} exceeds " + f"the maximum allowed value of {MAX_DIMENSIONALITY}."
+            )
+
+        result = await self._client.create_collection(collection_name, dimension_num, distance_type)
+        if result is True:
+            logger.info(f"Collection {collection_name} created.")
+
+    async def delete_collection(self, collection_name: str) -> None:
+        """Deletes a collection.
+
+        Args:
+            collection_name (str): The name of the collection to delete.
+
+        Returns:
+            None
+        """
+        result = await self._client.delete_collection(collection_name)
+        logger.log(
+            logging.INFO if result is True else logging.WARNING,
+            f"Collection {collection_name} {'deleted.' if result is True else 'does not exist.'}",
+        )
+
+    async def does_collection_exist(self, collection_name: str) -> bool:
+        """Checks if a collection exists.
+
+        Args:
+            collection_name (str): The name of the collection to check.
+
+        Returns:
+            bool: True if the collection exists; otherwise, False.
+        """
+        return await self._client.find_collection(collection_name)
+
+    async def upsert(self, collection_name: str, record: MemoryRecord) -> str:
+        """Upsert a memory record into the data store.
+
+        Does not guarantee that the collection exists.
+        If the record already exists, it will be updated.
+        If the record does not exist, it will be created.
+
+        Args:
+            collection_name (str): The name associated with a collection of embeddings.
+            record (MemoryRecord): The memory record to upsert.
+
+        Returns:
+            str: The unique identifier for the memory record.
+        """
+        filter = {"_id": record._id}
+        update = {"$set": build_payload(record)}
+        status = await self._client.update_document(collection_name, filter, update, True)
+
+        return status.get("upsertedId", record._id)
+
+    async def upsert_batch(self, collection_name: str, records: list[MemoryRecord]) -> list[str]:
+        """Upsert a batch of memory records into the data store.
+
+        Does not guarantee that the collection exists.
+        If the record already exists, it will be updated.
+        If the record does not exist, it will be created.
+
+        Args:
+            collection_name (str): The name associated with a collection of embeddings.
+            records (List[MemoryRecord]): The memory records to upsert.
+
+        Returns:
+            List[str]: The unique identifiers for the memory record.
+        """
+        return await asyncio.gather(*[self.upsert(collection_name, record) for record in records])
+
+    async def get(self, collection_name: str, key: str, with_embedding: bool = False) -> MemoryRecord:
+        """Gets a record. Does not guarantee that the collection exists.
+
+        Args:
+            collection_name (str): The name of the collection to get the record from.
+            key (str): The unique database key of the record.
+            with_embedding (bool): Whether to include the embedding in the result. (default: {False})
+
+        Returns:
+            MemoryRecord: The record.
+        """
+        filter = {"_id": key}
+        documents = await self._client.find_documents(
+            collection_name=collection_name,
+            filter=filter,
+            include_vector=with_embedding,
+        )
+
+        if len(documents) == 0:
+            raise KeyError(f"Record with key '{key}' does not exist")
+
+        return parse_payload(documents[0])
+
+    async def get_batch(
+        self, collection_name: str, keys: list[str], with_embeddings: bool = False
+    ) -> list[MemoryRecord]:
+        """Gets a batch of records. Does not guarantee that the collection exists.
+
+        Args:
+            collection_name (str): The name of the collection to get the records from.
+            keys (List[str]): The unique database keys of the records.
+            with_embeddings (bool): Whether to include the embeddings in the results. (default: {False})
+
+        Returns:
+            List[MemoryRecord]: The records.
+        """
+        filter = {"_id": {"$in": keys}}
+        documents = await self._client.find_documents(
+            collection_name=collection_name,
+            filter=filter,
+            include_vector=with_embeddings,
+        )
+        return [parse_payload(document) for document in documents]
+
+    async def remove(self, collection_name: str, key: str) -> None:
+        """Removes a memory record from the data store. Does not guarantee that the collection exists.
+
+        Args:
+            collection_name (str): The name of the collection to remove the record from.
+            key (str): The unique id associated with the memory record to remove.
+
+        Returns:
+            None
+        """
+        filter = {"_id": key}
+        await self._client.delete_documents(collection_name, filter)
+
+    async def remove_batch(self, collection_name: str, keys: list[str]) -> None:
+        """Removes a batch of records. Does not guarantee that the collection exists.
+
+        Args:
+            collection_name (str): The name of the collection to remove the records from.
+            keys (List[str]): The unique ids associated with the memory records to remove.
+
+        Returns:
+            None
+        """
+        filter = {"_id": {"$in": keys}}
+        await self._client.delete_documents(collection_name, filter)
+
+    async def get_nearest_match(
+        self,
+        collection_name: str,
+        embedding: ndarray,
+        min_relevance_score: float = 0.0,
+        with_embedding: bool = False,
+    ) -> tuple[MemoryRecord, float]:
+        """Gets the nearest match to an embedding using cosine similarity.
+
+        Args:
+            collection_name (str): The name of the collection to get the nearest matches from.
+            embedding (ndarray): The embedding to find the nearest matches to.
+            min_relevance_score (float): The minimum relevance score of the matches. (default: {0.0})
+            with_embedding (bool): Whether to include the embeddings in the results. (default: {False})
+
+        Returns:
+            Tuple[MemoryRecord, float]: The record and the relevance score.
+        """
+        matches = await self.get_nearest_matches(
+            collection_name=collection_name,
+            embedding=embedding,
+            limit=1,
+            min_relevance_score=min_relevance_score,
+            with_embeddings=with_embedding,
+        )
+        return matches[0]
+
+    async def get_nearest_matches(
+        self,
+        collection_name: str,
+        embedding: ndarray,
+        limit: int,
+        min_relevance_score: float = 0.0,
+        with_embeddings: bool = False,
+    ) -> list[tuple[MemoryRecord, float]]:
+        """Gets the nearest matches to an embedding using cosine similarity.
+
+        Args:
+            collection_name (str): The name of the collection to get the nearest matches from.
+            embedding (ndarray): The embedding to find the nearest matches to.
+            limit (int): The maximum number of matches to return.
+            min_relevance_score (float): The minimum relevance score of the matches. (default: {0.0})
+            with_embeddings (bool): Whether to include the embeddings in the results. (default: {False})
+
+        Returns:
+            List[Tuple[MemoryRecord, float]]: The records and their relevance scores.
+        """
+        matches = await self._client.find_documents(
+            collection_name=collection_name,
+            vector=embedding.tolist(),
+            limit=limit,
+            include_similarity=True,
+            include_vector=with_embeddings,
+        )
+
+        if min_relevance_score:
+            matches = [match for match in matches if match["$similarity"] >= min_relevance_score]
+
+        return (
+            [
+                (
+                    parse_payload(match),
+                    match["$similarity"],
+                )
+                for match in matches
+            ]
+            if len(matches) > 0
+            else []
+        )

semantic_kernel-1.8.3/semantic_kernel/connectors/memory/astradb/astradb_settings.py

@@ -0,0 +1,27 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import ClassVar
+
+from pydantic import SecretStr
+
+from semantic_kernel.kernel_pydantic import KernelBaseSettings
+from semantic_kernel.utils.experimental_decorator import experimental_class
+
+
+@experimental_class
+class AstraDBSettings(KernelBaseSettings):
+    """AstraDB model settings.
+
+    Settings for AstraDB connection:
+    - app_token: SecretStr | None - AstraDB token (Env var ASTRADB_APP_TOKEN)
+    - db_id: str | None - AstraDB database ID (Env var ASTRADB_DB_ID)
+    - region: str | None - AstraDB region (Env var ASTRADB_REGION)
+    - keyspace: str | None - AstraDB keyspace (Env var ASTRADB_KEYSPACE)
+    """
+
+    env_prefix: ClassVar[str] = "ASTRADB_"
+
+    app_token: SecretStr
+    db_id: str
+    region: str
+    keyspace: str

semantic_kernel-1.8.3/semantic_kernel/connectors/memory/astradb/utils.py

@@ -0,0 +1,51 @@
+# Copyright (c) Microsoft. All rights reserved.
+from typing import Any
+
+import aiohttp
+import numpy
+
+from semantic_kernel.memory.memory_record import MemoryRecord
+
+
+class AsyncSession:
+    """A wrapper around aiohttp.ClientSession that can be used as an async context manager."""
+
+    def __init__(self, session: aiohttp.ClientSession = None):
+        """Initializes a new instance of the AsyncSession class."""
+        self._session = session if session else aiohttp.ClientSession()
+
+    async def __aenter__(self):
+        """Enter the session."""
+        return await self._session.__aenter__()
+
+    async def __aexit__(self, *args, **kwargs):
+        """Close the session."""
+        await self._session.close()
+
+
+def build_payload(record: MemoryRecord) -> dict[str, Any]:
+    """Builds a metadata payload to be sent to AstraDb from a MemoryRecord."""
+    payload: dict[str, Any] = {}
+    payload["$vector"] = record.embedding.tolist()
+    if record._text:
+        payload["text"] = record._text
+    if record._description:
+        payload["description"] = record._description
+    if record._additional_metadata:
+        payload["additional_metadata"] = record._additional_metadata
+    return payload
+
+
+def parse_payload(document: dict[str, Any]) -> MemoryRecord:
+    """Parses a record from AstraDb into a MemoryRecord."""
+    text = document.get("text")
+    description = document.get("description")
+    additional_metadata = document.get("additional_metadata")
+
+    return MemoryRecord.local_record(
+        id=document["_id"],
+        description=description,
+        text=text,
+        additional_metadata=additional_metadata,
+        embedding=document["$vector"] if "$vector" in document else numpy.array([]),
+    )

semantic_kernel-1.8.3/semantic_kernel/connectors/memory/azure_ai_search/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from semantic_kernel.connectors.memory.azure_ai_search.azure_ai_search_collection import AzureAISearchCollection
+from semantic_kernel.connectors.memory.azure_ai_search.azure_ai_search_store import AzureAISearchStore
+from semantic_kernel.connectors.memory.azure_cognitive_search.azure_ai_search_settings import AzureAISearchSettings
+
+__all__ = ["AzureAISearchCollection", "AzureAISearchSettings", "AzureAISearchStore"]

semantic_kernel-1.8.3/semantic_kernel/connectors/memory/azure_ai_search/azure_ai_search_collection.py

@@ -0,0 +1,214 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import sys
+from collections.abc import Sequence
+from typing import Any, ClassVar, Generic, TypeVar
+
+if sys.version_info >= (3, 12):
+    from typing import override  # pragma: no cover
+else:
+    from typing_extensions import override  # pragma: no cover
+
+from azure.search.documents.aio import SearchClient
+from azure.search.documents.indexes.aio import SearchIndexClient
+from azure.search.documents.indexes.models import SearchIndex
+from pydantic import ValidationError
+
+from semantic_kernel.connectors.memory.azure_ai_search.utils import (
+    data_model_definition_to_azure_ai_search_index,
+    get_search_client,
+    get_search_index_client,
+)
+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 MemoryConnectorException, MemoryConnectorInitializationError
+from semantic_kernel.utils.experimental_decorator import experimental_class
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+TModel = TypeVar("TModel")
+
+
+@experimental_class
+class AzureAISearchCollection(VectorStoreRecordCollection[str, TModel], Generic[TModel]):
+    """Azure AI Search collection implementation."""
+
+    search_client: SearchClient
+    search_index_client: SearchIndexClient
+    supported_key_types: ClassVar[list[str] | None] = ["str"]
+    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,
+        search_index_client: SearchIndexClient | None = None,
+        search_client: SearchClient | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initializes a new instance of the AzureAISearchCollection class.
+
+        Args:
+            data_model_type (type[TModel]): The type of the data model.
+            data_model_definition (VectorStoreRecordDefinition): The model definition, optional.
+            collection_name (str): The name of the collection, optional.
+            search_index_client (SearchIndexClient): The search index client for interacting with Azure AI Search,
+                used for creating and deleting indexes.
+            search_client (SearchClient): The search client for interacting with Azure AI Search,
+                used for record operations.
+            **kwargs: Additional keyword arguments, including:
+                The same keyword arguments used for AzureAISearchVectorStore:
+                    search_endpoint: str | None = None,
+                    api_key: str | None = None,
+                    azure_credentials: AzureKeyCredential | None = None,
+                    token_credentials: AsyncTokenCredential | TokenCredential | None = None,
+                    env_file_path: str | None = None,
+                    env_file_encoding: str | None = None
+
+        """
+        if search_client and search_index_client:
+            if not collection_name:
+                collection_name = search_client._index_name
+            elif search_client._index_name != collection_name:
+                raise MemoryConnectorInitializationError(
+                    "Search client and search index client have different index names."
+                )
+            super().__init__(
+                data_model_type=data_model_type,
+                data_model_definition=data_model_definition,
+                collection_name=collection_name,
+                search_client=search_client,
+                search_index_client=search_index_client,
+            )
+            return
+
+        if search_index_client:
+            if not collection_name:
+                raise MemoryConnectorInitializationError("Collection name is required.")
+            super().__init__(
+                data_model_type=data_model_type,
+                data_model_definition=data_model_definition,
+                collection_name=collection_name,
+                search_client=get_search_client(
+                    search_index_client=search_index_client, collection_name=collection_name
+                ),
+                search_index_client=search_index_client,
+            )
+            return
+
+        from semantic_kernel.connectors.memory.azure_ai_search.azure_ai_search_settings import (
+            AzureAISearchSettings,
+        )
+
+        try:
+            azure_ai_search_settings = AzureAISearchSettings.create(
+                env_file_path=kwargs.get("env_file_path", None),
+                endpoint=kwargs.get("search_endpoint", None),
+                api_key=kwargs.get("api_key", None),
+                env_file_encoding=kwargs.get("env_file_encoding", None),
+                index_name=collection_name,
+            )
+        except ValidationError as exc:
+            raise MemoryConnectorInitializationError("Failed to create Azure Cognitive Search settings.") from exc
+        search_index_client = get_search_index_client(
+            azure_ai_search_settings=azure_ai_search_settings,
+            azure_credential=kwargs.get("azure_credentials", None),
+            token_credential=kwargs.get("token_credentials", None),
+        )
+        if not azure_ai_search_settings.index_name:
+            raise MemoryConnectorInitializationError("Collection name is required.")
+
+        super().__init__(
+            data_model_type=data_model_type,
+            data_model_definition=data_model_definition,
+            collection_name=azure_ai_search_settings.index_name,
+            search_client=get_search_client(
+                search_index_client=search_index_client, collection_name=azure_ai_search_settings.index_name
+            ),
+            search_index_client=search_index_client,
+        )
+
+    @override
+    async def _inner_upsert(
+        self,
+        records: Sequence[Any],
+        **kwargs: Any,
+    ) -> Sequence[str]:
+        if not isinstance(records, list):
+            records = list(records)
+        results = await self.search_client.merge_or_upload_documents(documents=records, **kwargs)
+        return [result.key for result in results]  # type: ignore
+
+    @override
+    async def _inner_get(self, keys: Sequence[str], **kwargs: Any) -> Sequence[dict[str, Any]]:
+        client = self.search_client
+        if "selected_fields" in kwargs:
+            selected_fields = kwargs["selected_fields"]
+        elif "include_vector" in kwargs and not kwargs["include_vector"]:
+            selected_fields = [
+                name
+                for name, field in self.data_model_definition.fields.items()
+                if not isinstance(field, VectorStoreRecordVectorField)
+            ]
+        else:
+            selected_fields = ["*"]
+
+        result = await asyncio.gather(
+            *[client.get_document(key=key, selected_fields=selected_fields) for key in keys],
+            return_exceptions=True,
+        )
+        return [res for res in result if not isinstance(res, BaseException)]
+
+    @override
+    async def _inner_delete(self, keys: Sequence[str], **kwargs: Any) -> None:
+        await self.search_client.delete_documents(documents=[{self._key_field_name: key} for key in keys])
+
+    @override
+    def _serialize_dicts_to_store_models(self, records: Sequence[dict[str, Any]], **kwargs: Any) -> Sequence[Any]:
+        return records
+
+    @override
+    def _deserialize_store_models_to_dicts(self, records: Sequence[Any], **kwargs: Any) -> Sequence[dict[str, Any]]:
+        return records
+
+    @override
+    async def create_collection(self, **kwargs) -> None:
+        """Create a new collection in Azure AI Search.
+
+        Args:
+            **kwargs: Additional keyword arguments.
+                index (SearchIndex): The search index to create, if this is supplied
+                    this is used instead of a index created based on the definition.
+                encryption_key (SearchResourceEncryptionKey): The encryption key to use,
+                    not used when index is supplied.
+                other kwargs are passed to the create_index method.
+        """
+        if index := kwargs.pop("index", None):
+            if isinstance(index, SearchIndex):
+                await self.search_index_client.create_index(index=index, **kwargs)
+                return
+            raise MemoryConnectorException("Invalid index type supplied.")
+        await self.search_index_client.create_index(
+            index=data_model_definition_to_azure_ai_search_index(
+                collection_name=self.collection_name,
+                definition=self.data_model_definition,
+                encryption_key=kwargs.pop("encryption_key", None),
+            ),
+            **kwargs,
+        )
+
+    @override
+    async def does_collection_exist(self, **kwargs) -> bool:
+        if "params" not in kwargs:
+            kwargs["params"] = {"select": ["name"]}
+        return self.collection_name in [
+            index_name async for index_name in self.search_index_client.list_index_names(**kwargs)
+        ]
+
+    @override
+    async def delete_collection(self, **kwargs) -> None:
+        await self.search_index_client.delete_index(self.collection_name, **kwargs)

semantic_kernel-1.8.3/semantic_kernel/connectors/memory/azure_ai_search/azure_ai_search_settings.py

@@ -0,0 +1,25 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import ClassVar
+
+from pydantic import SecretStr
+
+from semantic_kernel.kernel_pydantic import HttpsUrl, KernelBaseSettings
+from semantic_kernel.utils.experimental_decorator import experimental_class
+
+
+@experimental_class
+class AzureAISearchSettings(KernelBaseSettings):
+    """Azure AI Search model settings currently used by the AzureCognitiveSearchMemoryStore connector.
+
+    Args:
+    - api_key: SecretStr - Azure AI Search API key (Env var AZURE_AI_SEARCH_API_KEY)
+    - endpoint: HttpsUrl - Azure AI Search endpoint (Env var AZURE_AI_SEARCH_ENDPOINT)
+    - index_name: str - Azure AI Search index name (Env var AZURE_AI_SEARCH_INDEX_NAME)
+    """
+
+    env_prefix: ClassVar[str] = "AZURE_AI_SEARCH_"
+
+    api_key: SecretStr | None = None
+    endpoint: HttpsUrl
+    index_name: str | None = None