semantic-kernel 1.8.3__zip

semantic_kernel-1.8.3/semantic_kernel/connectors/memory/azure_cognitive_search/azure_cognitive_search_memory_store.py

@@ -0,0 +1,419 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import contextlib
+import logging
+import uuid
+from inspect import isawaitable
+
+from azure.core.credentials import AzureKeyCredential, TokenCredential
+from azure.core.exceptions import ResourceNotFoundError
+from azure.search.documents.indexes.aio import SearchIndexClient
+from azure.search.documents.indexes.models import (
+    HnswAlgorithmConfiguration,
+    HnswParameters,
+    SearchIndex,
+    SearchResourceEncryptionKey,
+    VectorSearch,
+    VectorSearchProfile,
+)
+from azure.search.documents.models import VectorizedQuery
+from numpy import ndarray
+from pydantic import ValidationError
+
+from semantic_kernel.connectors.memory.azure_cognitive_search.utils import (
+    SEARCH_FIELD_EMBEDDING,
+    SEARCH_FIELD_ID,
+    dict_to_memory_record,
+    encode_id,
+    get_field_selection,
+    get_index_schema,
+    get_search_index_async_client,
+    memory_record_to_search_record,
+)
+from semantic_kernel.exceptions import MemoryConnectorInitializationError, MemoryConnectorResourceNotFound
+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 AzureCognitiveSearchMemoryStore(MemoryStoreBase):
+    """Azure Cognitive Search Memory Store."""
+
+    _search_index_client: SearchIndexClient = None
+    _vector_size: int = None
+
+    def __init__(
+        self,
+        vector_size: int,
+        search_endpoint: str | None = None,
+        admin_key: str | None = None,
+        azure_credentials: AzureKeyCredential | None = None,
+        token_credentials: TokenCredential | None = None,
+        env_file_path: str | None = None,
+        env_file_encoding: str | None = None,
+    ) -> None:
+        """Initializes a new instance of the AzureCognitiveSearchMemoryStore class.
+
+        Instantiate using Async Context Manager:
+            async with AzureCognitiveSearchMemoryStore(<...>) as memory:
+                await memory.<...>
+
+        Args:
+            vector_size (int): Embedding vector size.
+            search_endpoint (str | None): The endpoint of the Azure Cognitive Search service
+                (default: {None}).
+            admin_key (str | None): Azure Cognitive Search API key (default: {None}).
+            azure_credentials (AzureKeyCredential | None): Azure Cognitive Search credentials (default: {None}).
+            token_credentials (TokenCredential | None): Azure Cognitive Search token credentials
+                (default: {None}).
+            env_file_path (str | None): Use the environment settings file as a fallback
+                to environment variables
+            env_file_encoding (str | None): The encoding of the environment settings file
+
+        """
+        from semantic_kernel.connectors.memory.azure_cognitive_search.azure_ai_search_settings import (
+            AzureAISearchSettings,
+        )
+
+        try:
+            acs_memory_settings = AzureAISearchSettings.create(
+                env_file_path=env_file_path,
+                endpoint=search_endpoint,
+                api_key=admin_key,
+                env_file_encoding=env_file_encoding,
+            )
+        except ValidationError as exc:
+            raise MemoryConnectorInitializationError("Failed to create Azure Cognitive Search settings.") from exc
+
+        self._vector_size = vector_size
+        self._search_index_client = get_search_index_async_client(
+            search_endpoint=str(acs_memory_settings.endpoint),
+            admin_key=acs_memory_settings.api_key.get_secret_value() if acs_memory_settings.api_key else None,
+            azure_credential=azure_credentials,
+            token_credential=token_credentials,
+        )
+
+    async def close(self):
+        """Async close connection, invoked by MemoryStoreBase.__aexit__()."""
+        if self._search_index_client is not None:
+            await self._search_index_client.close()
+
+    async def create_collection(
+        self,
+        collection_name: str,
+        vector_config: HnswAlgorithmConfiguration | None = None,
+        search_resource_encryption_key: SearchResourceEncryptionKey | None = None,
+    ) -> None:
+        """Creates a new collection if it does not exist.
+
+        Args:
+            collection_name (str): The name of the collection to create.
+            vector_config (HnswVectorSearchAlgorithmConfiguration): Optional search algorithm configuration
+                                                                      (default: {None}).
+            semantic_config (SemanticConfiguration): Optional search index configuration (default: {None}).
+            search_resource_encryption_key (SearchResourceEncryptionKey): Optional Search Encryption Key
+                                                                                       (default: {None}).
+
+        Returns:
+            None
+        """
+        vector_search_profile_name = "az-vector-config"
+        if vector_config:
+            vector_search_profile = VectorSearchProfile(
+                name=vector_search_profile_name, algorithm_configuration_name=vector_config.name
+            )
+            vector_search = VectorSearch(profiles=[vector_search_profile], algorithms=[vector_config])
+        else:
+            vector_search_algorithm_name = "az-vector-hnsw-config"
+            vector_search_profile = VectorSearchProfile(
+                name=vector_search_profile_name, algorithm_configuration_name=vector_search_algorithm_name
+            )
+            vector_search = VectorSearch(
+                profiles=[vector_search_profile],
+                algorithms=[
+                    HnswAlgorithmConfiguration(
+                        name=vector_search_algorithm_name,
+                        kind="hnsw",
+                        parameters=HnswParameters(
+                            m=4,  # Number of bidirectional links, typically between 4 and 10
+                            ef_construction=400,  # Size during indexing, range: 100-1000
+                            ef_search=500,  # Size during search, range: 100-1000
+                            metric="cosine",  # Can be "cosine", "dotProduct", or "euclidean"
+                        ),
+                    )
+                ],
+            )
+
+        if not self._search_index_client:
+            raise MemoryConnectorInitializationError("Error: self._search_index_client not set 1.")
+
+        # Check to see if collection exists
+        collection_index = None
+        with contextlib.suppress(ResourceNotFoundError):
+            collection_index = await self._search_index_client.get_index(collection_name.lower())
+
+        if not collection_index:
+            # Create the search index with the semantic settings
+            index = SearchIndex(
+                name=collection_name.lower(),
+                fields=get_index_schema(self._vector_size, vector_search_profile_name),
+                vector_search=vector_search,
+                encryption_key=search_resource_encryption_key,
+            )
+
+            await self._search_index_client.create_index(index)
+
+    async def get_collections(self) -> list[str]:
+        """Gets the list of collections.
+
+        Returns:
+            List[str]: The list of collections.
+        """
+        results_list = []
+        items = self._search_index_client.list_index_names()
+        if isawaitable(items):
+            items = await items
+
+        async for result in items:
+            results_list.append(result)
+
+        return results_list
+
+    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
+        """
+        await self._search_index_client.delete_index(index=collection_name.lower())
+
+    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.
+        """
+        try:
+            collection_result = await self._search_index_client.get_index(name=collection_name.lower())
+
+            return bool(collection_result)
+        except ResourceNotFoundError:
+            return False
+
+    async def upsert(self, collection_name: str, record: MemoryRecord) -> str:
+        """Upsert a record.
+
+        Args:
+            collection_name (str): The name of the collection to upsert the record into.
+            record (MemoryRecord): The record to upsert.
+
+        Returns:
+            str: The unique record id of the record.
+        """
+        result = await self.upsert_batch(collection_name, [record])
+        if result:
+            return result[0]
+        return None
+
+    async def upsert_batch(self, collection_name: str, records: list[MemoryRecord]) -> list[str]:
+        """Upsert a batch of records.
+
+        Args:
+            collection_name (str): The name of the collection to upsert the records into.
+            records (List[MemoryRecord]): The records to upsert.
+
+        Returns:
+            List[str]: The unique database keys of the records.
+        """
+        # Initialize search client here
+        # Look up Search client class to see if exists or create
+        search_client = self._search_index_client.get_search_client(collection_name.lower())
+
+        search_records = []
+        search_ids = []
+
+        for record in records:
+            # Note:
+            # * Document id     = user provided value
+            # * MemoryRecord.id = base64(Document id)
+            if not record._id:
+                record._id = str(uuid.uuid4())
+
+            search_record = memory_record_to_search_record(record)
+            search_records.append(search_record)
+            search_ids.append(record._id)
+
+        result = await search_client.upload_documents(documents=search_records)
+        await search_client.close()
+
+        if result[0].succeeded:
+            return search_ids
+        return None
+
+    async def get(self, collection_name: str, key: str, with_embedding: bool = False) -> MemoryRecord:
+        """Gets a record.
+
+        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.
+        """
+        # Look up Search client class to see if exists or create
+        search_client = self._search_index_client.get_search_client(collection_name.lower())
+
+        try:
+            search_result = await search_client.get_document(
+                key=encode_id(key), selected_fields=get_field_selection(with_embedding)
+            )
+        except ResourceNotFoundError as exc:
+            await search_client.close()
+            raise MemoryConnectorResourceNotFound("Memory record not found") from exc
+
+        await search_client.close()
+
+        # Create Memory record from document
+        return dict_to_memory_record(search_result, with_embedding)
+
+    async def get_batch(
+        self, collection_name: str, keys: list[str], with_embeddings: bool = False
+    ) -> list[MemoryRecord]:
+        """Gets a batch of records.
+
+        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.
+        """
+        search_results = []
+
+        for key in keys:
+            search_result = await self.get(
+                collection_name=collection_name.lower(),
+                key=key,
+                with_embedding=with_embeddings,
+            )
+            search_results.append(search_result)
+
+        return search_results
+
+    async def remove_batch(self, collection_name: str, keys: list[str]) -> None:
+        """Removes a batch of records.
+
+        Args:
+            collection_name (str): The name of the collection to remove the records from.
+            keys (List[str]): The unique database keys of the records to remove.
+
+        Returns:
+            None
+        """
+        for record_id in keys:
+            await self.remove(collection_name=collection_name.lower(), key=encode_id(record_id))
+
+    async def remove(self, collection_name: str, key: str) -> None:
+        """Removes a record.
+
+        Args:
+            collection_name (str): The name of the collection to remove the record from.
+            key (str): The unique database key of the record to remove.
+
+        Returns:
+            None
+        """
+        # Look up Search client class to see if exists or create
+        search_client = self._search_index_client.get_search_client(collection_name.lower())
+        docs_to_delete = {SEARCH_FIELD_ID: encode_id(key)}
+
+        await search_client.delete_documents(documents=[docs_to_delete])
+        await search_client.close()
+
+    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 vector configuration parameters.
+
+        Args:
+            collection_name (str): The name of the collection to get the nearest match from.
+            embedding (ndarray): The embedding to find the nearest match to.
+            min_relevance_score (float): The minimum relevance score of the match. (default: {0.0})
+            with_embedding (bool): Whether to include the embedding in the result. (default: {False})
+
+        Returns:
+            Tuple[MemoryRecord, float]: The record and the relevance score.
+        """
+        memory_records = await self.get_nearest_matches(
+            collection_name=collection_name,
+            embedding=embedding,
+            min_relevance_score=min_relevance_score,
+            with_embeddings=with_embedding,
+            limit=1,
+        )
+
+        if len(memory_records) > 0:
+            return memory_records[0]
+        return None
+
+    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 vector configuration.
+
+        Parameters:
+            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.
+        """
+        # Look up Search client class to see if exists or create
+        search_client = self._search_index_client.get_search_client(collection_name.lower())
+
+        vector = VectorizedQuery(vector=embedding.flatten(), k_nearest_neighbors=limit, fields=SEARCH_FIELD_EMBEDDING)
+
+        search_results = await search_client.search(
+            search_text="*",
+            select=get_field_selection(with_embeddings),
+            vector_queries=[vector],
+        )
+
+        if not search_results or search_results is None:
+            await search_client.close()
+            return []
+
+        # Convert the results to MemoryRecords
+        nearest_results = []
+        async for search_record in search_results:
+            if search_record["@search.score"] < min_relevance_score:
+                continue
+
+            memory_record = dict_to_memory_record(search_record, with_embeddings)
+            nearest_results.append((memory_record, search_record["@search.score"]))
+
+        await search_client.close()
+        return nearest_results

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

@@ -0,0 +1,232 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import base64
+import os
+
+from azure.core.credentials import AzureKeyCredential, TokenCredential
+from azure.search.documents.indexes.aio import SearchIndexClient
+from azure.search.documents.indexes.models import SearchableField, SearchField, SearchFieldDataType, SimpleField
+from dotenv import load_dotenv
+
+from semantic_kernel.const import USER_AGENT
+from semantic_kernel.exceptions import ServiceInitializationError
+from semantic_kernel.memory.memory_record import MemoryRecord
+
+SEARCH_FIELD_ID = "Id"
+SEARCH_FIELD_TEXT = "Text"
+SEARCH_FIELD_EMBEDDING = "Embedding"
+SEARCH_FIELD_SRC = "ExternalSourceName"
+SEARCH_FIELD_DESC = "Description"
+SEARCH_FIELD_METADATA = "AdditionalMetadata"
+SEARCH_FIELD_IS_REF = "IsReference"
+
+
+def get_search_index_async_client(
+    search_endpoint: str | None = None,
+    admin_key: str | None = None,
+    azure_credential: AzureKeyCredential | None = None,
+    token_credential: TokenCredential | None = None,
+):
+    """Return a client for Azure Cognitive Search.
+
+    Args:
+        search_endpoint (str): Optional endpoint (default: {None}).
+        admin_key (str): Optional API key (default: {None}).
+        azure_credential (AzureKeyCredential): Optional Azure credentials (default: {None}).
+        token_credential (TokenCredential): Optional Token credential (default: {None}).
+    """
+    ENV_VAR_ENDPOINT = "AZURE_COGNITIVE_SEARCH_ENDPOINT"
+    ENV_VAR_API_KEY = "AZURE_COGNITIVE_SEARCH_ADMIN_KEY"
+
+    # Load environment variables
+    load_dotenv()
+
+    # Service endpoint
+    if search_endpoint:
+        service_endpoint = search_endpoint
+    elif os.getenv(ENV_VAR_ENDPOINT):
+        service_endpoint = os.getenv(ENV_VAR_ENDPOINT)
+    else:
+        raise ServiceInitializationError("Error: missing Azure Cognitive Search client endpoint.")
+
+    if service_endpoint is None:
+        print(service_endpoint)
+        raise ServiceInitializationError("Error: Azure Cognitive Search client not set.")
+
+    # Credentials
+    if admin_key:
+        azure_credential = AzureKeyCredential(admin_key)
+    elif azure_credential:
+        azure_credential = azure_credential
+    elif token_credential:
+        token_credential = token_credential
+    elif os.getenv(ENV_VAR_API_KEY):
+        azure_credential = AzureKeyCredential(os.getenv(ENV_VAR_API_KEY))
+    else:
+        raise ServiceInitializationError("Error: missing Azure Cognitive Search client credentials.")
+
+    if azure_credential is None and token_credential is None:
+        raise ServiceInitializationError("Error: Azure Cognitive Search credentials not set.")
+
+    sk_headers = {USER_AGENT: "Semantic-Kernel"}
+
+    if azure_credential:
+        return SearchIndexClient(endpoint=service_endpoint, credential=azure_credential, headers=sk_headers)
+
+    if token_credential:
+        return SearchIndexClient(endpoint=service_endpoint, credential=token_credential, headers=sk_headers)
+
+    raise ValueError("Error: unable to create Azure Cognitive Search client.")
+
+
+def get_index_schema(vector_size: int, vector_search_profile_name: str) -> list:
+    """Return the schema of search indexes.
+
+    Args:
+        vector_size (int): The size of the vectors being stored in collection/index.
+        vector_search_profile_name (str): The name of the vector search profile.
+
+    Returns:
+        list: The Azure Cognitive Search schema as list type.
+    """
+    return [
+        SimpleField(
+            name=SEARCH_FIELD_ID,
+            type=SearchFieldDataType.String,
+            searchable=True,
+            filterable=True,
+            retrievable=True,
+            key=True,
+        ),
+        SearchableField(
+            name=SEARCH_FIELD_TEXT,
+            type=SearchFieldDataType.String,
+            searchable=True,
+            filterable=True,
+            retrievable=True,
+        ),
+        SearchField(
+            name=SEARCH_FIELD_EMBEDDING,
+            type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
+            searchable=True,
+            vector_search_dimensions=vector_size,
+            vector_search_profile_name=vector_search_profile_name,
+        ),
+        SimpleField(
+            name=SEARCH_FIELD_SRC,
+            type=SearchFieldDataType.String,
+            searchable=True,
+            filterable=True,
+            retrievable=True,
+        ),
+        SimpleField(
+            name=SEARCH_FIELD_DESC,
+            type=SearchFieldDataType.String,
+            searchable=True,
+            filterable=True,
+            retrievable=True,
+        ),
+        SimpleField(
+            name=SEARCH_FIELD_METADATA,
+            type=SearchFieldDataType.String,
+            searchable=True,
+            filterable=True,
+            retrievable=True,
+        ),
+        SimpleField(
+            name=SEARCH_FIELD_IS_REF,
+            type=SearchFieldDataType.Boolean,
+            searchable=True,
+            filterable=True,
+            retrievable=True,
+        ),
+    ]
+
+
+def get_field_selection(with_embeddings: bool) -> list[str]:
+    """Get the list of fields to search and load.
+
+    Args:
+        with_embeddings (bool): Whether to include the embedding vector field.
+
+    Returns:
+        List[str]: List of fields.
+    """
+    field_selection = [
+        SEARCH_FIELD_ID,
+        SEARCH_FIELD_TEXT,
+        SEARCH_FIELD_SRC,
+        SEARCH_FIELD_DESC,
+        SEARCH_FIELD_METADATA,
+        SEARCH_FIELD_IS_REF,
+    ]
+
+    if with_embeddings:
+        field_selection.append(SEARCH_FIELD_EMBEDDING)
+
+    return field_selection
+
+
+def dict_to_memory_record(data: dict, with_embeddings: bool) -> MemoryRecord:
+    """Converts a search result to a MemoryRecord.
+
+    Args:
+        data (dict): Azure Cognitive Search result data.
+        with_embeddings (bool): Whether to include the embedding vector field.
+
+    Returns:
+        MemoryRecord: The MemoryRecord from Azure Cognitive Search Data Result.
+    """
+    return MemoryRecord(
+        id=decode_id(data[SEARCH_FIELD_ID]),
+        key=data[SEARCH_FIELD_ID],
+        text=data[SEARCH_FIELD_TEXT],
+        external_source_name=data[SEARCH_FIELD_SRC],
+        description=data[SEARCH_FIELD_DESC],
+        additional_metadata=data[SEARCH_FIELD_METADATA],
+        is_reference=data[SEARCH_FIELD_IS_REF],
+        embedding=data[SEARCH_FIELD_EMBEDDING] if with_embeddings else None,
+        timestamp=None,
+    )
+
+
+def memory_record_to_search_record(record: MemoryRecord) -> dict:
+    """Convert a MemoryRecord to a dictionary.
+
+    Args:
+        record (MemoryRecord): The MemoryRecord from Azure Cognitive Search Data Result.
+
+    Returns:
+        data (dict): Dictionary data.
+    """
+    return {
+        SEARCH_FIELD_ID: encode_id(record._id),
+        SEARCH_FIELD_TEXT: str(record._text),
+        SEARCH_FIELD_SRC: record._external_source_name or "",
+        SEARCH_FIELD_DESC: record._description or "",
+        SEARCH_FIELD_METADATA: record._additional_metadata or "",
+        SEARCH_FIELD_IS_REF: str(record._is_reference),
+        SEARCH_FIELD_EMBEDDING: record._embedding.tolist(),
+    }
+
+
+def encode_id(id: str) -> str:
+    """Encode a record id to ensure compatibility with Azure Cognitive Search.
+
+    Azure Cognitive Search keys can contain only letters, digits, underscore, dash,
+    equal sign, recommending to encode values with a URL-safe algorithm.
+    """
+    id_bytes = id.encode("ascii")
+    base64_bytes = base64.b64encode(id_bytes)
+    return base64_bytes.decode("ascii")
+
+
+def decode_id(base64_id: str) -> str:
+    """Decode a record id to the original value.
+
+    Azure Cognitive Search keys can contain only letters, digits, underscore, dash,
+    equal sign, recommending to encode values with a URL-safe algorithm.
+    """
+    base64_bytes = base64_id.encode("ascii")
+    message_bytes = base64.b64decode(base64_bytes)
+    return message_bytes.decode("ascii")

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

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from semantic_kernel.connectors.memory.azure_cosmosdb.azure_cosmos_db_memory_store import (
+    AzureCosmosDBMemoryStore,
+)
+from semantic_kernel.connectors.memory.azure_cosmosdb.azure_cosmosdb_settings import AzureCosmosDBSettings
+
+__all__ = ["AzureCosmosDBMemoryStore", "AzureCosmosDBSettings"]