semantic-kernel 1.8.3__zip

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

@@ -0,0 +1,347 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import sys
+from typing import Any
+
+if sys.version >= "3.12":
+    from typing import override  # pragma: no cover
+else:
+    from typing_extensions import override  # pragma: no cover
+
+import numpy as np
+
+from semantic_kernel.connectors.memory.azure_cosmosdb.azure_cosmos_db_store_api import AzureCosmosDBStoreApi
+from semantic_kernel.connectors.memory.azure_cosmosdb.utils import CosmosDBSimilarityType, CosmosDBVectorSearchType
+from semantic_kernel.memory.memory_record import MemoryRecord
+from semantic_kernel.utils.experimental_decorator import experimental_class
+
+
+@experimental_class
+class MongoStoreApi(AzureCosmosDBStoreApi):
+    """MongoStoreApi class for the Azure Cosmos DB Mongo store."""
+
+    database = None
+    collection_name: str
+    index_name = None
+    vector_dimensions = None
+    num_lists = None
+    similarity = None
+    collection = None
+    kind = None
+    m = None
+    ef_construction = None
+    ef_search = None
+
+    """
+    Args:
+        collection_name: Name of the collection for the azure cosmos db mongo store
+        index_name: Index for the collection
+        vector_dimensions: Number of dimensions for vector similarity.
+            The maximum number of supported dimensions is 2000
+        num_lists: This integer is the number of clusters that the
+            inverted file (IVF) index uses to group the vector data.
+            We recommend that numLists is set to documentCount/1000
+            for up to 1 million documents and to sqrt(documentCount)
+            for more than 1 million documents.
+            Using a numLists value of 1 is akin to performing
+            brute-force search, which has limited performance
+        similarity: Similarity metric to use with the IVF index.
+            Possible options are:
+                - CosmosDBSimilarityType.COS (cosine distance),
+                - CosmosDBSimilarityType.L2 (Euclidean distance), and
+                - CosmosDBSimilarityType.IP (inner product).
+        collection:
+        kind: Type of vector index to create.
+            Possible options are:
+                - vector-ivf
+                - vector-hnsw: available as a preview feature only,
+                               to enable visit https://learn.microsoft.com/en-us/azure/azure-resource-manager/management/preview-features
+        m: The max number of connections per layer (16 by default, minimum
+           value is 2, maximum value is 100). Higher m is suitable for datasets
+           with high dimensionality and/or high accuracy requirements.
+        ef_construction: the size of the dynamic candidate list for constructing
+                        the graph (64 by default, minimum value is 4, maximum
+                        value is 1000). Higher ef_construction will result in
+                        better index quality and higher accuracy, but it will
+                        also increase the time required to build the index.
+                        ef_construction has to be at least 2 * m
+       ef_search: The size of the dynamic candidate list for search (40 by default).
+                  A higher value provides better recall at  the cost of speed.
+       database: The Mongo Database object of the azure cosmos db mongo store
+    """
+
+    def __init__(
+        self,
+        collection_name: str,
+        index_name: str,
+        vector_dimensions: int,
+        num_lists: int,
+        similarity: CosmosDBSimilarityType,
+        kind: CosmosDBVectorSearchType,
+        m: int,
+        ef_construction: int,
+        ef_search: int,
+        database=None,
+    ):
+        """Initializes a new instance of the MongoStoreApi class."""
+        self.database = database
+        self.collection_name = collection_name
+        self.index_name = index_name
+        self.num_lists = num_lists
+        self.similarity = similarity
+        self.vector_dimensions = vector_dimensions
+        self.kind = kind
+        self.m = m
+        self.ef_construction = ef_construction
+        self.ef_search = ef_search
+
+    @override
+    async def create_collection(self, collection_name: str) -> None:
+        if (
+            not await self.does_collection_exist(collection_name)
+            and self.index_name not in self.database[collection_name].list_indexes()
+        ):
+            # check the kind of vector search to be performed
+            # prepare the command accordingly
+            create_index_commands = {}
+            if self.kind == CosmosDBVectorSearchType.VECTOR_IVF:
+                create_index_commands = self._get_vector_index_ivf(
+                    collection_name, self.kind, self.num_lists, self.similarity, self.vector_dimensions
+                )
+            elif self.kind == CosmosDBVectorSearchType.VECTOR_HNSW:
+                create_index_commands = self._get_vector_index_hnsw(
+                    collection_name,
+                    self.kind,
+                    self.m,
+                    self.ef_construction,
+                    self.similarity,
+                    self.vector_dimensions,
+                )
+            # invoke the command from the database object
+            self.database.command(create_index_commands)
+        self.collection = self.database[collection_name]
+
+    def _get_vector_index_ivf(
+        self, collection_name: str, kind: str, num_lists: int, similarity: str, dimensions: int
+    ) -> dict[str, Any]:
+        return {
+            "createIndexes": collection_name,
+            "indexes": [
+                {
+                    "name": self.index_name,
+                    "key": {"embedding": "cosmosSearch"},
+                    "cosmosSearchOptions": {
+                        "kind": kind,
+                        "numLists": num_lists,
+                        "similarity": similarity,
+                        "dimensions": dimensions,
+                    },
+                }
+            ],
+        }
+
+    def _get_vector_index_hnsw(
+        self, collection_name: str, kind: str, m: int, ef_construction: int, similarity: str, dimensions: int
+    ) -> dict[str, Any]:
+        return {
+            "createIndexes": collection_name,
+            "indexes": [
+                {
+                    "name": self.index_name,
+                    "key": {"embedding": "cosmosSearch"},
+                    "cosmosSearchOptions": {
+                        "kind": kind,
+                        "m": m,
+                        "efConstruction": ef_construction,
+                        "similarity": similarity,
+                        "dimensions": dimensions,
+                    },
+                }
+            ],
+        }
+
+    @override
+    async def get_collections(self) -> list[str]:
+        return self.database.list_collection_names()
+
+    @override
+    async def delete_collection(self, collection_name: str) -> None:
+        return self.collection.drop()
+
+    @override
+    async def does_collection_exist(self, collection_name: str) -> bool:
+        return collection_name in self.database.list_collection_names()
+
+    @override
+    async def upsert(self, collection_name: str, record: MemoryRecord) -> str:
+        result = await self.upsert_batch(collection_name, [record])
+        return result[0]
+
+    @override
+    async def upsert_batch(self, collection_name: str, records: list[MemoryRecord]) -> list[str]:
+        doc_ids: list[str] = []
+        cosmosRecords: list[dict] = []
+        for record in records:
+            cosmosRecord: dict = {
+                "_id": record.id,
+                "embedding": record.embedding.tolist(),
+                "text": record.text,
+                "description": record.description,
+                "metadata": self.__serialize_metadata(record),
+            }
+            if record.timestamp is not None:
+                cosmosRecord["timestamp"] = record.timestamp
+
+            doc_ids.append(cosmosRecord["_id"])
+            cosmosRecords.append(cosmosRecord)
+        self.collection.insert_many(cosmosRecords)
+        return doc_ids
+
+    @override
+    async def get(self, collection_name: str, key: str, with_embedding: bool) -> MemoryRecord:
+        if not with_embedding:
+            result = self.collection.find_one({"_id": key}, {"embedding": 0})
+        else:
+            result = self.collection.find_one({"_id": key})
+        return MemoryRecord.local_record(
+            id=result["_id"],
+            embedding=np.array(result["embedding"]) if with_embedding else np.array([]),
+            text=result["text"],
+            description=result["description"],
+            additional_metadata=result["metadata"],
+            timestamp=result.get("timestamp", None),
+        )
+
+    @override
+    async def get_batch(self, collection_name: str, keys: list[str], with_embeddings: bool) -> list[MemoryRecord]:
+        if not with_embeddings:
+            results = self.collection.find({"_id": {"$in": keys}}, {"embedding": 0})
+        else:
+            results = self.collection.find({"_id": {"$in": keys}})
+
+        return [
+            MemoryRecord.local_record(
+                id=result["_id"],
+                embedding=np.array(result["embedding"]) if with_embeddings else np.array([]),
+                text=result["text"],
+                description=result["description"],
+                additional_metadata=result["metadata"],
+                timestamp=result.get("timestamp", None),
+            )
+            for result in results
+        ]
+
+    @override
+    async def remove(self, collection_name: str, key: str) -> None:
+        self.collection.delete_one({"_id": key})
+
+    @override
+    async def remove_batch(self, collection_name: str, keys: list[str]) -> None:
+        self.collection.delete_many({"_id": {"$in": keys}})
+
+    @override
+    async def get_nearest_matches(
+        self,
+        collection_name: str,
+        embedding: np.ndarray,
+        limit: int,
+        min_relevance_score: float,
+        with_embeddings: bool,
+    ) -> list[tuple[MemoryRecord, float]]:
+        pipeline: list[dict[str, Any]] = []
+        if self.kind == CosmosDBVectorSearchType.VECTOR_IVF:
+            pipeline = self._get_pipeline_vector_ivf(embedding.tolist(), limit)
+        elif self.kind == CosmosDBVectorSearchType.VECTOR_HNSW:
+            pipeline = self._get_pipeline_vector_hnsw(embedding.tolist(), limit, self.ef_search)
+
+        cursor = self.collection.aggregate(pipeline)
+
+        nearest_results = []
+        # Perform vector search
+        for aggResult in cursor:
+            score = aggResult["similarityScore"]
+            if score < min_relevance_score:
+                continue
+            result = MemoryRecord.local_record(
+                id=aggResult["_id"],
+                embedding=np.array(aggResult["document"]["embedding"]) if with_embeddings else np.array([]),
+                text=aggResult["document"]["text"],
+                description=aggResult["document"]["description"],
+                additional_metadata=aggResult["document"]["metadata"],
+                timestamp=aggResult["document"].get("timestamp", None),
+            )
+            nearest_results.append((result, aggResult["similarityScore"]))
+        return nearest_results
+
+    def _get_pipeline_vector_ivf(self, embeddings: list[float], k: int = 4) -> list[dict[str, Any]]:
+        pipeline: list[dict[str, Any]] = [
+            {
+                "$search": {
+                    "cosmosSearch": {
+                        "vector": embeddings,
+                        "path": "embedding",
+                        "k": k,
+                    },
+                    "returnStoredSource": True,
+                }
+            },
+            {
+                "$project": {
+                    "similarityScore": {"$meta": "searchScore"},
+                    "document": "$$ROOT",
+                }
+            },
+        ]
+        return pipeline
+
+    def _get_pipeline_vector_hnsw(
+        self, embeddings: list[float], k: int = 4, ef_search: int = 40
+    ) -> list[dict[str, Any]]:
+        pipeline: list[dict[str, Any]] = [
+            {
+                "$search": {
+                    "cosmosSearch": {
+                        "vector": embeddings,
+                        "path": "embedding",
+                        "k": k,
+                        "efSearch": ef_search,
+                    },
+                }
+            },
+            {
+                "$project": {
+                    "similarityScore": {"$meta": "searchScore"},
+                    "document": "$$ROOT",
+                }
+            },
+        ]
+        return pipeline
+
+    @override
+    async def get_nearest_match(
+        self,
+        collection_name: str,
+        embedding: np.ndarray,
+        min_relevance_score: float,
+        with_embedding: bool,
+    ) -> tuple[MemoryRecord, float]:
+        nearest_results = 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(nearest_results) > 0:
+            return nearest_results[0]
+        return None
+
+    @staticmethod
+    def __serialize_metadata(record: MemoryRecord) -> str:
+        return json.dumps({
+            "text": record.text,
+            "description": record.description,
+            "additional_metadata": record.additional_metadata,
+        })

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

@@ -0,0 +1,27 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from enum import Enum
+
+from semantic_kernel.utils.experimental_decorator import experimental_function
+
+
+@experimental_function
+class CosmosDBSimilarityType(str, Enum):
+    """Cosmos DB Similarity Type as enumerator."""
+
+    COS = "COS"
+    """CosineSimilarity"""
+    IP = "IP"
+    """inner - product"""
+    L2 = "L2"
+    """Euclidean distance"""
+
+
+@experimental_function
+class CosmosDBVectorSearchType(str, Enum):
+    """Cosmos DB Vector Search Type as enumerator."""
+
+    VECTOR_IVF = "vector-ivf"
+    """IVF vector index"""
+    VECTOR_HNSW = "vector-hnsw"
+    """HNSW vector index"""

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

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from semantic_kernel.connectors.memory.azure_cosmosdb_no_sql.azure_cosmosdb_no_sql_memory_store import (
+    AzureCosmosDBNoSQLMemoryStore,
+)
+
+__all__ = ["AzureCosmosDBNoSQLMemoryStore"]

semantic_kernel-1.8.3/semantic_kernel/connectors/memory/azure_cosmosdb_no_sql/azure_cosmosdb_no_sql_memory_store.py

@@ -0,0 +1,195 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import sys
+from typing import Any
+
+if sys.version_info >= (3, 12):
+    from typing import override  # pragma: no cover
+else:
+    from typing_extensions import override  # pragma: no cover
+
+import numpy as np
+from azure.cosmos.aio import ContainerProxy, CosmosClient, DatabaseProxy
+from numpy import ndarray
+
+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
+
+
+@experimental_class
+class AzureCosmosDBNoSQLMemoryStore(MemoryStoreBase):
+    """You can read more about vector search using AzureCosmosDBNoSQL here: https://aka.ms/CosmosVectorSearch."""
+
+    cosmos_client: CosmosClient = None
+    database: DatabaseProxy
+    container: ContainerProxy
+    database_name: str = None
+    partition_key: str = None
+    vector_embedding_policy: dict[str, Any] | None = None
+    indexing_policy: dict[str, Any] | None = None
+    cosmos_container_properties: dict[str, Any] | None = None
+
+    def __init__(
+        self,
+        cosmos_client: CosmosClient,
+        database_name: str,
+        partition_key: str,
+        vector_embedding_policy: dict[str, Any] | None = None,
+        indexing_policy: dict[str, Any] | None = None,
+        cosmos_container_properties: dict[str, Any] | None = None,
+    ):
+        """Initializes a new instance of the AzureCosmosDBNoSQLMemoryStore class."""
+        if indexing_policy["vectorIndexes"] is None or len(indexing_policy["vectorIndexes"]) == 0:
+            raise ValueError("vectorIndexes cannot be null or empty in the indexing_policy.")
+        if vector_embedding_policy is None or len(vector_embedding_policy["vectorEmbeddings"]) == 0:
+            raise ValueError("vectorEmbeddings cannot be null or empty in the vector_embedding_policy.")
+
+        self.cosmos_client = cosmos_client
+        self.database_name = database_name
+        self.partition_key = partition_key
+        self.vector_embedding_policy = vector_embedding_policy
+        self.indexing_policy = indexing_policy
+        self.cosmos_container_properties = cosmos_container_properties
+
+    @override
+    async def create_collection(self, collection_name: str) -> None:
+        # Create the database if it already doesn't exist
+        self.database = await self.cosmos_client.create_database_if_not_exists(id=self.database_name)
+
+        # Create the collection if it already doesn't exist
+        self.container = await self.database.create_container_if_not_exists(
+            id=collection_name,
+            partition_key=self.cosmos_container_properties["partition_key"],
+            indexing_policy=self.indexing_policy,
+            vector_embedding_policy=self.vector_embedding_policy,
+        )
+
+    @override
+    async def get_collections(self) -> list[str]:
+        return [container["id"] async for container in self.database.list_containers()]
+
+    @override
+    async def delete_collection(self, collection_name: str) -> None:
+        return await self.database.delete_container(collection_name)
+
+    @override
+    async def does_collection_exist(self, collection_name: str) -> bool:
+        return collection_name in [container["id"] async for container in self.database.list_containers()]
+
+    @override
+    async def upsert(self, collection_name: str, record: MemoryRecord) -> str:
+        result = await self.upsert_batch(collection_name, [record])
+        return result[0]
+
+    @override
+    async def upsert_batch(self, collection_name: str, records: list[MemoryRecord]) -> list[str]:
+        doc_ids: list[str] = []
+        for record in records:
+            cosmosRecord: dict = {
+                "id": record.id,
+                "embedding": record.embedding.tolist(),
+                "text": record.text,
+                "description": record.description,
+                "metadata": self.__serialize_metadata(record),
+            }
+            if record.timestamp is not None:
+                cosmosRecord["timeStamp"] = record.timestamp
+
+            await self.container.create_item(cosmosRecord)
+            doc_ids.append(cosmosRecord["id"])
+        return doc_ids
+
+    @override
+    async def get(self, collection_name: str, key: str, with_embedding: bool) -> MemoryRecord:
+        item = await self.container.read_item(key, partition_key=key)
+        return MemoryRecord.local_record(
+            id=item["id"],
+            embedding=np.array(item["embedding"]) if with_embedding else np.array([]),
+            text=item["text"],
+            description=item["description"],
+            additional_metadata=item["metadata"],
+            timestamp=item.get("timestamp", None),
+        )
+
+    @override
+    async def get_batch(self, collection_name: str, keys: list[str], with_embeddings: bool) -> list[MemoryRecord]:
+        query = "SELECT * FROM c WHERE ARRAY_CONTAINS(@ids, c.id)"
+        parameters = [{"name": "@ids", "value": keys}]
+
+        all_results = []
+        items = [item async for item in self.container.query_items(query, parameters=parameters)]
+        for item in items:
+            MemoryRecord.local_record(
+                id=item["id"],
+                embedding=np.array(item["embedding"]) if with_embeddings else np.array([]),
+                text=item["text"],
+                description=item["description"],
+                additional_metadata=item["metadata"],
+                timestamp=item.get("timestamp", None),
+            )
+            all_results.append(item)
+        return all_results
+
+    @override
+    async def remove(self, collection_name: str, key: str) -> None:
+        await self.container.delete_item(key, partition_key=key)
+
+    @override
+    async def remove_batch(self, collection_name: str, keys: list[str]) -> None:
+        for key in keys:
+            await self.container.delete_item(key, partition_key=key)
+
+    @override
+    async def get_nearest_matches(
+        self, collection_name: str, embedding: ndarray, limit: int, min_relevance_score: float, with_embeddings: bool
+    ) -> list[tuple[MemoryRecord, float]]:
+        embedding_key = self.vector_embedding_policy["vectorEmbeddings"][0]["path"][1:]
+        query = (
+            f"SELECT TOP {limit} c.id, c.{embedding_key}, c.text, c.description, c.metadata, "  # nosec
+            f"c.timestamp, VectorDistance(c.{embedding_key}, {embedding.tolist()}) AS SimilarityScore FROM c ORDER BY "  # nosec
+            f"VectorDistance(c.{embedding_key}, {embedding.tolist()})"  # nosec
+        )
+
+        items = [item async for item in self.container.query_items(query=query)]
+        nearest_results = []
+        for item in items:
+            score = item["SimilarityScore"]
+            if score < min_relevance_score:
+                continue
+            result = MemoryRecord.local_record(
+                id=item["id"],
+                embedding=np.array(item["embedding"]) if with_embeddings else np.array([]),
+                text=item["text"],
+                description=item["description"],
+                additional_metadata=item["metadata"],
+                timestamp=item.get("timestamp", None),
+            )
+            nearest_results.append((result, score))
+        return nearest_results
+
+    @override
+    async def get_nearest_match(
+        self, collection_name: str, embedding: ndarray, min_relevance_score: float, with_embedding: bool
+    ) -> tuple[MemoryRecord, float]:
+        nearest_results = await self.get_nearest_matches(
+            collection_name=collection_name,
+            embedding=embedding,
+            limit=1,
+            min_relevance_score=min_relevance_score,
+            with_embeddings=with_embedding,
+        )
+        if len(nearest_results) > 0:
+            return nearest_results[0]
+        return None
+
+    @staticmethod
+    def __serialize_metadata(record: MemoryRecord) -> str:
+        return json.dumps(
+            {
+                "text": record.text,
+                "description": record.description,
+                "additional_metadata": record.additional_metadata,
+            }
+        )

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

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from semantic_kernel.connectors.memory.chroma.chroma_memory_store import (
+    ChromaMemoryStore,
+)
+
+__all__ = ["ChromaMemoryStore"]