chore: import upstream snapshot with attribution
This commit is contained in:
@@ -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([]),
|
||||
)
|
||||
+422
@@ -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")
|
||||
+297
@@ -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)
|
||||
+188
@@ -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
|
||||
+353
@@ -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"]
|
||||
+197
@@ -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.
|
||||
+333
@@ -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]
|
||||
Reference in New Issue
Block a user