chore: import upstream snapshot with attribution
CodeQL / Analyze (csharp) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 13:21:23 +08:00
commit b957a53def
5423 changed files with 863745 additions and 0 deletions
@@ -0,0 +1,182 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import aiohttp
from semantic_kernel.connectors.memory_stores.astradb.utils import AsyncSession
from semantic_kernel.exceptions import ServiceResponseException
from semantic_kernel.utils.feature_stage_decorator import experimental
from semantic_kernel.utils.telemetry.user_agent import APP_INFO
ASTRA_CALLER_IDENTITY: str
SEMANTIC_KERNEL_VERSION = APP_INFO.get("Semantic-Kernel-Version")
ASTRA_CALLER_IDENTITY = f"semantic-kernel/{SEMANTIC_KERNEL_VERSION}" if SEMANTIC_KERNEL_VERSION else "semantic-kernel"
@experimental
class AstraClient:
"""AstraClient."""
def __init__(
self,
astra_id: str,
astra_region: str,
astra_application_token: str,
keyspace_name: str,
embedding_dim: int,
similarity_function: str,
session: aiohttp.ClientSession | None = None,
):
"""Initializes a new instance of the AstraClient class."""
self.astra_id = astra_id
self.astra_application_token = astra_application_token
self.astra_region = astra_region
self.keyspace_name = keyspace_name
self.embedding_dim = embedding_dim
self.similarity_function = similarity_function
self.request_base_url = (
f"https://{self.astra_id}-{self.astra_region}.apps.astra.datastax.com/api/json/v1/{self.keyspace_name}"
)
self.request_header = {
"x-cassandra-token": self.astra_application_token,
"Content-Type": "application/json",
"User-Agent": ASTRA_CALLER_IDENTITY,
}
self._session = session
async def _run_query(self, request_url: str, query: dict):
async with (
AsyncSession(self._session) as session,
session.post(request_url, data=json.dumps(query), headers=self.request_header) as response,
):
if response.status == 200:
response_dict = await response.json()
if "errors" in response_dict:
raise ServiceResponseException(f"Astra DB request error - {response_dict['errors']}")
return response_dict
raise ServiceResponseException(f"Astra DB not available. Status : {response}")
async def find_collections(self, include_detail: bool = True):
"""Finds all collections in the keyspace."""
query = {"findCollections": {"options": {"explain": include_detail}}}
result = await self._run_query(self.request_base_url, query)
return result["status"]["collections"]
async def find_collection(self, collection_name: str):
"""Finds a collection in the keyspace."""
collections = await self.find_collections(False)
found = False
for collection in collections:
if collection == collection_name:
found = True
break
return found
async def create_collection(
self,
collection_name: str,
embedding_dim: int | None = None,
similarity_function: str | None = None,
):
"""Creates a new collection in the keyspace."""
query = {
"createCollection": {
"name": collection_name,
"options": {
"vector": {
"dimension": embedding_dim if embedding_dim is not None else self.embedding_dim,
"metric": similarity_function if similarity_function is not None else self.similarity_function,
}
},
}
}
result = await self._run_query(self.request_base_url, query)
return result["status"]["ok"] == 1
async def delete_collection(self, collection_name: str):
"""Deletes a collection from the keyspace."""
query = {"deleteCollection": {"name": collection_name}}
result = await self._run_query(self.request_base_url, query)
return result["status"]["ok"] == 1
def _build_request_collection_url(self, collection_name: str):
return f"{self.request_base_url}/{collection_name}"
async def find_documents(
self,
collection_name: str,
filter: dict | None = None,
vector: list[float] | None = None,
limit: int | None = None,
include_vector: bool | None = None,
include_similarity: bool | None = None,
) -> list[dict]:
"""Finds all documents in the collection."""
find_query = {}
if filter is not None:
find_query["filter"] = filter
if vector is not None:
find_query["sort"] = {"$vector": vector}
if include_vector is not None and include_vector is False:
find_query["projection"] = {"$vector": 0}
else:
find_query["projection"] = {"*": 1}
if limit is not None:
find_query["options"] = {"limit": limit}
if include_similarity is not None:
if "options" in find_query:
find_query["options"]["includeSimilarity"] = int(include_similarity)
else:
find_query["options"] = {"includeSimilarity": int(include_similarity)}
query = {"find": find_query}
result = await self._run_query(self._build_request_collection_url(collection_name), query)
return result["data"]["documents"]
async def insert_document(self, collection_name: str, document: dict) -> str:
"""Inserts a document into the collection."""
query = {"insertOne": {"document": document}}
result = await self._run_query(self._build_request_collection_url(collection_name), query)
return result["status"]["insertedIds"][0]
async def insert_documents(self, collection_name: str, documents: list[dict]) -> list[str]:
"""Inserts multiple documents into the collection."""
query = {"insertMany": {"documents": documents}}
result = await self._run_query(self._build_request_collection_url(collection_name), query)
return result["status"]["insertedIds"]
async def update_document(self, collection_name: str, filter: dict, update: dict, upsert: bool = True) -> dict:
"""Updates a document in the collection."""
query = {
"findOneAndUpdate": {
"filter": filter,
"update": update,
"options": {"returnDocument": "after", "upsert": upsert},
}
}
result = await self._run_query(self._build_request_collection_url(collection_name), query)
return result["status"]
async def update_documents(self, collection_name: str, filter: dict, update: dict):
"""Updates multiple documents in the collection."""
query = {
"updateMany": {
"filter": filter,
"update": update,
}
}
result = await self._run_query(self._build_request_collection_url(collection_name), query)
return result["status"]
async def delete_documents(self, collection_name: str, filter: dict) -> int:
"""Deletes documents from the collection."""
query = {"deleteMany": {"filter": filter}}
result = await self._run_query(self._build_request_collection_url(collection_name), query)
return result["status"]["deletedCount"]
@@ -0,0 +1,335 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import logging
import sys
import aiohttp
from numpy import ndarray
from pydantic import ValidationError
from semantic_kernel.connectors.memory_stores.astradb.astra_client import AstraClient
from semantic_kernel.connectors.memory_stores.astradb.astradb_settings import AstraDBSettings
from semantic_kernel.connectors.memory_stores.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.feature_stage_decorator import experimental
if sys.version_info >= (3, 13):
from warnings import deprecated
else:
from typing_extensions import deprecated
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__)
@deprecated("This class has been deprecated and will be removed in a future release. ")
@experimental
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(
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 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_similarity",
) -> 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 []
)
@@ -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.feature_stage_decorator import experimental
@experimental
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
@@ -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([]),
)
@@ -0,0 +1,422 @@
# Copyright (c) Microsoft. All rights reserved.
import contextlib
import logging
import sys
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_stores.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
if sys.version_info >= (3, 13):
from warnings import deprecated
else:
from typing_extensions import deprecated
logger: logging.Logger = logging.getLogger(__name__)
@deprecated("This class will be removed in a future version. Use AzureAISearchStore and Collection instead.")
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.azure_ai_search import AzureAISearchSettings
try:
acs_memory_settings = AzureAISearchSettings(
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_distance"
),
)
],
)
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
@@ -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")
@@ -0,0 +1,297 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
import sys
from typing import Literal
from numpy import ndarray
from pymongo import MongoClient
from semantic_kernel.connectors.azure_cosmos_db import AzureCosmosDBforMongoDBSettings
from semantic_kernel.connectors.memory_stores.azure_cosmosdb.azure_cosmos_db_store_api import AzureCosmosDBStoreApi
from semantic_kernel.connectors.memory_stores.azure_cosmosdb.mongo_vcore_store_api import MongoStoreApi
from semantic_kernel.connectors.memory_stores.azure_cosmosdb.utils import (
CosmosDBSimilarityType,
CosmosDBVectorSearchType,
)
from semantic_kernel.exceptions import MemoryConnectorInitializationError
from semantic_kernel.memory.memory_record import MemoryRecord
from semantic_kernel.memory.memory_store_base import MemoryStoreBase
if sys.version_info >= (3, 13):
from warnings import deprecated
else:
from typing_extensions import deprecated
logger: logging.Logger = logging.getLogger(__name__)
@deprecated("This class will be removed in a future release, use AzureCosmosDBforMongoDBStore and Collection instead.")
class AzureCosmosDBMemoryStore(MemoryStoreBase):
"""A memory store that uses AzureCosmosDB for MongoDB vCore.
To perform vector similarity search on a fully managed MongoDB compatible database service.
https://learn.microsoft.com/en-us/azure/cosmos-db/mongodb/vcore/vector-search.
"""
# Right now this only supports Mongo, but set up to support more later.
api_store: AzureCosmosDBStoreApi = None
mongodb_client = None
database = None
index_name = None
vector_dimensions = None
num_lists = None
similarity = None
collection_name = None
kind = None
m = None
ef_construction = None
ef_search = None
def __init__(
self,
cosmosStore: AzureCosmosDBStoreApi,
database_name: str,
index_name: str,
vector_dimensions: int,
num_lists: int = 100,
similarity: CosmosDBSimilarityType = CosmosDBSimilarityType.COS,
kind: CosmosDBVectorSearchType = CosmosDBVectorSearchType.VECTOR_HNSW,
m: int = 16,
ef_construction: int = 64,
ef_search: int = 40,
):
"""Initializes a new instance of the AzureCosmosDBMemoryStore class."""
if vector_dimensions <= 0:
raise MemoryConnectorInitializationError("Vector dimensions must be a positive number.")
if database_name is None:
raise MemoryConnectorInitializationError("Database Name cannot be empty.")
if index_name is None:
raise MemoryConnectorInitializationError("Index Name cannot be empty.")
self.cosmos_store = cosmosStore
self.index_name = index_name
self.num_lists = num_lists
self.similarity = similarity
self.kind = kind
self.m = m
self.ef_construction = ef_construction
self.ef_search = ef_search
@staticmethod
async def create(
database_name: str,
collection_name: str,
vector_dimensions: int,
num_lists: int,
similarity: CosmosDBSimilarityType,
kind: CosmosDBVectorSearchType,
m: int,
ef_construction: int,
ef_search: int,
index_name: str | None = None,
cosmos_connstr: str | None = None,
application_name: str | None = None,
cosmos_api: Literal["mongo-vcore"] = "mongo-vcore",
env_file_path: str | None = None,
) -> MemoryStoreBase:
"""Creates the underlying data store based on the API definition."""
# Right now this only supports Mongo, but set up to support more later.
api_store: AzureCosmosDBStoreApi = None
if cosmos_api == "mongo-vcore":
cosmosdb_settings = AzureCosmosDBforMongoDBSettings(
env_file_path=env_file_path,
connection_string=cosmos_connstr,
)
mongodb_client = MongoClient(
cosmosdb_settings.connection_string.get_secret_value() if cosmosdb_settings.connection_string else None,
appname=application_name,
)
database = mongodb_client[database_name]
api_store = MongoStoreApi(
collection_name=collection_name,
index_name=index_name,
vector_dimensions=vector_dimensions,
num_lists=num_lists,
similarity=similarity,
database=database,
kind=kind,
m=m,
ef_construction=ef_construction,
ef_search=ef_search,
)
else:
raise MemoryConnectorInitializationError(f"API type {cosmos_api} is not supported.")
store = AzureCosmosDBMemoryStore(
api_store,
database_name,
index_name,
vector_dimensions,
num_lists,
similarity,
kind,
m,
ef_construction,
ef_search,
)
await store.create_collection(collection_name)
return store
async def create_collection(self, collection_name: str) -> None:
"""Creates a new collection in the data store.
Args:
collection_name (str): The name associated with a collection of embeddings.
Returns:
None
"""
return await self.cosmos_store.create_collection(collection_name)
async def get_collections(self) -> list[str]:
"""Gets the list of collections.
Returns:
List[str]: The list of collections.
"""
return await self.cosmos_store.get_collections()
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
"""
return await self.cosmos_store.delete_collection("")
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.cosmos_store.does_collection_exist("")
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.
"""
return await self.cosmos_store.upsert("", record)
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.
"""
return await self.cosmos_store.upsert_batch("", records)
async def get(self, collection_name: str, key: str, with_embedding: bool) -> 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.
"""
return await self.cosmos_store.get("", key, with_embedding)
async def get_batch(self, collection_name: str, keys: list[str], with_embeddings: bool) -> 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.
"""
return await self.cosmos_store.get_batch("", keys, with_embeddings)
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
"""
return await self.cosmos_store.remove("", key)
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
"""
return await self.cosmos_store.remove_batch("", keys)
async def get_nearest_matches(
self,
collection_name: str,
embedding: ndarray,
limit: int,
min_relevance_score: float,
with_embeddings: bool,
) -> 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.
"""
return await self.cosmos_store.get_nearest_matches("", embedding, limit, min_relevance_score, with_embeddings)
async def get_nearest_match(
self,
collection_name: str,
embedding: ndarray,
min_relevance_score: float,
with_embedding: bool,
) -> 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.
"""
return await self.cosmos_store.get_nearest_match("", embedding, min_relevance_score, with_embedding)
@@ -0,0 +1,188 @@
# Copyright (c) Microsoft. All rights reserved.
import sys
from abc import ABC, abstractmethod
from numpy import ndarray
from semantic_kernel.memory.memory_record import MemoryRecord
if sys.version_info >= (3, 13):
from warnings import deprecated
else:
from typing_extensions import deprecated
# Abstract class similar to the original data store that allows API level abstraction
@deprecated(
"This class will be removed in a future release, use the AzureCosmosDBNoSQLStore and "
"Collection or AzureCosmosDBMongoDBStore and collection instead."
)
class AzureCosmosDBStoreApi(ABC):
"""AzureCosmosDBStoreApi."""
@abstractmethod
async def create_collection(self, collection_name: str) -> None:
"""Creates a new collection in the data store.
Args:
collection_name (str): The name associated with a collection of embeddings.
"""
raise NotImplementedError
@abstractmethod
async def get_collections(self) -> list[str]:
"""Gets all collection names in the data store.
Returns:
List[str]: A group of collection names.
"""
raise NotImplementedError
@abstractmethod
async def delete_collection(self, collection_name: str) -> None:
"""Deletes a collection from the data store.
Args:
collection_name (str): The name associated with a collection of embeddings.
"""
raise NotImplementedError
@abstractmethod
async def does_collection_exist(self, collection_name: str) -> bool:
"""Determines if a collection exists in the data store.
Args:
collection_name (str): The name associated with a collection of embeddings.
Returns:
bool: True if given collection exists, False if not.
"""
raise NotImplementedError
@abstractmethod
async def upsert(self, collection_name: str, record: MemoryRecord) -> str:
"""Upserts 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.
"""
raise NotImplementedError
@abstractmethod
async def upsert_batch(self, collection_name: str, records: list[MemoryRecord]) -> list[str]:
"""Upserts a group 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 (MemoryRecord): The memory records to upsert.
Returns:
List[str]: The unique identifiers for the memory records.
"""
raise NotImplementedError
@abstractmethod
async def get(self, collection_name: str, key: str, with_embedding: bool) -> MemoryRecord:
"""Gets a memory record from the data store. Does not guarantee that the collection exists.
Args:
collection_name (str): The name associated with a collection of embeddings.
key (str): The unique id associated with the memory record to get.
with_embedding (bool): If true, the embedding will be returned in the memory record.
Returns:
MemoryRecord: The memory record if found
"""
raise NotImplementedError
@abstractmethod
async def get_batch(self, collection_name: str, keys: list[str], with_embeddings: bool) -> list[MemoryRecord]:
"""Gets a batch of memory records from the data store. Does not guarantee that the collection exists.
Args:
collection_name (str): The name associated with a collection of embeddings.
keys (List[str]): The unique ids associated with the memory records to get.
with_embeddings (bool): If true, the embedding will be returned in the memory records.
Returns:
List[MemoryRecord]: The memory records associated with the unique keys provided.
"""
raise NotImplementedError
@abstractmethod
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 associated with a collection of embeddings.
key (str): The unique id associated with the memory record to remove.
"""
raise NotImplementedError
@abstractmethod
async def remove_batch(self, collection_name: str, keys: list[str]) -> None:
"""Removes a batch of memory records from the data store. Does not guarantee that the collection exists.
Args:
collection_name (str): The name associated with a collection of embeddings.
keys (List[str]): The unique ids associated with the memory records to remove.
"""
raise NotImplementedError
@abstractmethod
async def get_nearest_matches(
self,
collection_name: str,
embedding: ndarray,
limit: int,
min_relevance_score: float,
with_embeddings: bool,
) -> list[tuple[MemoryRecord, float]]:
"""Gets the nearest matches to an embedding of type float. Does not guarantee that the collection exists.
Args:
collection_name (str): The name associated with a collection of embeddings.
embedding (ndarray): The embedding to compare the collection's embeddings with.
limit (int): The maximum number of similarity results to return.
min_relevance_score (float): The minimum relevance threshold for returned results.
with_embeddings (bool): If true, the embeddings will be returned in the memory records.
Returns:
List[Tuple[MemoryRecord, float]]: A list of tuples where item1 is a MemoryRecord and item2
is its similarity score as a float.
"""
raise NotImplementedError
@abstractmethod
async def get_nearest_match(
self,
collection_name: str,
embedding: ndarray,
min_relevance_score: float,
with_embedding: bool,
) -> tuple[MemoryRecord, float]:
"""Gets the nearest match to an embedding of type float. Does not guarantee that the collection exists.
Args:
collection_name (str): The name associated with a collection of embeddings.
embedding (ndarray): The embedding to compare the collection's embeddings with.
min_relevance_score (float): The minimum relevance threshold for returned result.
with_embedding (bool): If true, the embeddings will be returned in the memory record.
Returns:
Tuple[MemoryRecord, float]: A tuple consisting of the MemoryRecord and the similarity score as a float.
"""
raise NotImplementedError
@@ -0,0 +1,353 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import sys
from typing import Any
import numpy as np
from semantic_kernel.connectors.memory_stores.azure_cosmosdb.azure_cosmos_db_store_api import AzureCosmosDBStoreApi
from semantic_kernel.connectors.memory_stores.azure_cosmosdb.utils import (
CosmosDBSimilarityType,
CosmosDBVectorSearchType,
)
from semantic_kernel.memory.memory_record import MemoryRecord
if sys.version >= "3.12":
from typing import override # pragma: no cover
else:
from typing_extensions import override # pragma: no cover
if sys.version_info >= (3, 13):
from warnings import deprecated
else:
from typing_extensions import deprecated
@deprecated("This class will be removed in a future release.")
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,
})
@@ -0,0 +1,27 @@
# Copyright (c) Microsoft. All rights reserved.
from enum import Enum
from semantic_kernel.utils.feature_stage_decorator import experimental
@experimental
class CosmosDBSimilarityType(str, Enum):
"""Cosmos DB Similarity Type as enumerator."""
COS = "COS"
"""CosineSimilarity"""
IP = "IP"
"""inner - product"""
L2 = "L2"
"""Euclidean distance"""
@experimental
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"""
@@ -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"]
@@ -0,0 +1,197 @@
# 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
if sys.version_info >= (3, 13):
from warnings import deprecated
else:
from typing_extensions import deprecated
@deprecated("This class will be removed in a future version. Use AzureComosDBNoSQLStore and Collection instead.")
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,
})
@@ -0,0 +1,351 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
import sys
from typing import TYPE_CHECKING, Any, Optional
from numpy import array, ndarray
if sys.version_info >= (3, 12):
from typing import override # pragma: no cover
else:
from typing_extensions import override # pragma: no cover
from semantic_kernel.connectors.memory_stores.chroma.utils import (
chroma_compute_similarity_scores,
query_results_to_records,
)
from semantic_kernel.exceptions import ServiceInitializationError, ServiceResourceNotFoundError
from semantic_kernel.memory.memory_record import MemoryRecord
from semantic_kernel.memory.memory_store_base import MemoryStoreBase
if TYPE_CHECKING:
import chromadb
import chromadb.config
from chromadb.api.models.Collection import Collection
if sys.version_info >= (3, 12):
from warnings import deprecated
else:
from typing_extensions import deprecated
logger: logging.Logger = logging.getLogger(__name__)
@deprecated(
"ChromaMemoryStore is deprecated and will be removed in a future version. Use ChromaStore and Collection instead."
)
class ChromaMemoryStore(MemoryStoreBase):
"""ChromaMemoryStore provides an interface to store and retrieve data using ChromaDB."""
_client: "chromadb.Client"
def __init__(
self,
persist_directory: str | None = None,
client_settings: Optional["chromadb.config.Settings"] = None,
**kwargs: Any,
) -> None:
"""ChromaMemoryStore provides an interface to store and retrieve data using ChromaDB.
Collection names with uppercase characters are not supported by ChromaDB, they will be automatically converted.
Args:
persist_directory (Optional[str], optional): Path to the directory where data will be persisted.
Defaults to None, which means the default settings for ChromaDB will be used.
client_settings (Optional["chromadb.config.Settings"], optional): A Settings instance to configure
the ChromaDB client. Defaults to None, which means the default settings for ChromaDB will be used.
similarity_fetch_limit (int, optional): The maximum number of results to calculate cosine-similarity.
**kwargs: Additional keyword arguments.
Example:
# Create a ChromaMemoryStore with a local specified directory for data persistence
chroma_local_data_store = ChromaMemoryStore(persist_directory='/path/to/persist/directory')
# Create a ChromaMemoryStore with a custom Settings instance
chroma_remote_data_store = ChromaMemoryStore(
client_settings=Settings(
chroma_api_impl="rest",
chroma_server_host="xxx.xxx.xxx.xxx",
chroma_server_http_port="8000"
)
)
"""
try:
import chromadb
import chromadb.config
except ImportError as exc:
raise ServiceInitializationError(
"Could not import chromadb python package. Please install it with `pip install chromadb`."
) from exc
if client_settings:
self._client_settings = client_settings
else:
self._client_settings = chromadb.config.Settings()
if persist_directory is not None:
self._client_settings = chromadb.config.Settings(
is_persistent=True, persist_directory=persist_directory
)
self._client = chromadb.Client(self._client_settings)
self._persist_directory = persist_directory
self._default_query_includes = ["embeddings", "metadatas", "documents"]
async def create_collection(self, collection_name: str) -> None:
"""Creates a new collection in Chroma if it does not exist.
To prevent downloading model file from embedding_function,
embedding_function is set to "DoNotUseChromaEmbeddingFunction".
Args:
collection_name (str): The name of the collection to create.
The name of the collection will be converted to snake case.
Returns:
None
"""
self._client.create_collection(name=collection_name)
@override
async def get_collection(self, collection_name: str) -> Optional["Collection"]:
try:
# Current version of ChromeDB rejects camel case collection names.
return self._client.get_collection(name=collection_name)
except ValueError:
return None
async def get_collections(self) -> list[str]:
"""Gets the list of collections.
Returns:
List[str]: The list of collections.
"""
return [collection.name for collection in self._client.list_collections()]
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
"""
self._client.delete_collection(name=collection_name)
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.get_collection(collection_name) is not None
async def upsert(self, collection_name: str, record: MemoryRecord) -> str:
"""Upsert a single MemoryRecord.
Args:
collection_name (str): The name of the collection to upsert the record into.
record (MemoryRecord): The record to upsert.
Returns:
List[str]: The unique database key of the record.
"""
collection = await self.get_collection(collection_name)
if collection is None:
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
record._key = record._id
metadata = {
"timestamp": record._timestamp or "",
"is_reference": str(record._is_reference),
"external_source_name": record._external_source_name or "",
"description": record._description or "",
"additional_metadata": record._additional_metadata or "",
"id": record._id or "",
}
collection.add(
metadatas=metadata,
# by providing embeddings, we can skip the chroma's embedding function call
embeddings=record.embedding.tolist(),
documents=record._text,
ids=record._key,
)
return record._key
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. In Pinecone, these are the record IDs.
"""
# upsert is checking collection existence
return [await 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.
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.
"""
records = await self.get_batch(collection_name, [key], with_embedding)
try:
return records[0]
except IndexError as exc:
raise ServiceResourceNotFoundError(
f"Record with key '{key}' does not exist in collection '{collection_name}'"
) from exc
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.
"""
collection = await self.get_collection(collection_name)
if collection is None:
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
query_includes = ["embeddings", "metadatas", "documents"] if with_embeddings else ["metadatas", "documents"]
value = collection.get(ids=keys, include=query_includes)
return query_results_to_records(value, with_embeddings)
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
"""
await self.remove_batch(collection_name, [key])
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
"""
collection = await self.get_collection(collection_name=collection_name)
if collection is not None:
collection.delete(ids=keys)
async def get_nearest_matches(
self,
collection_name: str,
embedding: ndarray,
limit: int,
min_relevance_score: float = 0.0,
with_embeddings: bool = True,
) -> 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.
"""
if with_embeddings is False:
logger.warning(
"Chroma returns distance score not cosine similarity score.\
So embeddings are automatically queried from database for calculation."
)
collection = await self.get_collection(collection_name)
if collection is None:
return []
query_results = collection.query(
query_embeddings=embedding.tolist(),
n_results=limit,
include=self._default_query_includes,
)
# Convert the collection of embeddings into a numpy array (stacked)
embedding_array = array(query_results["embeddings"][0])
embedding_array = embedding_array.reshape(embedding_array.shape[0], -1)
# If the query embedding has shape (1, embedding_size),
# reshape it to (embedding_size,)
if len(embedding.shape) == 2:
embedding = embedding.reshape(
embedding.shape[1],
)
similarity_score = chroma_compute_similarity_scores(embedding, embedding_array)
# Convert query results into memory records
record_list = [
(record, distance)
for record, distance in zip(
query_results_to_records(query_results, with_embeddings),
similarity_score,
)
]
sorted_results = sorted(
record_list,
key=lambda x: x[1],
reverse=True,
)
filtered_results = [x for x in sorted_results if x[1] >= min_relevance_score]
return filtered_results[:limit]
async def get_nearest_match(
self,
collection_name: str,
embedding: ndarray,
min_relevance_score: float = 0.0,
with_embedding: bool = True,
) -> 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 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.
"""
results = await self.get_nearest_matches(
collection_name=collection_name,
embedding=embedding,
limit=1,
min_relevance_score=min_relevance_score,
with_embeddings=with_embedding,
)
return results[0]
@@ -0,0 +1,124 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
from typing import TYPE_CHECKING, Any
from numpy import array, linalg, ndarray
from semantic_kernel.memory.memory_record import MemoryRecord
if TYPE_CHECKING:
from chromadb.api.types import QueryResult
logger: logging.Logger = logging.getLogger(__name__)
def camel_to_snake(camel_str):
"""Convert camel case to snake case."""
snake_str = ""
for i, char in enumerate(camel_str):
if char.isupper():
if i != 0 and camel_str[i - 1].islower():
snake_str += "_"
if i != len(camel_str) - 1 and camel_str[i + 1].islower():
snake_str += "_"
snake_str += char.lower()
return snake_str
def query_results_to_records(results: "QueryResult", with_embedding: bool) -> list[MemoryRecord]:
"""Turn query results into Memory Records.
If results has only one record, it will be a list instead of a nested list
this is to make sure that results is always a nested list
{'ids': ['test_id1'], 'embeddings': [[...]], 'documents': ['sample text1'], 'metadatas': [{...}]}
=> {'ids': [['test_id1']], 'embeddings': [[[...]]], 'documents': [['sample text1']], 'metadatas': [[{...}]]}
"""
try:
if isinstance(results["ids"][0], str):
for k, v in results.items():
results[k] = [v]
except IndexError:
return []
if with_embedding:
memory_records = [
(
MemoryRecord(
is_reference=(metadata["is_reference"] == "True"),
external_source_name=metadata["external_source_name"],
id=metadata["id"],
description=metadata["description"],
text=document,
embedding=embedding,
additional_metadata=metadata["additional_metadata"],
key=id,
timestamp=metadata["timestamp"],
)
)
for id, document, embedding, metadata in zip(
results["ids"][0],
results["documents"][0],
results["embeddings"][0],
results["metadatas"][0],
)
]
else:
memory_records = [
(
MemoryRecord(
is_reference=(metadata["is_reference"] == "True"),
external_source_name=metadata["external_source_name"],
id=metadata["id"],
description=metadata["description"],
text=document,
embedding=None,
additional_metadata=metadata["additional_metadata"],
key=id,
timestamp=metadata["timestamp"],
)
)
for id, document, metadata in zip(
results["ids"][0],
results["documents"][0],
results["metadatas"][0],
)
]
return memory_records
def chroma_compute_similarity_scores(embedding: ndarray, embedding_array: ndarray, **kwargs: Any) -> ndarray:
"""Computes the cosine similarity scores between a query embedding and a group of embeddings.
Args:
embedding (ndarray): The query embedding.
embedding_array (ndarray): The group of embeddings.
**kwargs: Additional keyword arguments.
Returns:
ndarray: The cosine similarity scores.
"""
query_norm = linalg.norm(embedding)
collection_norm = linalg.norm(embedding_array, axis=1)
# Compute indices for which the similarity scores can be computed
valid_indices = (query_norm != 0) & (collection_norm != 0)
# Initialize the similarity scores with -1 to distinguish the cases
# between zero similarity from orthogonal vectors and invalid similarity
similarity_scores = array([-1.0] * embedding_array.shape[0])
if valid_indices.any():
similarity_scores[valid_indices] = embedding.dot(embedding_array[valid_indices].T) / (
query_norm * collection_norm[valid_indices]
)
if not valid_indices.all():
logger.warning(
"Some vectors in the embedding collection are zero vectors."
"Ignoring cosine similarity score computation for those vectors."
)
else:
raise ValueError(
f"Invalid vectors, cannot compute cosine similarity scoresfor zero vectors{embedding_array} or {embedding}"
)
return similarity_scores
@@ -0,0 +1,463 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
import sys
from datetime import datetime
from typing import Any
from numpy import array, expand_dims, ndarray
from pymilvus import Collection, CollectionSchema, DataType, FieldSchema, connections, utility
from semantic_kernel.exceptions import ServiceResourceNotFoundError, ServiceResponseException
from semantic_kernel.memory.memory_record import MemoryRecord
from semantic_kernel.memory.memory_store_base import MemoryStoreBase
from semantic_kernel.utils.feature_stage_decorator import experimental
if sys.version_info >= (3, 13):
from warnings import deprecated
else:
from typing_extensions import deprecated
logger: logging.Logger = logging.getLogger(__name__)
# Index parameters
_INDEX_TYPE = "IVF_FLAT"
_NLIST = 1024
SEARCH_FIELD_ID = "id"
SEARCH_FIELD_TEXT = "text"
SEARCH_FIELD_EMBEDDING = "embedding"
SEARCH_FIELD_SRC = "external_source_name"
SEARCH_FIELD_DESC = "description"
SEARCH_FIELD_METADATA = "additional_metadata"
SEARCH_FIELD_IS_REF = "is_reference"
SEARCH_FIELD_TIMESTAMP = "timestamp"
OUTPUT_FIELDS_W_EMBEDDING = [
SEARCH_FIELD_ID,
SEARCH_FIELD_TEXT,
SEARCH_FIELD_SRC,
SEARCH_FIELD_DESC,
SEARCH_FIELD_METADATA,
SEARCH_FIELD_IS_REF,
SEARCH_FIELD_EMBEDDING,
SEARCH_FIELD_TIMESTAMP,
]
OUTPUT_FIELDS_WO_EMBEDDING = [
SEARCH_FIELD_ID,
SEARCH_FIELD_TEXT,
SEARCH_FIELD_SRC,
SEARCH_FIELD_DESC,
SEARCH_FIELD_METADATA,
SEARCH_FIELD_IS_REF,
SEARCH_FIELD_TIMESTAMP,
]
@experimental
def memoryrecord_to_milvus_dict(mem: MemoryRecord) -> dict[str, Any]:
"""Convert a memoryrecord into a dict.
Args:
mem (MemoryRecord): MemoryRecord to convert.
Returns:
dict: Dict result.
"""
ret_dict = {}
# Grab all the class vars
for key, val in vars(mem).items():
if val is not None:
# Remove underscore
if isinstance(val, datetime):
val = val.isoformat()
ret_dict[key[1:]] = val
return ret_dict
@experimental
def milvus_dict_to_memoryrecord(milvus_dict: dict[str, Any]) -> MemoryRecord:
"""Convert Milvus search result dict into MemoryRecord.
Args:
milvus_dict (dict): Search hit
Returns:
MemoryRecord
"""
# Embedding needs conversion to numpy array
embedding = milvus_dict.get(SEARCH_FIELD_EMBEDDING)
if embedding is not None:
embedding = array(embedding)
return MemoryRecord(
is_reference=milvus_dict.get(SEARCH_FIELD_IS_REF),
external_source_name=milvus_dict.get(SEARCH_FIELD_SRC),
id=milvus_dict.get(SEARCH_FIELD_ID),
description=milvus_dict.get(SEARCH_FIELD_DESC),
text=milvus_dict.get(SEARCH_FIELD_TEXT),
additional_metadata=milvus_dict.get(SEARCH_FIELD_METADATA),
embedding=embedding,
key=milvus_dict.get("key"),
timestamp=milvus_dict.get(SEARCH_FIELD_TIMESTAMP),
)
@experimental
def create_fields(dimensions: int) -> list[FieldSchema]:
"""Create the fields for the Milvus collection."""
return [
FieldSchema(
name=SEARCH_FIELD_ID,
dtype=DataType.VARCHAR,
is_primary=True,
auto_id=False,
max_length=100,
),
FieldSchema(
name=SEARCH_FIELD_TEXT,
dtype=DataType.VARCHAR,
max_length=65535,
),
FieldSchema(
name=SEARCH_FIELD_SRC,
dtype=DataType.VARCHAR,
max_length=400,
),
FieldSchema(
name=SEARCH_FIELD_DESC,
dtype=DataType.VARCHAR,
max_length=800,
),
FieldSchema(
name=SEARCH_FIELD_METADATA,
dtype=DataType.VARCHAR,
max_length=800,
),
FieldSchema(
name=SEARCH_FIELD_EMBEDDING,
dtype=DataType.FLOAT_VECTOR,
dim=dimensions,
),
FieldSchema(
name=SEARCH_FIELD_IS_REF,
dtype=DataType.BOOL,
),
FieldSchema(
name=SEARCH_FIELD_TIMESTAMP,
dtype=DataType.VARCHAR,
max_length=100,
),
]
@deprecated("This class will be removed in a future version.")
class MilvusMemoryStore(MemoryStoreBase):
"""Memory store based on Milvus."""
def __init__(
self,
uri: str = "http://localhost:19530",
token: str | None = None,
**kwargs: Any,
) -> None:
"""Memory store based on Milvus.
For more details on how to get the service started, take a look here:
- Milvus: https://milvus.io/docs/get_started.md
- Zilliz Cloud: https://docs.zilliz.com/docs/quick-start
Args:
uri (str, optional): The uri of the cluster. Defaults to
"http://localhost:19530".
token (Optional[str], optional): The token to connect to the cluster if
authentication is required. Defaults to None.
**kwargs (Any): Unused.
"""
connections.connect("default", uri=uri, token=token)
self.collections: dict[str, Collection] = {}
async def create_collection(
self,
collection_name: str,
dimension_num: int = 1536,
distance_type: str | None = "IP",
overwrite: bool = False,
consistency: str = "Session",
) -> None:
"""Create a Milvus collection.
Args:
collection_name (str): The name of the collection.
dimension_num (Optional[int], optional): The size of the embeddings being
stored. Defaults to 1536.
distance_type (Optional[str], optional): Which distance function, at the
moment only "IP" and "L2" are supported. Defaults to "IP".
overwrite (bool, optional): Whether to overwrite any existing collection
with the same name. Defaults to False.
consistency (str, optional): Which consistency level to use:
Strong, Session, Bounded, Eventually. Defaults to "Session".
"""
schema = CollectionSchema(
create_fields(dimension_num), "Semantic Kernel Milvus Collection", enable_dynamic_field=True
)
index_param = {"index_type": _INDEX_TYPE, "params": {"nlist": _NLIST}, "metric_type": distance_type}
if utility.has_collection(collection_name) and overwrite:
utility.drop_collection(collection_name=collection_name)
self.collections[collection_name] = Collection(
name=collection_name,
schema=schema,
consistency_level=consistency,
)
self.collections[collection_name].create_index(SEARCH_FIELD_EMBEDDING, index_param)
async def get_collections(
self,
) -> list[str]:
"""Return a list of present collections.
Returns:
List[str]: List of collection names.
"""
return utility.list_collections()
async def delete_collection(self, collection_name: str | None = None, all: bool = False) -> None:
"""Delete the specified collection.
If all is True, all collections in the cluster will be removed.
Args:
collection_name (str, optional): The name of the collection to delete. Defaults to "".
all (bool, optional): Whether to delete all collections. Defaults to False.
"""
if collection_name and utility.has_collection(collection_name):
utility.drop_collection(collection_name)
del self.collections[collection_name]
return
if all:
for collection in utility.list_collections():
utility.drop_collection(collection)
self.collections = {}
async def does_collection_exist(self, collection_name: str) -> bool:
"""Return if the collection exists in the cluster.
Args:
collection_name (str): The name of the collection.
Returns:
bool: True if it exists, False otherwise.
"""
return utility.has_collection(collection_name)
async def upsert(self, collection_name: str, record: MemoryRecord) -> str:
"""Upsert a single MemoryRecord into the collection.
Args:
collection_name (str): The name of the collection.
record (MemoryRecord): The record to store.
Returns:
str: The ID of the inserted record.
"""
# Use the batch insert with a total batch
res = await self.upsert_batch(
collection_name=collection_name,
records=[record],
batch_size=0,
)
return res[0]
async def upsert_batch(self, collection_name: str, records: list[MemoryRecord], batch_size=100) -> list[str]:
"""_summary_.
Args:
collection_name (str): The collection name.
records (List[MemoryRecord]): A list of memory records.
batch_size (int, optional): Batch size of the insert, 0 is a batch
size of total size. Defaults to 100.
Raises:
Exception: Collection doesnt exist.
e: Failed to upsert a record.
Returns:
List[str]: A list of inserted ID's.
"""
# Check if the collection exists.
if collection_name not in utility.list_collections():
logger.debug(f"Collection {collection_name} does not exist, cannot insert.")
raise ServiceResourceNotFoundError(f"Collection {collection_name} does not exist, cannot insert.")
# Convert the records to dicts
insert_list = [memoryrecord_to_milvus_dict(record) for record in records]
try:
ids = self.collections[collection_name].upsert(data=insert_list).primary_keys
self.collections[collection_name].flush()
return ids
except Exception as e:
logger.debug(f"Upsert failed due to: {e}")
raise ServiceResponseException(f"Upsert failed due to: {e}") from e
async def get(self, collection_name: str, key: str, with_embedding: bool) -> MemoryRecord:
"""Get the MemoryRecord corresponding to the key.
Args:
collection_name (str): The collection to get from.
key (str): The ID to grab.
with_embedding (bool): Whether to include the embedding in the results.
Returns:
MemoryRecord: The MemoryRecord for the key.
"""
res = await self.get_batch(collection_name=collection_name, keys=[key], with_embeddings=with_embedding)
return res[0]
async def get_batch(self, collection_name: str, keys: list[str], with_embeddings: bool) -> list[MemoryRecord]:
"""Get the MemoryRecords corresponding to the keys.
Args:
collection_name (str): _description_
keys (List[str]): _description_
with_embeddings (bool): _description_
Raises:
Exception: _description_
e: _description_
Returns:
List[MemoryRecord]: _description_
"""
# Check if the collection exists
if not utility.has_collection(collection_name):
logger.debug(f"Collection {collection_name} does not exist, cannot get.")
raise ServiceResourceNotFoundError(f"Collection {collection_name} does not exist, cannot get.")
try:
self.collections[collection_name].load()
gets = self.collections[collection_name].query(
expr=f"{SEARCH_FIELD_ID} in {keys}",
output_fields=OUTPUT_FIELDS_W_EMBEDDING if with_embeddings else OUTPUT_FIELDS_WO_EMBEDDING,
)
except Exception as e:
logger.debug(f"Get failed due to: {e}")
raise ServiceResponseException(f"Get failed due to: {e}") from e
return [milvus_dict_to_memoryrecord(get) for get in gets]
async def remove(self, collection_name: str, key: str) -> None:
"""Remove the specified record based on key.
Args:
collection_name (str): Collection to remove from.
key (str): The key to remove.
"""
await self.remove_batch(collection_name=collection_name, keys=[key])
async def remove_batch(self, collection_name: str, keys: list[str]) -> None:
"""Remove multiple records based on keys.
Args:
collection_name (str): Collection to remove from
keys (List[str]): The list of keys.
Raises:
Exception: Collection doesnt exist.
e: Failure to remove key.
"""
if collection_name not in utility.list_collections():
logger.debug(f"Collection {collection_name} does not exist, cannot remove.")
raise ServiceResourceNotFoundError(f"Collection {collection_name} does not exist, cannot remove.")
try:
self.collections[collection_name].load()
result = self.collections[collection_name].delete(
expr=f"{SEARCH_FIELD_ID} in {keys}",
)
self.collections[collection_name].flush()
except Exception as e:
logger.debug(f"Remove failed due to: {e}")
raise ServiceResponseException(f"Remove failed due to: {e}") from e
if result.delete_count != len(keys):
logger.debug(f"Failed to remove all keys, {result.delete_count} removed out of {len(keys)}")
raise ServiceResponseException(
f"Failed to remove all keys, {result.delete_count} removed out of {len(keys)}"
)
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]]:
"""Find the nearest `limit` matches for an embedding.
Args:
collection_name (str): The collection to search.
embedding (ndarray): The embedding to search.
limit (int): The total results to display.
min_relevance_score (float, optional): Minimum distance to include. Defaults to None.
with_embeddings (bool, optional): Whether to include embeddings in result. Defaults to False.
Raises:
Exception: Missing collection
e: Failure to search
Returns:
List[Tuple[MemoryRecord, float]]: MemoryRecord and distance tuple.
"""
# Check if collection exists
if collection_name not in utility.list_collections():
logger.debug(f"Collection {collection_name} does not exist, cannot search.")
raise ServiceResourceNotFoundError(f"Collection {collection_name} does not exist, cannot search.")
# Search requests takes a list of requests.
if len(embedding.shape) == 1:
embedding = expand_dims(embedding, axis=0)
try:
self.collections[collection_name].load()
metric = self.collections[collection_name].index(index_name=SEARCH_FIELD_EMBEDDING).params["metric_type"]
# Try with passed in metric
results = self.collections[collection_name].search(
data=embedding,
anns_field=SEARCH_FIELD_EMBEDDING,
limit=limit,
output_fields=OUTPUT_FIELDS_W_EMBEDDING if with_embeddings else OUTPUT_FIELDS_WO_EMBEDDING,
param={"metric_type": metric},
)[0]
except Exception as e:
logger.debug(f"Search failed: {e}")
raise ServiceResponseException(f"Search failed: {e}") from e
return [
(milvus_dict_to_memoryrecord(result.fields), result.distance)
for result in results
if result.distance >= min_relevance_score
]
async def get_nearest_match(
self,
collection_name: str,
embedding: ndarray,
min_relevance_score: float = 0.0,
with_embedding: bool = False,
) -> tuple[MemoryRecord, float] | None:
"""Find the nearest match for an embedding.
Args:
collection_name (str): The collection to search.
embedding (ndarray): The embedding to search for.
min_relevance_score (float, optional): T. Defaults to 0.0.
with_embedding (bool, optional): Whether to include embedding in result. Defaults to False.
Returns:
Tuple[MemoryRecord, float]: A tuple of record and distance.
"""
m = await self.get_nearest_matches(
collection_name,
embedding,
1,
min_relevance_score,
with_embedding,
)
if len(m) > 0:
return m[0]
return None
@@ -0,0 +1,52 @@
# microsoft.semantic_kernel.connectors.memory.mongodb_atlas
This connector uses [MongoDB Atlas Vector Search](https://www.mongodb.com/products/platform/atlas-vector-search) to implement Semantic Memory.
## Quick Start
1. Create [Atlas cluster](https://www.mongodb.com/docs/atlas/getting-started/)
2. Create a collection
3. Create [Vector Search Index](https://www.mongodb.com/docs/atlas/atlas-search/field-types/knn-vector/) for the collection.
The index has to be defined on a field called ```embedding```. For example:
```
{
"mappings": {
"dynamic": true,
"fields": {
"embedding": {
"dimension": 1024,
"similarity": "cosine_similarity",
"type": "knnVector"
}
}
}
}
```
4. Create the MongoDB memory store
```python
import semantic_kernel as sk
import semantic_kernel.connectors.ai.open_ai
from semantic_kernel.connectors.memory.mongodb_atlas import (
MongoDBAtlasMemoryStore
)
kernel = sk.Kernel()
...
kernel.register_memory_store(memory_store=MongoDBAtlasMemoryStore(
# connection_string = if not provided pull from .env
))
...
```
## Important Notes
### Vector search indexes
In this version, vector search index management is outside of ```MongoDBAtlasMemoryStore``` scope.
Creation and maintenance of the indexes have to be done by the user. Please note that deleting a collection
(```memory_store.delete_collection_async```) will delete the index as well.
@@ -0,0 +1,333 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
import sys
from collections.abc import Mapping
from importlib import metadata
from typing import Any
from motor import core, motor_asyncio
from numpy import ndarray
from pydantic import ValidationError
from pymongo import DeleteOne, ReadPreference, UpdateOne, results
from pymongo.driver_info import DriverInfo
from semantic_kernel.connectors.memory_stores.mongodb_atlas.utils import (
MONGODB_FIELD_EMBEDDING,
MONGODB_FIELD_ID,
document_to_memory_record,
memory_record_to_mongo_document,
)
from semantic_kernel.connectors.mongodb import NUM_CANDIDATES_SCALAR
from semantic_kernel.exceptions import ServiceResourceNotFoundError
from semantic_kernel.exceptions.memory_connector_exceptions import MemoryConnectorInitializationError
from semantic_kernel.memory.memory_record import MemoryRecord
from semantic_kernel.memory.memory_store_base import MemoryStoreBase
if sys.version_info >= (3, 12):
from warnings import deprecated
else:
from typing_extensions import deprecated
logger: logging.Logger = logging.getLogger(__name__)
@deprecated("MongoDBAtlasMemoryStore is deprecated. Use MongoDBStore and Collection instead.")
class MongoDBAtlasMemoryStore(MemoryStoreBase):
"""Memory Store for MongoDB Atlas Vector Search Connections."""
def __init__(
self,
index_name: str | None = None,
connection_string: str | None = None,
database_name: str | None = None,
read_preference: ReadPreference | None = ReadPreference.PRIMARY,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
):
"""Create the MongoDB Atlas Memory Store.
Args:
index_name (str): The name of the index.
connection_string (str): The connection string for the MongoDB Atlas instance.
database_name (str): The name of the database.
read_preference (ReadPreference): The read preference for the connection.
env_file_path (str): The path to the .env file containing the connection string.
env_file_encoding (str): The encoding of the .env file.
"""
from semantic_kernel.connectors.mongodb import MongoDBAtlasSettings
try:
mongodb_settings = MongoDBAtlasSettings(
database_name=database_name,
index_name=index_name,
connection_string=connection_string,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
except ValidationError as ex:
raise MemoryConnectorInitializationError("Failed to create MongoDB Atlas settings.") from ex
self.mongo_client: motor_asyncio.AsyncIOMotorClient = motor_asyncio.AsyncIOMotorClient(
mongodb_settings.connection_string.get_secret_value(),
read_preference=read_preference,
driver=DriverInfo("Microsoft Semantic Kernel", metadata.version("semantic-kernel")),
)
self.database_name: str = mongodb_settings.database_name
self.index_name: str = mongodb_settings.index_name
@property
def database(self) -> core.AgnosticDatabase:
"""The database object."""
return self.mongo_client[self.database_name]
@property
def num_candidates(self) -> int:
"""The number of candidates to return."""
return self.__num_candidates
async def close(self):
"""Async close connection, invoked by MemoryStoreBase.__aexit__()."""
if self.mongo_client:
self.mongo_client.close()
del self.mongo_client
async def create_collection(self, collection_name: str) -> None:
"""Creates a new collection in the data store.
Args:
collection_name (str): The name associated with a collection of embeddings.
Returns:
None
"""
if not await self.does_collection_exist(collection_name):
await self.database.create_collection(collection_name)
async def get_collections(
self,
) -> list[str]:
"""Gets all collection names in the data store.
Returns:
List[str]: A group of collection names.
"""
return await self.database.list_collection_names()
async def delete_collection(self, collection_name: str) -> None:
"""Deletes a collection from the data store.
Args:
collection_name (str): The name associated with a collection of embeddings.
Returns:
None
"""
await self.database[collection_name].drop()
async def does_collection_exist(self, collection_name: str) -> bool:
"""Determines if a collection exists in the data store.
Args:
collection_name (str): The name associated with a collection of embeddings.
Returns:
bool: True if given collection exists, False if not.
"""
return collection_name in (await self.get_collections())
async def upsert(self, collection_name: str, record: MemoryRecord) -> str:
"""Upserts 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.
"""
document: Mapping[str, Any] = memory_record_to_mongo_document(record)
update_result: results.UpdateResult = await self.database[collection_name].update_one(
document, {"$set": document}, upsert=True
)
if not update_result.acknowledged:
raise ValueError("Upsert failed")
return record._id
async def upsert_batch(self, collection_name: str, records: list[MemoryRecord]) -> list[str]:
"""Upserts a group 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 (MemoryRecord): The memory records to upsert.
Returns:
List[str]: The unique identifiers for the memory records.
"""
upserts: list[UpdateOne] = []
for record in records:
document = memory_record_to_mongo_document(record)
upserts.append(UpdateOne(document, {"$set": document}, upsert=True))
bulk_update_result: results.BulkWriteResult = await self.database[collection_name].bulk_write(
upserts, ordered=False
)
# Assert the number matched and the number upserted equal the total batch updated
logger.debug(
"matched_count=%s, upserted_count=%s",
bulk_update_result.matched_count,
bulk_update_result.upserted_count,
)
if bulk_update_result.matched_count + bulk_update_result.upserted_count != len(records):
raise ValueError("Batch upsert failed")
return [record._id for record in records]
async def get(self, collection_name: str, key: str, with_embedding: bool) -> MemoryRecord:
"""Gets a memory record from the data store. Does not guarantee that the collection exists.
Args:
collection_name (str): The name associated with a collection of embeddings.
key (str): The unique id associated with the memory record to get.
with_embedding (bool): If true, the embedding will be returned in the memory record.
Returns:
MemoryRecord: The memory record if found
"""
document = await self.database[collection_name].find_one({MONGODB_FIELD_ID: key})
return document_to_memory_record(document, with_embedding) if document else None
async def get_batch(self, collection_name: str, keys: list[str], with_embeddings: bool) -> list[MemoryRecord]:
"""Gets a batch of memory records from the data store. Does not guarantee that the collection exists.
Args:
collection_name (str): The name associated with a collection of embeddings.
keys (List[str]): The unique ids associated with the memory records to get.
with_embeddings (bool): If true, the embedding will be returned in the memory records.
Returns:
List[MemoryRecord]: The memory records associated with the unique keys provided.
"""
results = self.database[collection_name].find({MONGODB_FIELD_ID: {"$in": keys}})
return [
document_to_memory_record(result, with_embeddings) for result in await results.to_list(length=len(keys))
]
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 associated with a collection of embeddings.
key (str): The unique id associated with the memory record to remove.
Returns:
None
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f"collection {collection_name} not found")
await self.database[collection_name].delete_one({MONGODB_FIELD_ID: key})
async def remove_batch(self, collection_name: str, keys: list[str]) -> None:
"""Removes a batch of memory records from the data store. Does not guarantee that the collection exists.
Args:
collection_name (str): The name associated with a collection of embeddings.
keys (List[str]): The unique ids associated with the memory records to remove.
Returns:
None
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f"collection {collection_name} not found")
deletes: list[DeleteOne] = [DeleteOne({MONGODB_FIELD_ID: key}) for key in keys]
bulk_write_result = await self.database[collection_name].bulk_write(deletes, ordered=False)
logger.debug("%s entries deleted", bulk_write_result.deleted_count)
async def get_nearest_matches(
self,
collection_name: str,
embedding: ndarray,
limit: int,
with_embeddings: bool,
min_relevance_score: float | None = None,
) -> list[tuple[MemoryRecord, float]]:
"""Gets the nearest matches to an embedding of type float. Does not guarantee that the collection exists.
Args:
collection_name (str): The name associated with a collection of embeddings.
embedding (ndarray): The embedding to compare the collection's embeddings with.
limit (int): The maximum number of similarity results to return, defaults to 1.
min_relevance_score (float): The minimum relevance threshold for returned results.
with_embeddings (bool): If true, the embeddings will be returned in the memory records.
Returns:
List[Tuple[MemoryRecord, float]]: A list of tuples where item1 is a MemoryRecord and item2
is its similarity score as a float.
"""
pipeline: list[dict[str, Any]] = []
vector_search_query: list[Mapping[str, Any]] = {
"$vectorSearch": {
"queryVector": embedding.tolist(),
"limit": limit,
"numCandidates": limit * NUM_CANDIDATES_SCALAR,
"path": MONGODB_FIELD_EMBEDDING,
"index": self.index_name,
}
}
pipeline.append(vector_search_query)
# add meta search scoring
pipeline.append({"$set": {"score": {"$meta": "vectorSearchScore"}}})
if min_relevance_score is not None:
pipeline.append({"$match": {"score": {"$gte": min_relevance_score}}})
cursor = self.database[collection_name].aggregate(pipeline)
return [
(
document_to_memory_record(doc, with_embeddings=with_embeddings),
doc["score"],
)
for doc in await cursor.to_list(length=limit)
]
async def get_nearest_match(
self,
collection_name: str,
embedding: ndarray,
with_embedding: bool,
min_relevance_score: float | None = None,
) -> tuple[MemoryRecord, float]:
"""Gets the nearest match to an embedding of type float. Does not guarantee that the collection exists.
Args:
collection_name (str): The name associated with a collection of embeddings.
embedding (ndarray): The embedding to compare the collection's embeddings with.
min_relevance_score (float): The minimum relevance threshold for returned result.
with_embedding (bool): If true, the embeddings will be returned in the memory record.
Returns:
Tuple[MemoryRecord, float]: A tuple consisting of the MemoryRecord and the similarity score as a float.
"""
matches: list[tuple[MemoryRecord, float]] = 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] if matches else None
@@ -0,0 +1,64 @@
# Copyright (c) Microsoft. All rights reserved.
from numpy import array
from semantic_kernel.memory.memory_record import MemoryRecord
MONGODB_FIELD_ID = "_id"
MONGODB_FIELD_TEXT = "text"
MONGODB_FIELD_EMBEDDING = "embedding"
MONGODB_FIELD_SRC = "externalSourceName"
MONGODB_FIELD_DESC = "description"
MONGODB_FIELD_METADATA = "additionalMetadata"
MONGODB_FIELD_IS_REF = "isReference"
MONGODB_FIELD_KEY = "key"
MONGODB_FIELD_TIMESTAMP = "timestamp"
def document_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 embeddings.
Returns:
MemoryRecord: The MemoryRecord from Azure Cognitive Search Data Result.
"""
meta = data.get(MONGODB_FIELD_METADATA, {})
return MemoryRecord(
id=meta.get(MONGODB_FIELD_ID),
text=meta.get(MONGODB_FIELD_TEXT),
external_source_name=meta.get(MONGODB_FIELD_SRC),
description=meta.get(MONGODB_FIELD_DESC),
additional_metadata=meta.get(MONGODB_FIELD_METADATA),
is_reference=meta.get(MONGODB_FIELD_IS_REF),
embedding=array(data.get(MONGODB_FIELD_EMBEDDING)) if with_embeddings else None,
timestamp=data.get(MONGODB_FIELD_TIMESTAMP),
key=meta.get(MONGODB_FIELD_ID),
)
def memory_record_to_mongo_document(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 {
MONGODB_FIELD_ID: record._id,
MONGODB_FIELD_METADATA: {
MONGODB_FIELD_ID: record._id,
MONGODB_FIELD_TEXT: record._text,
MONGODB_FIELD_SRC: record._external_source_name or "",
MONGODB_FIELD_DESC: record._description or "",
MONGODB_FIELD_METADATA: record._additional_metadata or "",
MONGODB_FIELD_IS_REF: record._is_reference,
},
MONGODB_FIELD_EMBEDDING: record._embedding.tolist(),
MONGODB_FIELD_TIMESTAMP: record._timestamp,
}
@@ -0,0 +1,402 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
import sys
from typing import NamedTuple
from numpy import ndarray
from pinecone import FetchResponse, IndexList, IndexModel, Pinecone, ServerlessSpec
from pydantic import ValidationError
from semantic_kernel.connectors.memory_stores.pinecone.utils import build_payload, parse_payload
from semantic_kernel.connectors.pinecone import PineconeSettings
from semantic_kernel.exceptions import (
ServiceInitializationError,
ServiceInvalidRequestError,
ServiceResourceNotFoundError,
ServiceResponseException,
)
from semantic_kernel.exceptions.memory_connector_exceptions import MemoryConnectorInitializationError
from semantic_kernel.memory.memory_record import MemoryRecord
from semantic_kernel.memory.memory_store_base import MemoryStoreBase
if sys.version_info >= (3, 13):
from warnings import deprecated
else:
from typing_extensions import deprecated
# Limitations set by Pinecone at https://docs.pinecone.io/reference/known-limitations
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__)
@deprecated("This class will be removed in a future version. Use PineconeStore and Collection instead.")
class PineconeMemoryStore(MemoryStoreBase):
"""A memory store that uses Pinecone as the backend."""
_default_dimensionality: int
DEFAULT_INDEX_SPEC: ServerlessSpec = ServerlessSpec(
cloud="aws",
region="us-east-1",
)
def __init__(
self,
api_key: str,
default_dimensionality: int,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initializes a new instance of the PineconeMemoryStore class.
Args:
api_key (str): The Pinecone API key.
default_dimensionality (int): The default dimensionality to use for new collections.
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)
"""
if default_dimensionality > MAX_DIMENSIONALITY:
raise MemoryConnectorInitializationError(
f"Dimensionality of {default_dimensionality} exceeds the maximum allowed value of {MAX_DIMENSIONALITY}."
)
try:
pinecone_settings = PineconeSettings(
api_key=api_key,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
except ValidationError as ex:
raise MemoryConnectorInitializationError("Failed to create Pinecone settings.", ex) from ex
self._default_dimensionality = default_dimensionality
self.pinecone = Pinecone(api_key=pinecone_settings.api_key.get_secret_value())
self.collection_names_cache = set()
async def create_collection(
self,
collection_name: str,
dimension_num: int | None = None,
distance_type: str | None = "cosine",
index_spec: NamedTuple = DEFAULT_INDEX_SPEC,
) -> None:
"""Creates a new collection in Pinecone if it does not exist.
This function creates an index, by default the following index
settings are used: metric = cosine, cloud = aws, region = us-east-1.
Args:
collection_name (str): The name of the collection to create.
In Pinecone, a collection is represented as an index. The concept
of "collection" in Pinecone is just a static copy of an index.
dimension_num (int, optional): The dimensionality of the embeddings.
distance_type (str, optional): The distance metric to use for the index.
(default: {"cosine"})
index_spec (NamedTuple, optional): The index spec to use for the index.
"""
if dimension_num is None:
dimension_num = self._default_dimensionality
if dimension_num > MAX_DIMENSIONALITY:
raise ServiceInitializationError(
f"Dimensionality of {dimension_num} exceeds " + f"the maximum allowed value of {MAX_DIMENSIONALITY}."
)
if not await self.does_collection_exist(collection_name):
self.pinecone.create_index(
name=collection_name, dimension=dimension_num, metric=distance_type, spec=index_spec
)
self.collection_names_cache.add(collection_name)
async def describe_collection(self, collection_name: str) -> IndexModel | None:
"""Gets the description of the index.
Args:
collection_name (str): The name of the index to get.
Returns:
Optional[dict]: The index.
"""
if await self.does_collection_exist(collection_name):
return self.pinecone.describe_index(collection_name)
return None
async def get_collections(
self,
) -> IndexList:
"""Gets the list of collections.
Returns:
IndexList: The list of collections.
"""
return self.pinecone.list_indexes()
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
"""
if await self.does_collection_exist(collection_name):
self.pinecone.delete_index(collection_name)
self.collection_names_cache.discard(collection_name)
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.
"""
if collection_name in self.collection_names_cache:
return True
index_collection_names = self.pinecone.list_indexes().names()
self.collection_names_cache |= set(index_collection_names)
return collection_name in index_collection_names
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 database key of the record. In Pinecone, this is the record ID.
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
collection = self.pinecone.Index(collection_name)
upsert_response = collection.upsert(
vectors=[(record._id, record.embedding.tolist(), build_payload(record))],
namespace="",
)
if upsert_response.upserted_count is None:
raise ServiceResponseException(f"Error upserting record: {upsert_response.message}")
return record._id
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.
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
collection = self.pinecone.Index(collection_name)
vectors = [
(
record._id,
record.embedding.tolist(),
build_payload(record),
)
for record in records
]
upsert_response = collection.upsert(vectors, namespace="", batch_size=MAX_UPSERT_BATCH_SIZE)
if upsert_response.upserted_count is None:
raise ServiceResponseException(f"Error upserting record: {upsert_response.message}")
return [record._id for record in records]
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.
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
collection = self.pinecone.Index(collection_name)
fetch_response = collection.fetch([key])
if len(fetch_response.vectors) == 0:
raise ServiceResourceNotFoundError(f"Record with key '{key}' does not exist")
return parse_payload(fetch_response.vectors[key], 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.
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
fetch_response = await self.__get_batch(collection_name, keys, with_embeddings)
return [parse_payload(fetch_response.vectors[key], with_embeddings) for key in fetch_response.vectors]
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
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
collection = self.pinecone.Index(collection_name)
collection.delete([key])
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
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
collection = self.pinecone.Index(collection_name)
for i in range(0, len(keys), MAX_DELETE_BATCH_SIZE):
collection.delete(keys[i : i + MAX_DELETE_BATCH_SIZE])
collection.delete(keys)
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 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.
"""
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.
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
collection = self.pinecone.Index(collection_name)
if limit > MAX_QUERY_WITHOUT_METADATA_BATCH_SIZE:
raise ServiceInvalidRequestError(
"Limit must be less than or equal to " + f"{MAX_QUERY_WITHOUT_METADATA_BATCH_SIZE}"
)
if limit > MAX_QUERY_WITH_METADATA_BATCH_SIZE:
query_response = collection.query(
vector=embedding.tolist(),
top_k=limit,
include_values=False,
include_metadata=False,
)
keys = [match.id for match in query_response.matches]
fetch_response = await self.__get_batch(collection_name, keys, with_embeddings)
vectors = fetch_response.vectors
for match in query_response.matches:
vectors[match.id].update(match)
matches = [vectors[key] for key in vectors]
else:
query_response = collection.query(
vector=embedding.tolist(),
top_k=limit,
include_values=with_embeddings,
include_metadata=True,
)
matches = query_response.matches
if min_relevance_score is not None:
matches = [match for match in matches if match.score >= min_relevance_score]
return (
[
(
parse_payload(match, with_embeddings),
match["score"],
)
for match in matches
]
if len(matches) > 0
else []
)
async def __get_batch(
self, collection_name: str, keys: list[str], with_embeddings: bool = False
) -> "FetchResponse":
index = self.pinecone.Index(collection_name)
if len(keys) > MAX_FETCH_BATCH_SIZE:
fetch_response = index.fetch(keys[0:MAX_FETCH_BATCH_SIZE])
for i in range(MAX_FETCH_BATCH_SIZE, len(keys), MAX_FETCH_BATCH_SIZE):
fetch_response.vectors.update(index.fetch(keys[i : i + MAX_FETCH_BATCH_SIZE]).vectors)
else:
fetch_response = index.fetch(keys)
return fetch_response
@@ -0,0 +1,33 @@
# Copyright (c) Microsoft. All rights reserved.
import numpy
from pinecone import Vector
from semantic_kernel.memory.memory_record import MemoryRecord
def build_payload(record: MemoryRecord) -> dict:
"""Builds a metadata payload to be sent to Pinecone from a MemoryRecord."""
payload: dict = {}
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(record: Vector, with_embeddings: bool) -> MemoryRecord:
"""Parses a record from Pinecone into a MemoryRecord."""
payload = record.metadata
description = payload.get("description", None)
text = payload.get("text", None)
additional_metadata = payload.get("additional_metadata", None)
return MemoryRecord.local_record(
id=record.id,
description=description,
text=text,
additional_metadata=additional_metadata,
embedding=record.values if with_embeddings else numpy.array([]),
)
@@ -0,0 +1,516 @@
# Copyright (c) Microsoft. All rights reserved.
import atexit
import json
import logging
import sys
import numpy as np
from numpy import ndarray
from psycopg import Cursor
from psycopg.sql import SQL, Identifier
from psycopg_pool import ConnectionPool
from pydantic import ValidationError
from semantic_kernel.connectors.postgres import DEFAULT_SCHEMA, MAX_DIMENSIONALITY, PostgresSettings
from semantic_kernel.exceptions import (
ServiceInitializationError,
ServiceResourceNotFoundError,
ServiceResponseException,
)
from semantic_kernel.exceptions.memory_connector_exceptions import MemoryConnectorInitializationError
from semantic_kernel.memory.memory_record import MemoryRecord
from semantic_kernel.memory.memory_store_base import MemoryStoreBase
if sys.version_info >= (3, 13):
from warnings import deprecated
else:
from typing_extensions import deprecated
logger: logging.Logger = logging.getLogger(__name__)
@deprecated("This class will be removed in a future version. Use PostgresStore and Collection instead.")
class PostgresMemoryStore(MemoryStoreBase):
"""A memory store that uses Postgres with pgvector as the backend."""
_connection_string: str
_connection_pool: ConnectionPool
_default_dimensionality: int
_schema: str
def __init__(
self,
connection_string: str,
default_dimensionality: int,
min_pool: int | None = None,
max_pool: int | None = None,
schema: str = DEFAULT_SCHEMA,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initializes a new instance of the PostgresMemoryStore class.
Args:
connection_string: The connection string to the Postgres database.
default_dimensionality: The default dimensionality of the embeddings.
min_pool: The minimum number of connections in the connection pool.
max_pool: The maximum number of connections in the connection pool.
schema: The schema to use. (default: {"public"})
env_file_path: Use the environment settings file as a fallback
to environment variables. (Optional)
env_file_encoding: The encoding of the environment settings file.
"""
try:
postgres_settings = PostgresSettings(
connection_string=connection_string,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
except ValidationError as ex:
raise MemoryConnectorInitializationError("Failed to create Postgres settings.", ex) from ex
min_pool = min_pool or postgres_settings.min_pool
max_pool = max_pool or postgres_settings.max_pool
self._check_dimensionality(default_dimensionality)
self._default_dimensionality = default_dimensionality
self._connection_pool = ConnectionPool(
min_size=min_pool, max_size=max_pool, open=True, kwargs=postgres_settings.get_connection_args()
)
self._schema = schema
atexit.register(self._connection_pool.close)
async def create_collection(
self,
collection_name: str,
dimension_num: int | None = None,
) -> None:
r"""Creates a new collection.
Args:
collection_name: The name of the collection to create.\n
dimension_num: The dimensionality of the embeddings. (default: {None})
Uses the default dimensionality when not provided
Returns:
None
"""
if dimension_num is None:
dimension_num = self._default_dimensionality
else:
self._check_dimensionality(dimension_num)
with self._connection_pool.connection() as conn, conn.cursor() as cur:
cur.execute(
SQL(
"""
CREATE TABLE IF NOT EXISTS {scm}.{tbl} (
key TEXT PRIMARY KEY,
embedding vector({dim}),
metadata JSONB,
timestamp TIMESTAMP
)"""
).format(
scm=Identifier(self._schema),
tbl=Identifier(collection_name),
dim=dimension_num,
),
(),
)
async def get_collections(self) -> list[str]:
"""Gets the list of collections.
Returns:
The list of collections.
"""
with self._connection_pool.connection() as conn, conn.cursor() as cur:
return await self.__get_collections(cur)
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
"""
with self._connection_pool.connection() as conn, conn.cursor() as cur:
cur.execute(
SQL("DROP TABLE IF EXISTS {scm}.{tbl} CASCADE").format(
scm=Identifier(self._schema), tbl=Identifier(collection_name)
),
)
async def does_collection_exist(self, collection_name: str) -> bool:
"""Checks if a collection exists.
Args:
collection_name: The name of the collection to check.
Returns:
True if the collection exists; otherwise, False.
"""
with self._connection_pool.connection() as conn, conn.cursor() as cur:
return await self.__does_collection_exist(cur, collection_name)
async def upsert(self, collection_name: str, record: MemoryRecord) -> str:
"""Upserts a record.
Args:
collection_name: The name of the collection to upsert the record into.
record: The record to upsert.
Returns:
The unique database key of the record. In Pinecone, this is the record ID.
"""
with self._connection_pool.connection() as conn, conn.cursor() as cur:
if not await self.__does_collection_exist(cur, collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
cur.execute(
SQL(
"""
INSERT INTO {scm}.{tbl} (key, embedding, metadata, timestamp)
VALUES (%s, %s, %s, %s)
ON CONFLICT (key) DO UPDATE
SET embedding = EXCLUDED.embedding,
metadata = EXCLUDED.metadata,
timestamp = EXCLUDED.timestamp
RETURNING key
"""
).format(
scm=Identifier(self._schema),
tbl=Identifier(collection_name),
),
(
record._id,
record.embedding.tolist(),
self.__serialize_metadata(record),
record._timestamp,
),
)
result = cur.fetchone()
if result is None:
raise ServiceResponseException("Upsert failed")
return result[0]
async def upsert_batch(self, collection_name: str, records: list[MemoryRecord]) -> list[str]:
"""Upserts a batch of records.
Args:
collection_name: The name of the collection to upsert the records into.
records: The records to upsert.
Returns:
List[str]: The unique database keys of the records.
"""
with self._connection_pool.connection() as conn, conn.cursor() as cur:
if not await self.__does_collection_exist(cur, collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
cur.nextset()
cur.executemany(
SQL(
"""
INSERT INTO {scm}.{tbl} (key, embedding, metadata, timestamp)
VALUES (%s, %s, %s, %s)
ON CONFLICT (key) DO UPDATE
SET embedding = EXCLUDED.embedding,
metadata = EXCLUDED.metadata,
timestamp = EXCLUDED.timestamp
RETURNING key
"""
).format(
scm=Identifier(self._schema),
tbl=Identifier(collection_name),
),
[
(
record._id,
record.embedding.tolist(),
self.__serialize_metadata(record),
record._timestamp,
)
for record in records
],
returning=True,
)
# collate results
results = [cur.fetchone()]
while cur.nextset():
results.append(cur.fetchone())
if None in results:
raise ServiceResponseException("Upsert failed")
return [result[0] for result in results if result is not None]
async def get(self, collection_name: str, key: str, with_embedding: bool = False) -> MemoryRecord:
"""Gets a record.
Args:
collection_name: The name of the collection to get the record from.
key: The unique database key of the record.
with_embedding: Whether to include the embedding in the result. (default: {False})
Returns:
The record.
"""
with self._connection_pool.connection() as conn, conn.cursor() as cur:
if not await self.__does_collection_exist(cur, collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
cur.execute(
SQL(
"""
SELECT key, embedding, metadata, timestamp
FROM {scm}.{tbl}
WHERE key = %s
"""
).format(
scm=Identifier(self._schema),
tbl=Identifier(collection_name),
),
(key,),
)
result = cur.fetchone()
if result is None:
raise ServiceResourceNotFoundError("Key not found")
return MemoryRecord.local_record(
id=result[0],
embedding=(
np.fromstring(result[1].strip("[]"), dtype=float, sep=",") if with_embedding else np.array([])
),
text=result[2]["text"],
description=result[2]["description"],
additional_metadata=result[2]["additional_metadata"],
timestamp=result[3],
)
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: The name of the collection to get the records from.
keys: The unique database keys of the records.
with_embeddings: Whether to include the embeddings in the results. (default: {False})
Returns:
The records that were found from list of keys, can be empty.
"""
with self._connection_pool.connection() as conn, conn.cursor() as cur:
if not await self.__does_collection_exist(cur, collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
cur.execute(
SQL(
"""
SELECT key, embedding, metadata, timestamp
FROM {scm}.{tbl}
WHERE key = ANY(%s)
"""
).format(
scm=Identifier(self._schema),
tbl=Identifier(collection_name),
),
(list(keys),),
)
results = cur.fetchall()
return [
MemoryRecord.local_record(
id=result[0],
embedding=(
np.fromstring(result[1].strip("[]"), dtype=float, sep=",") if with_embeddings else np.array([])
),
text=result[2]["text"],
description=result[2]["description"],
additional_metadata=result[2]["additional_metadata"],
timestamp=result[3],
)
for result in results
]
async def remove(self, collection_name: str, key: str) -> None:
"""Removes a record.
Args:
collection_name: The name of the collection to remove the record from.
key: The unique database key of the record to remove.
Returns:
None
"""
with self._connection_pool.connection() as conn, conn.cursor() as cur:
if not await self.__does_collection_exist(cur, collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
cur.execute(
SQL(
"""
DELETE FROM {scm}.{tbl}
WHERE key = %s
"""
).format(scm=Identifier(self._schema), tbl=Identifier(collection_name)),
(key,),
)
async def remove_batch(self, collection_name: str, keys: list[str]) -> None:
"""Removes a batch of records.
Args:
collection_name: The name of the collection to remove the records from.
keys: The unique database keys of the records to remove.
Returns:
None
"""
with self._connection_pool.connection() as conn, conn.cursor() as cur:
if not await self.__does_collection_exist(cur, collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
cur.execute(
SQL(
"""
DELETE FROM {scm}.{tbl}
WHERE key = ANY(%s)
"""
).format(scm=Identifier(self._schema), tbl=Identifier(collection_name)),
(list(keys),),
)
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: The name of the collection to get the nearest matches from.
embedding: The embedding to find the nearest matches to.
limit: The maximum number of matches to return.
min_relevance_score: The minimum relevance score of the matches. (default: {0.0})
with_embeddings: Whether to include the embeddings in the results. (default: {False})
Returns:
The records and their relevance scores.
"""
with self._connection_pool.connection() as conn, conn.cursor() as cur:
if not await self.__does_collection_exist(cur, collection_name):
raise ServiceResourceNotFoundError(f"Collection '{collection_name}' does not exist")
cur.execute(
SQL(
"""
SELECT key,
embedding,
metadata,
cosine_similarity,
timestamp
FROM (
SELECT key, embedding, metadata, 1 - (embedding <=> '[{emb}]') \
AS cosine_similarity, timestamp
FROM {scm}.{tbl}
) AS subquery
WHERE cosine_similarity >= {mrs}
ORDER BY cosine_similarity DESC
LIMIT {limit}
"""
).format(
scm=Identifier(self._schema),
tbl=Identifier(collection_name),
mrs=min_relevance_score,
limit=limit,
emb=SQL(",").join(embedding.tolist()),
)
)
results = cur.fetchall()
return [
(
MemoryRecord.local_record(
id=result[0],
embedding=(
np.fromstring(result[1].strip("[]"), dtype=float, sep=",")
if with_embeddings
else np.array([])
),
text=result[2]["text"],
description=result[2]["description"],
additional_metadata=result[2]["additional_metadata"],
timestamp=result[4],
),
result[3],
)
for result in results
]
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: The name of the collection to get the nearest match from.
embedding: The embedding to find the nearest match to.
min_relevance_score: The minimum relevance score of the match. (default: {0.0})
with_embedding: Whether to include the embedding in the result. (default: {False})
Returns:
Tuple[MemoryRecord, float]: The record and the relevance score.
"""
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(results) == 0:
raise ServiceResourceNotFoundError("No match found")
return results[0]
async def __does_collection_exist(self, cur: Cursor, collection_name: str) -> bool:
results = await self.__get_collections(cur)
return collection_name in results
async def __get_collections(self, cur: Cursor) -> list[str]:
cur.execute(
"""
SELECT table_name
FROM information_schema.tables
WHERE table_schema = %s
""",
(self._schema,),
)
return [row[0] for row in cur.fetchall()]
def _check_dimensionality(self, dimension_num):
if dimension_num > MAX_DIMENSIONALITY:
raise ServiceInitializationError(
f"Dimensionality of {dimension_num} exceeds " + f"the maximum allowed value of {MAX_DIMENSIONALITY}."
)
if dimension_num <= 0:
raise ServiceInitializationError("Dimensionality must be a positive integer. ")
def __serialize_metadata(self, record: MemoryRecord) -> str:
return json.dumps({
"text": record._text,
"description": record._description,
"additional_metadata": record._additional_metadata,
})
# Enable the connection pool to be closed when using as a context manager
def __enter__(self) -> "PostgresMemoryStore":
"""Enter the runtime context."""
return self
def __exit__(self, exc_type, exc_value, traceback) -> bool:
"""Exit the runtime context and dispose of the connection pool."""
self._connection_pool.close()
return False
@@ -0,0 +1,307 @@
# 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
if sys.version_info >= (3, 13):
from warnings import deprecated
else:
from typing_extensions import deprecated
logger: logging.Logger = logging.getLogger(__name__)
@deprecated("This class will be removed in a future version. Use QdrantStore and Collection instead.")
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)
@@ -0,0 +1,31 @@
# semantic_kernel.connectors.memory.redis
This connector uses Redis to implement Semantic Memory. It requires the [RediSearch](https://redis.io/docs/latest/develop/interact/search-and-query/advanced-concepts/vectors/) module to be enabled on Redis to implement vector similarity search.
See the [.net README](https://github.com/microsoft/semantic-kernel/blob/main/dotnet/src/Connectors/Connectors.Memory.Redis/README.md) for more information.
## Quick start
1. Run with Docker:
```bash
docker run -d --name redis-stack-server -p 6379:6379 redis/redis-stack-server:latest
```
2. To use Redis as a semantic memory store:
```python
import semantic_kernel as sk
import semantic_kernel.connectors.ai.open_ai as sk_oai
from semantic_kernel.connectors.memory.redis import RedisMemoryStore
kernel = sk.Kernel()
api_key, org_id = sk.openai_settings_from_dot_env()
kernel.add_service(sk_oai.OpenAITextCompletion(service_id="dv", ai_model_id="text-davinci-003", api_key=api_key, org_id=org_id))
embedding_generator = sk_oai.OpenAITextEmbedding(service_id="ada", ai_model_id="text-embedding-ada-002", api_key=api_key, org_id=org_id)
kernel.add_service(embedding_generator)
redis_connection_string = sk.redis_settings_from_dot_env()
kernel.use_memory(storage=RedisMemoryStore(connection_string=redis_connection_string), embeddings_generator=embedding_generator)
```
@@ -0,0 +1,402 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
import sys
from typing import ClassVar
import numpy as np
import redis
from numpy import ndarray
from pydantic import SecretStr, ValidationError
from redis.commands.search.field import TextField, VectorField
from redis.commands.search.indexDefinition import IndexDefinition, IndexType
from redis.commands.search.query import Query
from redis.exceptions import ResponseError
from semantic_kernel.connectors.memory_stores.redis.utils import (
deserialize_document_to_record,
deserialize_redis_to_record,
get_redis_key,
serialize_record_to_redis,
)
from semantic_kernel.exceptions import ServiceResourceNotFoundError, ServiceResponseException
from semantic_kernel.exceptions.memory_connector_exceptions import MemoryConnectorInitializationError
from semantic_kernel.kernel_pydantic import KernelBaseSettings
from semantic_kernel.memory.memory_record import MemoryRecord
from semantic_kernel.memory.memory_store_base import MemoryStoreBase
from semantic_kernel.utils.feature_stage_decorator import experimental
if sys.version_info >= (3, 13):
from warnings import deprecated
else:
from typing_extensions import deprecated
logger: logging.Logger = logging.getLogger(__name__)
@experimental
class RedisSettings(KernelBaseSettings):
"""Redis model settings.
Args:
- connection_string (str | None):
Redis connection string (Env var REDIS_CONNECTION_STRING)
"""
env_prefix: ClassVar[str] = "REDIS_"
connection_string: SecretStr
@deprecated("This class will be removed in a future version. Use RedisStore and RedisHashsetCollection instead.")
class RedisMemoryStore(MemoryStoreBase):
"""A memory store implementation using Redis."""
_database: "redis.Redis"
_ft: "redis.Redis.ft"
# Without RedisAI, it is currently not possible to retrieve index-specific vector attributes to have
# fully independent collections.
_query_dialect: int
_vector_distance_metric: str
_vector_index_algorithm: str
_vector_size: int
_vector_type: "np.dtype"
_vector_type_str: str
def __init__(
self,
connection_string: str,
vector_size: int = 1536,
vector_distance_metric: str = "COSINE",
vector_type: str = "FLOAT32",
vector_index_algorithm: str = "HNSW",
query_dialect: int = 2,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""RedisMemoryStore is an abstracted interface to interact with a Redis node connection.
See documentation about connections: https://redis-py.readthedocs.io/en/stable/connections.html
See documentation about vector attributes: https://redis.io/docs/stack/search/reference/vectors.
Args:
connection_string (str): Provide connection URL to a Redis instance
vector_size (str): Size of vectors, defaults to 1536
vector_distance_metric (str): Metric for measuring vector distances, defaults to COSINE
vector_type (str): Vector type, defaults to FLOAT32
vector_index_algorithm (str): Indexing algorithm for vectors, defaults to HNSW
query_dialect (int): Query dialect, must be 2 or greater for vector similarity searching, defaults to 2
env_file_path (str | None): Use the environment settings file as a fallback to
environment variables, defaults to False
env_file_encoding (str | None): Encoding of the environment settings file, defaults to "utf-8"
"""
try:
redis_settings = RedisSettings(
connection_string=connection_string,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
except ValidationError as ex:
raise MemoryConnectorInitializationError("Failed to create Redis settings.", ex) from ex
if vector_size <= 0:
raise MemoryConnectorInitializationError("Vector dimension must be a positive integer")
self._database = redis.Redis.from_url(redis_settings.connection_string.get_secret_value())
self._ft = self._database.ft
self._query_dialect = query_dialect
self._vector_distance_metric = vector_distance_metric
self._vector_index_algorithm = vector_index_algorithm
self._vector_type_str = vector_type
self._vector_type = np.float32 if vector_type == "FLOAT32" else np.float64
self._vector_size = vector_size
async def close(self):
"""Closes the Redis database connection."""
logger.info("Closing Redis connection")
self._database.close()
async def create_collection(self, collection_name: str) -> None:
"""Creates a collection.
Implemented as a Redis index containing hashes prefixed with "collection_name:".
If a collection of the name exists, it is left unchanged.
Args:
collection_name (str): Name for a collection of embeddings
"""
if await self.does_collection_exist(collection_name):
logger.info(f'Collection "{collection_name}" already exists.')
else:
index_def = IndexDefinition(prefix=f"{collection_name}:", index_type=IndexType.HASH)
schema = (
TextField(name="key"),
TextField(name="metadata"),
TextField(name="timestamp"),
VectorField(
name="embedding",
algorithm=self._vector_index_algorithm,
attributes={
"TYPE": self._vector_type_str,
"DIM": self._vector_size,
"DISTANCE_METRIC": self._vector_distance_metric,
},
),
)
try:
self._ft(collection_name).create_index(definition=index_def, fields=schema)
except Exception as e:
raise ServiceResponseException(f"Failed to create collection {collection_name}") from e
async def get_collections(self) -> list[str]:
"""Returns a list of names of all collection names present in the data store.
Returns:
List[str]: list of collection names
"""
# Note: FT._LIST is a temporary command that may be deprecated in the future according to Redis
return [name.decode() for name in self._database.execute_command("FT._LIST")]
async def delete_collection(self, collection_name: str, delete_records: bool = True) -> None:
"""Deletes a collection from the data store.
If the collection does not exist, the database is left unchanged.
Args:
collection_name (str): Name for a collection of embeddings
delete_records (bool): Delete all data associated with the collection, default to True
"""
if await self.does_collection_exist(collection_name):
self._ft(collection_name).dropindex(delete_documents=delete_records)
async def does_collection_exist(self, collection_name: str) -> bool:
"""Determines if a collection exists in the data store.
Args:
collection_name (str): Name for a collection of embeddings
Returns:
True if the collection exists, False if not
"""
try:
self._ft(collection_name).info()
return True
except ResponseError:
return False
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.
Note: if the record do not have the same dimensionality configured for the collection,
it will not be detected to belong to the collection in Redis.
Args:
collection_name (str): Name for a collection of embeddings
record (MemoryRecord): Memory record to upsert
Returns:
str: Redis key associated with the upserted memory record
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f'Collection "{collection_name}" does not exist')
# Typical Redis key structure: collection_name:{some identifier}
record._key = get_redis_key(collection_name, record._id)
# Overwrites previous data or inserts new key if not present
# Index registers any hash matching its schema and prefixed with collection_name:
try:
self._database.hset(
record._key,
mapping=serialize_record_to_redis(record, self._vector_type),
)
return record._key
except Exception as e:
raise ServiceResponseException("Could not upsert messages.") from e
async def upsert_batch(self, collection_name: str, records: list[MemoryRecord]) -> list[str]:
"""Upserts a group 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.
Note: if the records do not have the same dimensionality configured for the collection,
they will not be detected to belong to the collection in Redis.
Args:
collection_name (str): Name for a collection of embeddings
records (List[MemoryRecord]): List of memory records to upsert
Returns:
List[str]: Redis keys associated with the upserted memory records
"""
keys = list()
for record in records:
record_key = await self.upsert(collection_name, record)
keys.append(record_key)
return keys
async def get(self, collection_name: str, key: str, with_embedding: bool = False) -> MemoryRecord:
"""Gets a memory record from the data store. Does not guarantee that the collection exists.
Args:
collection_name (str): Name for a collection of embeddings
key (str): ID associated with the memory to get
with_embedding (bool): Include embedding with the memory record, default to False
Returns:
MemoryRecord: The memory record if found, else None
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f'Collection "{collection_name}" does not exist')
internal_key = get_redis_key(collection_name, key)
fields = self._database.hgetall(internal_key)
# Did not find the record
if len(fields) == 0:
return None
record = deserialize_redis_to_record(fields, self._vector_type, with_embedding)
record._key = internal_key
return record
async def get_batch(
self, collection_name: str, keys: list[str], with_embeddings: bool = False
) -> list[MemoryRecord]:
"""Gets a batch of memory records from the data store.
Does not guarantee that the collection exists.
Args:
collection_name (str): Name for a collection of embeddings
keys (List[str]): IDs associated with the memory records to get
with_embeddings (bool): Include embeddings with the memory records, default to False
Returns:
List[MemoryRecord]: The memory records if found, else an empty list
"""
records = list()
for key in keys:
record = await self.get(collection_name, key, with_embeddings)
if record:
records.append(record)
return records
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.
If the key does not exist, do nothing.
Args:
collection_name (str): Name for a collection of embeddings
key (str): ID associated with the memory to remove
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f'Collection "{collection_name}" does not exist')
self._database.delete(get_redis_key(collection_name, key))
async def remove_batch(self, collection_name: str, keys: list[str]) -> None:
"""Removes a batch of memory records from the data store. Does not guarantee that the collection exists.
Args:
collection_name (str): Name for a collection of embeddings
keys (List[str]): IDs associated with the memory records to remove
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f'Collection "{collection_name}" does not exist')
self._database.delete(*[get_redis_key(collection_name, key) for key in keys])
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]]:
"""Get the nearest matches to an embedding using the configured similarity algorithm.
Args:
collection_name (str): Name for a collection of embeddings
embedding (ndarray): Embedding to find the nearest matches to
limit (int): Maximum number of matches to return
min_relevance_score (float): Minimum relevance score of the matches, default to 0.0
with_embeddings (bool): Include embeddings in the resultant memory records, default to False
Returns:
List[Tuple[MemoryRecord, float]]: Records and their relevance scores by descending
order, or an empty list if no relevant matches are found
"""
if not await self.does_collection_exist(collection_name):
raise ServiceResourceNotFoundError(f'Collection "{collection_name}" does not exist')
# Perform a k-nearest neighbors query, score by similarity
query = (
Query(f"*=>[KNN {limit} @embedding $embedding AS vector_score]")
.dialect(self._query_dialect)
.paging(offset=0, num=limit)
.return_fields(
"metadata",
"timestamp",
"embedding",
"vector_score",
)
.sort_by("vector_score", asc=False)
)
query_params = {"embedding": embedding.astype(self._vector_type).tobytes()}
matches = self._ft(collection_name).search(query, query_params).docs
relevant_records = list()
for match in matches:
score = float(match["vector_score"])
# Sorted by descending order
if score < min_relevance_score:
break
record = deserialize_document_to_record(self._database, match, self._vector_type, with_embeddings)
relevant_records.append((record, score))
return relevant_records
async def get_nearest_match(
self,
collection_name: str,
embedding: ndarray,
min_relevance_score: float = 0.0,
with_embedding: bool = False,
) -> tuple[MemoryRecord, float]:
"""Get the nearest match to an embedding using the configured similarity algorithm.
Args:
collection_name (str): Name for a collection of embeddings
embedding (ndarray): Embedding to find the nearest match to
min_relevance_score (float): Minimum relevance score of the match, default to 0.0
with_embedding (bool): Include embedding in the resultant memory record, default to False
Returns:
Tuple[MemoryRecord, float]: Record and the relevance score, or None if not found
"""
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] if len(matches) else None
@@ -0,0 +1,111 @@
# Copyright (c) Microsoft. All rights reserved.
import json
from datetime import datetime
from typing import Any
import numpy as np
from redis.asyncio.client import Redis
from redis.commands.search.document import Document
from semantic_kernel.memory.memory_record import MemoryRecord
def get_redis_key(collection_name: str, record_id: str) -> str: # pragma: no cover
"""Returns the Redis key for an element called record_id within collection_name.
Args:
collection_name (str): Name for a collection of embeddings
record_id (str): ID associated with a memory record
Returns:
str: Redis key in the format collection_name:id
"""
return f"{collection_name}:{record_id}"
def split_redis_key(redis_key: str) -> tuple[str, str]: # pragma: no cover
"""Split a Redis key into its collection name and record ID.
Args:
redis_key (str): Redis key
Returns:
tuple[str, str]: Tuple of the collection name and ID
"""
collection, record_id = redis_key.split(":")
return collection, record_id
def serialize_record_to_redis(record: MemoryRecord, vector_type: np.dtype) -> dict[str, Any]: # pragma: no cover
"""Serialize a MemoryRecord to Redis fields."""
all_metadata = {
"is_reference": record._is_reference,
"external_source_name": record._external_source_name or "",
"id": record._id or "",
"description": record._description or "",
"text": record._text or "",
"additional_metadata": record._additional_metadata or "",
}
return {
"key": record._key or "",
"timestamp": record._timestamp.isoformat() if record._timestamp else "",
"metadata": json.dumps(all_metadata),
"embedding": (record._embedding.astype(vector_type).tobytes() if record._embedding is not None else ""),
}
def deserialize_redis_to_record(
fields: dict[str, Any], vector_type: np.dtype, with_embedding: bool
) -> MemoryRecord: # pragma: no cover
"""Deserialize Redis fields to a MemoryRecord."""
metadata = json.loads(fields[b"metadata"])
record = MemoryRecord(
id=metadata["id"],
is_reference=metadata["is_reference"] is True,
description=metadata["description"],
external_source_name=metadata["external_source_name"],
text=metadata["text"],
additional_metadata=metadata["additional_metadata"],
embedding=None,
)
if fields[b"timestamp"] != b"":
record._timestamp = datetime.fromisoformat(fields[b"timestamp"].decode())
if with_embedding:
# Extract using the vector type, then convert to regular Python float type
record._embedding = np.frombuffer(fields[b"embedding"], dtype=vector_type).astype(float)
return record
def deserialize_document_to_record(
database: Redis, doc: Document, vector_type: np.dtype, with_embedding: bool
) -> MemoryRecord: # pragma: no cover
"""Deserialize document to a MemoryRecord."""
# Document's ID refers to the Redis key
redis_key = doc["id"]
_, id_str = split_redis_key(redis_key)
metadata = json.loads(doc["metadata"])
record = MemoryRecord(
id=id_str,
is_reference=metadata["is_reference"] is True,
description=metadata["description"],
external_source_name=metadata["external_source_name"],
text=metadata["text"],
additional_metadata=metadata["additional_metadata"],
embedding=None,
)
if doc["timestamp"] != "":
record._timestamp = datetime.fromisoformat(doc["timestamp"])
if with_embedding:
# Some bytes are lost when retrieving a document, fetch raw embedding
eb = database.hget(redis_key, "embedding")
record._embedding = np.frombuffer(eb, dtype=vector_type).astype(float)
return record
@@ -0,0 +1,583 @@
# Copyright (c) Microsoft. All rights reserved.
import itertools
import logging
import os
import sys
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
from numpy import ndarray
from usearch.index import BatchMatches, CompiledMetric, Index, Matches, MetricKind, ScalarKind
from semantic_kernel.exceptions import (
ServiceInitializationError,
ServiceInvalidRequestError,
ServiceResourceNotFoundError,
)
from semantic_kernel.memory.memory_record import MemoryRecord
from semantic_kernel.memory.memory_store_base import MemoryStoreBase
if sys.version_info >= (3, 13):
from warnings import deprecated
else:
from typing_extensions import deprecated
logger: logging.Logger = logging.getLogger(__name__)
@dataclass
class _USearchCollection:
"""Represents a collection for USearch with embeddings and related data.
Attributes:
embeddings_index (Index): The index of embeddings.
embeddings_data_table (pa.Table): The PyArrow table holding embeddings data.
embeddings_id_to_label (Dict[str, int]): Mapping of embeddings ID to label.
"""
embeddings_index: Index
embeddings_data_table: pa.Table
embeddings_id_to_label: dict[str, int]
@staticmethod
def create_default(embeddings_index: Index) -> "_USearchCollection":
"""Create a default `_USearchCollection` using a given embeddings index.
Args:
embeddings_index (Index): The index of embeddings to be used for the default collection.
Returns:
_USearchCollection: A default `_USearchCollection` initialized with the given embeddings index.
"""
return _USearchCollection(
embeddings_index,
pa.Table.from_pandas(
pd.DataFrame(columns=_embeddings_data_schema.names),
schema=_embeddings_data_schema,
),
{},
)
# PyArrow Schema definition for the embeddings data from `MemoryRecord`.
_embeddings_data_schema = pa.schema([
pa.field("key", pa.string()),
pa.field("timestamp", pa.timestamp("us")),
pa.field("is_reference", pa.bool_()),
pa.field("external_source_name", pa.string()),
pa.field("id", pa.string()),
pa.field("description", pa.string()),
pa.field("text", pa.string()),
pa.field("additional_metadata", pa.string()),
])
class _CollectionFileType(Enum):
"""Enumeration of file types used for storing collections."""
USEARCH = 0
PARQUET = 1
# Mapping of collection file types to their file extensions.
_collection_file_extensions: dict[_CollectionFileType, str] = {
_CollectionFileType.USEARCH: ".usearch",
_CollectionFileType.PARQUET: ".parquet",
}
def memoryrecords_to_pyarrow_table(records: list[MemoryRecord]) -> pa.Table:
"""Convert a list of `MemoryRecord` to a PyArrow Table."""
records_pylist = [
{attr: getattr(record, "_" + attr) for attr in _embeddings_data_schema.names} for record in records
]
return pa.Table.from_pylist(records_pylist, schema=_embeddings_data_schema)
def pyarrow_table_to_memoryrecords(table: pa.Table, vectors: ndarray | None = None) -> list[MemoryRecord]:
"""Convert a PyArrow Table to a list of MemoryRecords.
Args:
table (pa.Table): The PyArrow Table to convert.
vectors (Optional[ndarray], optional): An array of vectors to include as embeddings in the MemoryRecords.
The length and order of the vectors should match the rows in the table. Defaults to None.
Returns:
List[MemoryRecord]: List of MemoryRecords constructed from the table.
"""
return [
MemoryRecord(**row.to_dict(), embedding=vectors[index] if vectors is not None else None)
for index, row in table.to_pandas().iterrows()
]
@deprecated("This class will be removed in a future version.")
class USearchMemoryStore(MemoryStoreBase):
"""Memory store for searching embeddings with USearch."""
def __init__(
self,
persist_directory: os.PathLike | None = None,
) -> None:
"""Create a USearchMemoryStore instance.
This store helps searching embeddings with USearch, keeping collections in memory.
To save collections to disk, provide the `persist_directory` param.
Collections are saved when `close` is called.
To both save collections and free up memory, call `close`.
When `USearchMemoryStore` is used with a context manager, this will happen automatically.
Otherwise, it should be called explicitly.
Args:
persist_directory (Optional[os.PathLike], default=None): Directory for loading and saving collections.
If None, collections are not loaded nor saved.
"""
self._persist_directory = Path(persist_directory) if persist_directory is not None else None
self._collections: dict[str, _USearchCollection] = {}
if self._persist_directory:
self._collections = self._read_collections_from_dir()
def _get_collection_path(self, collection_name: str, *, file_type: _CollectionFileType) -> Path:
"""Get the path for the given collection name and file type.
Args:
collection_name (str): Name of the collection.
file_type (_CollectionFileType): The file type.
Returns:
Path: Path to the collection file.
Raises:
ValueError: If persist directory path is not set.
"""
collection_name = collection_name.lower()
if self._persist_directory is None:
raise ServiceInitializationError("Path of persist directory is not set")
return self._persist_directory / (collection_name + _collection_file_extensions[file_type])
async def create_collection(
self,
collection_name: str,
ndim: int = 0,
metric: str | MetricKind | CompiledMetric = MetricKind.IP,
dtype: str | ScalarKind | None = None,
connectivity: int | None = None,
expansion_add: int | None = None,
expansion_search: int | None = None,
view: bool = False,
) -> None:
"""Create a new collection.
Args:
collection_name (str): Name of the collection. Case-insensitive.
Must have name that is valid file name for the current OS environment.
ndim (int, optional): Number of dimensions. Defaults to 0.
metric (Union[str, MetricKind, CompiledMetric], optional): Metric kind. Defaults to MetricKind.IP.
dtype (Optional[Union[str, ScalarKind]], optional): Data type. Defaults to None.
connectivity (int, optional): Connectivity parameter. Defaults to None.
expansion_add (int, optional): Expansion add parameter. Defaults to None.
expansion_search (int, optional): Expansion search parameter. Defaults to None.
view (bool, optional): Viewing flag. Defaults to False.
Raises:
ValueError: If collection with the given name already exists.
ValueError: If collection name is empty string.
"""
collection_name = collection_name.lower()
if not collection_name:
raise ServiceInvalidRequestError("Collection name can not be empty.")
if collection_name in self._collections:
raise ServiceInvalidRequestError(f"Collection with name {collection_name} already exists.")
embeddings_index_path = (
self._get_collection_path(collection_name, file_type=_CollectionFileType.USEARCH)
if self._persist_directory
else None
)
embeddings_index = Index(
ndim=ndim,
metric=metric,
dtype=dtype,
connectivity=connectivity,
expansion_add=expansion_add,
expansion_search=expansion_search,
path=embeddings_index_path,
view=view,
)
self._collections[collection_name] = _USearchCollection.create_default(embeddings_index)
return
def _read_embeddings_table(self, path: os.PathLike) -> tuple[pa.Table, dict[str, int]]:
"""Read embeddings from the provided path and generate an ID to label mapping.
Args:
path (os.PathLike): Path to the embeddings.
Returns:
Tuple of embeddings table and a dictionary mapping from record ID to its label.
"""
embeddings_table = pq.read_table(path, schema=_embeddings_data_schema)
embeddings_id_to_label: dict[str, int] = {
record_id: idx for idx, record_id in enumerate(embeddings_table.column("id").to_pylist())
}
return embeddings_table, embeddings_id_to_label
def _read_embeddings_index(self, path: Path) -> Index:
"""Read embeddings index."""
# str cast is temporarily fix for https://github.com/unum-cloud/usearch/issues/196
return Index.restore(str(path), view=False)
def _read_collections_from_dir(self) -> dict[str, _USearchCollection]:
"""Read all collections from directory to memory.
Raises:
ValueError: If files for a collection do not match expected amount.
Returns:
Dict[str, _USearchCollection]: Dictionary with collection names as keys and
their _USearchCollection as values.
"""
collections: dict[str, _USearchCollection] = {}
for collection_name, collection_files in self._get_all_storage_files().items():
expected_storage_files = len(_CollectionFileType)
if len(collection_files) != expected_storage_files:
raise ServiceInitializationError(
f"Expected {expected_storage_files} files for collection {collection_name}"
)
parquet_file, usearch_file = collection_files
if parquet_file.suffix == _collection_file_extensions[_CollectionFileType.USEARCH]:
parquet_file, usearch_file = usearch_file, parquet_file
embeddings_table, embeddings_id_to_label = self._read_embeddings_table(parquet_file)
embeddings_index = self._read_embeddings_index(usearch_file)
collections[collection_name] = _USearchCollection(
embeddings_index,
embeddings_table,
embeddings_id_to_label,
)
return collections
async def get_collections(self) -> list[str]:
"""Get list of existing collections.
Returns:
List[str]: List of collection names.
"""
return list(self._collections.keys())
async def delete_collection(self, collection_name: str) -> None:
"""Delete collection by name."""
collection_name = collection_name.lower()
collection = self._collections.pop(collection_name, None)
if collection:
collection.embeddings_index.reset()
return
async def does_collection_exist(self, collection_name: str) -> bool:
"""Check if collection exists."""
collection_name = collection_name.lower()
return collection_name in self._collections
async def upsert(self, collection_name: str, record: MemoryRecord) -> str:
"""Upsert single MemoryRecord and return its ID."""
collection_name = collection_name.lower()
res = await self.upsert_batch(collection_name=collection_name, records=[record])
return res[0]
async def upsert_batch(
self,
collection_name: str,
records: list[MemoryRecord],
*,
compact: bool = False,
copy: bool = True,
threads: int = 0,
log: str | bool = False,
batch_size: int = 0,
) -> list[str]:
"""Upsert a batch of MemoryRecords and return their IDs.
Args:
collection_name (str): Name of the collection to search within.
records (List[MemoryRecord]): Records to upsert.
compact (bool, optional): Removes links to removed nodes (expensive). Defaults to False.
copy (bool, optional): Should the index store a copy of vectors. Defaults to True.
threads (int, optional): Optimal number of cores to use. Defaults to 0.
log (Union[str, bool], optional): Whether to print the progress bar. Defaults to False.
batch_size (int, optional): Number of vectors to process at once. Defaults to 0.
Raises:
KeyError: If collection not exist
Returns:
List[str]: List of IDs.
"""
collection_name = collection_name.lower()
if collection_name not in self._collections:
raise ServiceResourceNotFoundError(f"Collection {collection_name} does not exist, cannot insert.")
ucollection = self._collections[collection_name]
all_records_id = [record._id for record in records]
# Remove vectors from index
remove_labels = [
ucollection.embeddings_id_to_label[id] for id in all_records_id if id in ucollection.embeddings_id_to_label
]
ucollection.embeddings_index.remove(remove_labels, compact=compact, threads=threads)
# Determine label insertion points
table_num_rows = ucollection.embeddings_data_table.num_rows
insert_labels = np.arange(table_num_rows, table_num_rows + len(records))
# Add embeddings to index
ucollection.embeddings_index.add(
keys=insert_labels,
vectors=np.stack([record.embedding for record in records]),
copy=copy,
threads=threads,
log=log,
)
# Update embeddings_table
ucollection.embeddings_data_table = pa.concat_tables([
ucollection.embeddings_data_table,
memoryrecords_to_pyarrow_table(records),
])
# Update embeddings_id_to_label
for index, record_id in enumerate(all_records_id):
ucollection.embeddings_id_to_label[record_id] = insert_labels[index]
return all_records_id
async def get(
self,
collection_name: str,
key: str,
with_embedding: bool,
dtype: ScalarKind = ScalarKind.F32,
) -> MemoryRecord:
"""Retrieve a single MemoryRecord using its key."""
collection_name = collection_name.lower()
result = await self.get_batch(
collection_name=collection_name,
keys=[key],
with_embeddings=with_embedding,
dtype=dtype,
)
if not result:
raise ServiceResourceNotFoundError(f"Key '{key}' not found in collection '{collection_name}'")
return result[0]
async def get_batch(
self,
collection_name: str,
keys: list[str],
with_embeddings: bool,
dtype: ScalarKind = ScalarKind.F32,
) -> list[MemoryRecord]:
"""Retrieve a batch of MemoryRecords using their keys."""
collection_name = collection_name.lower()
if collection_name not in self._collections:
raise ServiceResourceNotFoundError(f"Collection {collection_name} does not exist")
ucollection = self._collections[collection_name]
labels = [ucollection.embeddings_id_to_label[key] for key in keys if key in ucollection.embeddings_id_to_label]
if not labels:
return []
vectors = ucollection.embeddings_index.get(labels, dtype) if with_embeddings else None
return pyarrow_table_to_memoryrecords(ucollection.embeddings_data_table.take(pa.array(labels)), vectors)
async def remove(self, collection_name: str, key: str) -> None:
"""Remove a single MemoryRecord using its key."""
collection_name = collection_name.lower()
await self.remove_batch(collection_name=collection_name, keys=[key])
return
async def remove_batch(self, collection_name: str, keys: list[str]) -> None:
"""Remove a batch of MemoryRecords using their keys."""
collection_name = collection_name.lower()
if collection_name not in self._collections:
raise ServiceResourceNotFoundError(f"Collection {collection_name} does not exist, cannot insert.")
ucollection = self._collections[collection_name]
labels = [ucollection.embeddings_id_to_label[key] for key in keys]
ucollection.embeddings_index.remove(labels)
for key in keys:
del ucollection.embeddings_id_to_label[key]
return
async def get_nearest_match(
self,
collection_name: str,
embedding: ndarray,
min_relevance_score: float = 0.0,
with_embedding: bool = True,
exact: bool = False,
) -> tuple[MemoryRecord, float]:
"""Retrieve the nearest matching MemoryRecord for the provided embedding.
By default it is approximately search, see `exact` param description.
Measure of similarity between vectors is relevance score. It is from 0 to 1.
USearch returns distances for vectors. Distance is converted to relevance score by inverse function.
Args:
collection_name (str): Name of the collection to search within.
embedding (ndarray): The embedding vector to search for.
min_relevance_score (float, optional): The minimum relevance score for vectors. Supposed to be from 0 to 1.
Only vectors with greater or equal relevance score are returned. Defaults to 0.0.
with_embedding (bool, optional): If True, include the embedding in the result. Defaults to True.
exact (bool, optional): Perform exhaustive linear-time exact search. Defaults to False.
Returns:
Tuple[MemoryRecord, float]: The nearest matching record and its relevance score.
"""
collection_name = collection_name.lower()
results = await self.get_nearest_matches(
collection_name=collection_name,
embedding=embedding,
limit=1,
min_relevance_score=min_relevance_score,
with_embeddings=with_embedding,
exact=exact,
)
return results[0]
async def get_nearest_matches(
self,
collection_name: str,
embedding: ndarray,
limit: int,
min_relevance_score: float = 0.0,
with_embeddings: bool = True,
*,
threads: int = 0,
exact: bool = False,
log: str | bool = False,
batch_size: int = 0,
) -> list[tuple[MemoryRecord, float]]:
"""Get the nearest matches to a given embedding.
By default it is approximately search, see `exact` param description.
Measure of similarity between vectors is relevance score. It is from 0 to 1.
USearch returns distances for vectors. Distance is converted to relevance score by inverse function.
Args:
collection_name (str): Name of the collection to search within.
embedding (ndarray): The embedding vector to search for.
limit (int): maximum amount of embeddings to search for.
min_relevance_score (float, optional): The minimum relevance score for vectors. Supposed to be from 0 to 1.
Only vectors with greater or equal relevance score are returned. Defaults to 0.0.
with_embeddings (bool, optional): If True, include the embedding in the result. Defaults to True.
threads (int, optional): Optimal number of cores to use. Defaults to 0.
exact (bool, optional): Perform exhaustive linear-time exact search. Defaults to False.
log (Union[str, bool], optional): Whether to print the progress bar. Defaults to False.
batch_size (int, optional): Number of vectors to process at once. Defaults to 0.
Raises:
KeyError: if a collection with specified name does not exist
Returns:
List[Tuple[MemoryRecord, float]]: The nearest matching records and their relevance score.
"""
collection_name = collection_name.lower()
ucollection = self._collections[collection_name]
result: Matches | BatchMatches = ucollection.embeddings_index.search(
vectors=embedding,
count=limit,
threads=threads,
exact=exact,
log=log,
)
relevance_score = 1 / (result.distances + 1)
filtered_labels = result.keys[np.where(relevance_score >= min_relevance_score)[0]]
filtered_vectors: np.ndarray | None = None
if with_embeddings:
filtered_vectors = ucollection.embeddings_index.get(filtered_labels)
return [
(mem_rec, relevance_score[index].item())
for index, mem_rec in enumerate(
pyarrow_table_to_memoryrecords(
ucollection.embeddings_data_table.take(pa.array(filtered_labels)),
filtered_vectors,
)
)
]
def _get_all_storage_files(self) -> dict[str, list[Path]]:
"""Return storage files for each collection in `self._persist_directory`.
Collection name is derived from file name and converted to lowercase. Files with extensions that
do not match storage extensions are discarded.
Raises:
ValueError: If persist directory is not set.
Returns:
Dict[str, List[Path]]: Dictionary of collection names mapped to their respective files.
"""
if self._persist_directory is None:
raise ServiceInitializationError("Persist directory is not set")
storage_exts = _collection_file_extensions.values()
collection_storage_files: dict[str, list[Path]] = {}
for path in self._persist_directory.iterdir():
if path.is_file() and (path.suffix in storage_exts):
collection_name = path.stem.lower()
if collection_name in collection_storage_files:
collection_storage_files[collection_name].append(path)
else:
collection_storage_files[collection_name] = [path]
return collection_storage_files
def _dump_collections(self) -> None:
collection_storage_files = self._get_all_storage_files()
for file_path in itertools.chain.from_iterable(collection_storage_files.values()):
file_path.unlink()
for collection_name, ucollection in self._collections.items():
ucollection.embeddings_index.save(
self._get_collection_path(collection_name, file_type=_CollectionFileType.USEARCH)
)
pq.write_table(
ucollection.embeddings_data_table,
self._get_collection_path(collection_name, file_type=_CollectionFileType.PARQUET),
)
return
async def close(self) -> None:
"""Persist collection, clear.
Returns:
None
"""
if self._persist_directory:
self._dump_collections()
for collection_name in await self.get_collections():
await self.delete_collection(collection_name)
self._collections = {}
@@ -0,0 +1,17 @@
# Weaviate Memory Connector
[Weaviate](https://weaviate.io/developers/weaviate) is an open source vector database. Semantic Kernel provides a connector to allow you to store and retrieve information for you AI applications from a Weaviate database.
## Setup
There are a few ways you can deploy your Weaviate database:
- [Weaviate Cloud](https://weaviate.io/developers/weaviate/installation/weaviate-cloud-services)
- [Docker](https://weaviate.io/developers/weaviate/installation/docker-compose)
- [Embedded](https://weaviate.io/developers/weaviate/installation/embedded)
- Other cloud providers such as [Azure](https://azuremarketplace.microsoft.com/en-us/marketplace/apps/weaviatebv1686614539420.weaviate_1?tab=Overview), [AWS](https://weaviate.io/developers/weaviate/installation/aws-marketplace) or [GCP](https://weaviate.io/developers/weaviate/installation/gc-marketplace).
> Note that embedded mode is not supported on Windows yet: [GitHub issue](https://github.com/weaviate/weaviate/issues/3315) and it's still an experimental feature on Linux and MacOS.
## Using the Connector
Once the Weaviate database is up and running, and the environment variables are set, you can use the connector in your Semantic Kernel application. Please refer to this sample to see how to use the connector: [Complex Connector Sample](../../../../samples/concepts/memory/complex_memory.py)
@@ -0,0 +1,325 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import logging
import sys
from typing import Final
import numpy as np
import weaviate
from semantic_kernel.connectors.weaviate import WeaviateSettings
from semantic_kernel.exceptions.memory_connector_exceptions import MemoryConnectorInitializationError
from semantic_kernel.memory.memory_record import MemoryRecord
from semantic_kernel.memory.memory_store_base import MemoryStoreBase
if sys.version_info >= (3, 13):
from warnings import deprecated
else:
from typing_extensions import deprecated
logger: logging.Logger = logging.getLogger(__name__)
SCHEMA = {
"class": "MemoryRecord",
"description": "A document from semantic kernel.",
"properties": [
{
"name": "key",
"description": "The key of the record.",
"dataType": ["string"],
},
{
"name": "timestamp",
"description": "The timestamp of the record.",
"dataType": ["date"],
},
{
"name": "isReference",
"description": "Whether the record is a reference record.",
"dataType": ["boolean"],
},
{
"name": "externalSourceName",
"description": "The name of the external source.",
"dataType": ["string"],
},
{
"name": "skId",
"description": "A unique identifier for the record.",
"dataType": ["string"],
},
{
"name": "description",
"description": "The description of the record.",
"dataType": ["text"],
},
{
"name": "text",
"description": "The text of the record.",
"dataType": ["text"],
},
{
"name": "additionalMetadata",
"description": "Optional custom metadata of the record.",
"dataType": ["string"],
},
],
}
ALL_PROPERTIES = [property["name"] for property in SCHEMA["properties"]]
class FieldMapper:
"""This maps attribute names between the SK's memory record and weaviate's schema.
It provides methods for converting between the two naming conventions.
"""
SK_TO_WEAVIATE_MAPPING: Final[dict[str, str]] = {
"_key": "key",
"_timestamp": "timestamp",
"_is_reference": "isReference",
"_external_source_name": "externalSourceName",
"_id": "skId",
"_description": "description",
"_text": "text",
"_additional_metadata": "additionalMetadata",
"_embedding": "vector",
}
WEAVIATE_TO_SK_MAPPING: Final[dict[str, str]] = {v: k for k, v in SK_TO_WEAVIATE_MAPPING.items()}
@classmethod
def sk_to_weaviate(cls, sk_dict):
"""Used to convert a MemoryRecord to a dict of attribute-values that can be used by Weaviate."""
return {cls.SK_TO_WEAVIATE_MAPPING.get(k, k): v for k, v in sk_dict.items() if k in cls.SK_TO_WEAVIATE_MAPPING}
@classmethod
def weaviate_to_sk(cls, weaviate_dict):
"""Used to convert a Weaviate object to a dict that can be used to initialize a MemoryRecord."""
return {
cls.WEAVIATE_TO_SK_MAPPING.get(k, k): v for k, v in weaviate_dict.items() if k in cls.WEAVIATE_TO_SK_MAPPING
}
@classmethod
def remove_underscore_prefix(cls, sk_dict):
"""Used to initialize a MemoryRecord from a SK's dict of private attribute-values."""
return {key.lstrip("_"): value for key, value in sk_dict.items()}
@deprecated(
"WeaviateMemoryStore is deprecated and will be removed in a future version. "
"Please use WeaviateStore and Collection instead."
)
class WeaviateMemoryStore(MemoryStoreBase):
"""A memory store that uses Weaviate as the backend."""
def __init__(
self,
url: str | None = None,
api_key: str | None = None,
use_embed: bool = False,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
):
"""Initializes a new instance of the WeaviateMemoryStore.
Args:
url (str): The URL of the Weaviate instance.
api_key (str): The API key to use for authentication.
use_embed (bool): Whether to use the client embedding options.
env_file_path (str): Whether to use the environment settings (.env) file.
env_file_encoding (str): The encoding of the environment settings (.env) file. Defaults to 'utf-8'.
"""
self.settings = WeaviateSettings(
url=url,
api_key=api_key,
use_embed=use_embed,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
if self.settings.use_embed:
self.client = weaviate.Client(embedded_options=weaviate.EmbeddedOptions())
elif self.settings.api_key and self.settings.url:
self.client = weaviate.Client(
url=str(self.settings.url),
auth_client_secret=weaviate.auth.AuthApiKey(api_key=self.settings.api_key.get_secret_value()),
)
elif self.settings.url:
self.client = weaviate.Client(url=str(self.settings.url))
else:
raise MemoryConnectorInitializationError("WeaviateMemoryStore requires a URL or API key, or to use embed.")
async def create_collection(self, collection_name: str) -> None:
"""Creates a new collection in Weaviate."""
schema = SCHEMA.copy()
schema["class"] = collection_name
await asyncio.get_running_loop().run_in_executor(None, self.client.schema.create_class, schema)
async def get_collections(self) -> list[str]:
"""Returns a list of all collections in Weaviate."""
schemas = await asyncio.get_running_loop().run_in_executor(None, self.client.schema.get)
return [schema["class"] for schema in schemas["classes"]]
async def delete_collection(self, collection_name: str) -> bool:
"""Deletes a collection in Weaviate."""
await asyncio.get_running_loop().run_in_executor(None, self.client.schema.delete_class, collection_name)
async def does_collection_exist(self, collection_name: str) -> bool:
"""Checks if a collection exists in Weaviate."""
collections = await self.get_collections()
return collection_name in collections
async def upsert(self, collection_name: str, record: MemoryRecord) -> str:
"""Upserts a record into Weaviate."""
weaviate_record = FieldMapper.sk_to_weaviate(vars(record))
vector = weaviate_record.pop("vector", None)
weaviate_id = weaviate.util.generate_uuid5(weaviate_record, collection_name)
return await asyncio.get_running_loop().run_in_executor(
None,
self.client.data_object.create,
weaviate_record,
collection_name,
weaviate_id,
vector,
)
async def upsert_batch(self, collection_name: str, records: list[MemoryRecord]) -> list[str]:
"""Upserts a batch of records into Weaviate."""
def _upsert_batch_inner():
results = []
with self.client.batch as batch:
for record in records:
weaviate_record = FieldMapper.sk_to_weaviate(vars(record))
vector = weaviate_record.pop("vector", None)
weaviate_id = weaviate.util.generate_uuid5(weaviate_record, collection_name)
batch.add_data_object(
data_object=weaviate_record,
uuid=weaviate_id,
vector=vector,
class_name=collection_name,
)
results.append(weaviate_id)
return results
return await asyncio.get_running_loop().run_in_executor(None, _upsert_batch_inner)
async def get(self, collection_name: str, key: str, with_embedding: bool) -> MemoryRecord:
"""Gets a record from Weaviate by key."""
# Call the batched version with a single key
results = await self.get_batch(collection_name, [key], with_embedding)
return results[0] if results else None
async def get_batch(self, collection_name: str, keys: list[str], with_embedding: bool) -> list[MemoryRecord]:
"""Gets a batch of records from Weaviate by keys."""
queries = self._build_multi_get_query(collection_name, keys, with_embedding)
results = await asyncio.get_running_loop().run_in_executor(None, self.client.query.multi_get(queries).do)
get_dict = results.get("data", {}).get("Get", {})
return [self._convert_weaviate_doc_to_memory_record(doc) for docs in get_dict.values() for doc in docs]
def _build_multi_get_query(self, collection_name: str, keys: list[str], with_embedding: bool):
queries = []
for i, key in enumerate(keys):
query = self.client.query.get(collection_name, ALL_PROPERTIES).with_where({
"path": ["key"],
"operator": "Equal",
"valueString": key,
})
if with_embedding:
query = query.with_additional("vector")
query = query.with_alias(f"query_{i}")
queries.append(query)
return queries
def _convert_weaviate_doc_to_memory_record(self, weaviate_doc: dict) -> MemoryRecord:
weaviate_doc_copy = weaviate_doc.copy()
vector = weaviate_doc_copy.pop("_additional", {}).get("vector")
weaviate_doc_copy["vector"] = np.array(vector) if vector else None
sk_doc = FieldMapper.weaviate_to_sk(weaviate_doc_copy)
mem_vals = FieldMapper.remove_underscore_prefix(sk_doc)
return MemoryRecord(**mem_vals)
async def remove(self, collection_name: str, key: str) -> None:
"""Removes a record from Weaviate by key."""
await self.remove_batch(collection_name, [key])
async def remove_batch(self, collection_name: str, keys: list[str]) -> None:
"""Removes a batch of records from Weaviate by keys."""
for key in keys:
where = {
"path": ["key"],
"operator": "Equal",
"valueString": key,
}
await asyncio.get_running_loop().run_in_executor(
None, self.client.batch.delete_objects, collection_name, where
)
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]]:
"""Gets the nearest matches to an embedding in Weaviate."""
nearVector = {
"vector": embedding,
"certainty": min_relevance_score,
}
additional = ["certainty"]
if with_embeddings:
additional.append("vector")
query = (
self.client.query.get(collection_name, ALL_PROPERTIES)
.with_near_vector(nearVector)
.with_additional(additional)
.with_limit(limit)
)
results = await asyncio.get_running_loop().run_in_executor(None, query.do)
get_dict = results.get("data", {}).get("Get", {})
return [
(
self._convert_weaviate_doc_to_memory_record(doc),
item["_additional"]["certainty"],
)
for items in get_dict.values()
for item in items
for doc in [item]
]
async def get_nearest_match(
self,
collection_name: str,
embedding: np.ndarray,
min_relevance_score: float,
with_embedding: bool,
) -> tuple[MemoryRecord, float]:
"""Gets the nearest match to an embedding in Weaviate."""
results = await self.get_nearest_matches(
collection_name,
embedding,
limit=1,
min_relevance_score=min_relevance_score,
with_embeddings=with_embedding,
)
return results[0]