diff --git a/python/semantic_kernel/connectors/in_memory.py b/python/semantic_kernel/connectors/in_memory.py index f2e1dc30fbe1..e770a5dff425 100644 --- a/python/semantic_kernel/connectors/in_memory.py +++ b/python/semantic_kernel/connectors/in_memory.py @@ -6,7 +6,7 @@ from typing import Any, ClassVar, Final, Generic, TypeVar, cast from numpy import dot -from pydantic import Field +from pydantic import BaseModel, ConfigDict, Field from scipy.spatial.distance import cityblock, cosine, euclidean, hamming, sqeuclidean from typing_extensions import override @@ -27,7 +27,7 @@ ) from semantic_kernel.exceptions import VectorSearchExecutionException, VectorStoreModelValidationError from semantic_kernel.exceptions.vector_store_exceptions import VectorStoreModelException, VectorStoreOperationException -from semantic_kernel.kernel_types import OneOrMany +from semantic_kernel.kernel_types import OneOrList, OneOrMany from semantic_kernel.utils.feature_stage_decorator import release_candidate from semantic_kernel.utils.list_handler import empty_generator @@ -81,6 +81,15 @@ def __delattr__(self, name) -> None: raise AttributeError(name) +class _InMemoryRecord(dict[str, Any]): + """A serialized record with the name of its key field.""" + + def __init__(self, record: dict[str, Any], key_field_name: str): + """Keep the dumped fields and the name of their key field.""" + super().__init__(record) + self.key_field_name = key_field_name + + class ReadOnlyAttributeDict(Mapping[TAKey, TAValue], Generic[TAKey, TAValue]): """A read-only mapping that allows attribute access to keys.""" @@ -632,14 +641,41 @@ async def _inner_get( async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence[TKey]: updated_keys = [] for record in records: + key_field_name = ( + record.key_field_name if isinstance(record, _InMemoryRecord) else self._key_field_storage_name + ) record = AttributeDict(record) - self.inner_storage[record[self._key_field_name]] = record - updated_keys.append(record[self._key_field_name]) + if key_field_name not in record: + key_field_name = ( + self._key_field_storage_name if self._key_field_storage_name in record else self._key_field_name + ) + key = record[key_field_name] + self.inner_storage[key] = record + updated_keys.append(key) return updated_keys def _deserialize_store_models_to_dicts(self, records: Sequence[Any], **kwargs: Any) -> Sequence[dict[str, Any]]: return records + def _serialize_data_model_to_dict(self, record: TModel, **kwargs: Any) -> OneOrList[dict[str, Any]]: + """Keep track of key names before records lose their model type.""" + serialized = super()._serialize_data_model_to_dict(record, **kwargs) + key_field_name = self._key_field_name + if not self.definition.to_dict: + if not isinstance(record, BaseModel): + key_field_name = self._key_field_storage_name + elif ( + type(record).model_dump is BaseModel.model_dump + and not type(record).__pydantic_decorators__.model_serializers + and "serialize_by_alias" in ConfigDict.__annotations__ + and type(record).model_config.get("serialize_by_alias") + and (key_field := type(record).model_fields.get(key_field_name)) + ): + key_field_name = key_field.serialization_alias or key_field_name + if isinstance(serialized, list): + return [_InMemoryRecord(item, key_field_name) for item in serialized] + return _InMemoryRecord(serialized, key_field_name) + def _serialize_dicts_to_store_models(self, records: Sequence[dict[str, Any]], **kwargs: Any) -> Sequence[Any]: return records diff --git a/python/tests/unit/connectors/memory/test_in_memory.py b/python/tests/unit/connectors/memory/test_in_memory.py index 7955b7696632..5447f267a68f 100644 --- a/python/tests/unit/connectors/memory/test_in_memory.py +++ b/python/tests/unit/connectors/memory/test_in_memory.py @@ -2,6 +2,7 @@ import ast +from pydantic import ConfigDict, Field, model_serializer from pytest import fixture, mark, raises from semantic_kernel.connectors.in_memory import InMemoryCollection, InMemoryStore @@ -35,6 +36,209 @@ async def test_upsert(collection): assert collection.inner_storage == {"testid": record} +@mark.parametrize("batch", [False, True]) +@mark.parametrize("model_kind", ["dict", "dataclass", "pydantic"]) +async def test_upsert_with_key_storage_name(definition, dataclass_vector_data_model, record_type, batch, model_kind): + definition.key_field.storage_name = "stored_id" + model_type = {"dict": dict, "dataclass": dataclass_vector_data_model, "pydantic": record_type}[model_kind] + collection = InMemoryCollection(collection_name="test", record_type=model_type, definition=definition) + records = [ + model_type(id="first", content="first content", vector=[1.0, 0.0, 0.0, 0.0, 0.0]), + model_type(id="second", content="second content", vector=[0.0, 1.0, 0.0, 0.0, 0.0]), + ] + + keys = await collection.upsert(records if batch else records[0]) + assert keys == (["first", "second"] if batch else "first") + assert await collection.get("first", include_vectors=True) == records[0] + if batch: + assert await collection.get(["first", "second"], include_vectors=True) == records + + results = await collection.search(vector=[1.0, 0.0, 0.0, 0.0, 0.0], include_vectors=True) + found = [result.record async for result in results.results] + assert found == (records if batch else records[:1]) + + await collection.delete("first") + assert await collection.get("first") is None + if batch: + assert await collection.get("second", include_vectors=True) == records[1] + + +@mark.parametrize("batch", [False, True]) +async def test_upsert_when_storage_name_matches_logical_field(definition, record_type, batch): + definition.key_field.storage_name = "content" + collection = InMemoryCollection(collection_name="test", record_type=record_type, definition=definition) + record = record_type(id="first", content="unrelated", vector=[1.0, 0.0, 0.0, 0.0, 0.0]) + + keys = await collection.upsert([record] if batch else record) + assert keys == (["first"] if batch else "first") + assert "first" in collection.inner_storage + assert "unrelated" not in collection.inner_storage + + await collection.delete("first") + assert collection.inner_storage == {} + + +@mark.parametrize("batch", [False, True]) +@mark.parametrize("serializer", ["default", "serialize_none", "model_serialize_none"]) +async def test_upsert_when_key_storage_name_matches_another_logical_field(definition, record_type, batch, serializer): + definition.key_field.storage_name = "content" + if serializer == "serialize_none": + definition.serialize = lambda record, **kwargs: None + elif serializer == "model_serialize_none": + + class CustomRecord(record_type): + def serialize(self, **kwargs): + return None + + record_type = CustomRecord + collection = InMemoryCollection(collection_name="test", record_type=record_type, definition=definition) + records = [ + record_type(id="first", content="unrelated first", vector=[1.0, 0.0, 0.0, 0.0, 0.0]), + record_type(id="second", content="unrelated second", vector=[0.0, 1.0, 0.0, 0.0, 0.0]), + ] + + keys = await collection.upsert(records if batch else records[0]) + assert keys == (["first", "second"] if batch else "first") + assert set(collection.inner_storage) == ({"first", "second"} if batch else {"first"}) + assert collection.inner_storage["first"] == records[0].model_dump() + + await collection.delete(["first", "second"] if batch else "first") + assert collection.inner_storage == {} + + +@mark.parametrize("batch", [False, True]) +@mark.parametrize("model_kind", ["dict", "dataclass"]) +async def test_upsert_when_another_storage_name_matches_key(definition, dataclass_vector_data_model, batch, model_kind): + definition.key_field.storage_name = "stored_id" + definition.fields[1].storage_name = "id" + model_type = dict if model_kind == "dict" else dataclass_vector_data_model + collection = InMemoryCollection(collection_name="test", record_type=model_type, definition=definition) + records = [ + model_type(id="first", content="unrelated first", vector=[1.0, 0.0, 0.0, 0.0, 0.0]), + model_type(id="second", content="unrelated second", vector=[0.0, 1.0, 0.0, 0.0, 0.0]), + ] + + keys = await collection.upsert(records if batch else records[0]) + assert keys == (["first", "second"] if batch else "first") + assert await collection.get("first", include_vectors=True) == records[0] + if batch: + assert await collection.get(["first", "second"], include_vectors=True) == records + + await collection.delete(["first", "second"] if batch else "first") + assert collection.inner_storage == {} + + +@mark.parametrize("batch", [False, True]) +@mark.parametrize("serializer", ["to_dict", "serialize", "model_serialize", "model_dump", "model_serializer"]) +async def test_upsert_with_pydantic_custom_storage_serialization(definition, record_type, batch, serializer): + definition.key_field.storage_name = "stored_id" + content_name = "id" if serializer in {"serialize", "model_serialize"} else "stored_content" + definition.fields[1].storage_name = content_name + if serializer == "model_serialize": + + class CustomRecord(record_type): + def serialize(self, **kwargs): + return {"stored_id": self.id, content_name: self.content, "vector": self.vector} + + model_type = CustomRecord + elif serializer == "model_dump": + + class CustomDumpRecord(record_type): + def model_dump(self, **kwargs): + return {"stored_id": self.id, content_name: self.content, "vector": self.vector} + + model_type = CustomDumpRecord + elif serializer == "model_serializer": + + class CustomSerializedRecord(record_type): + @model_serializer + def serialize_model(self): + return {"stored_id": self.id, content_name: self.content, "vector": self.vector} + + model_type = CustomSerializedRecord + else: + model_type = record_type + setattr( + definition, + serializer, + lambda record, **kwargs: {"stored_id": record.id, content_name: record.content, "vector": record.vector}, + ) + collection = InMemoryCollection(collection_name="test", record_type=model_type, definition=definition) + records = [ + model_type(id="first", content="unrelated first", vector=[1.0, 0.0, 0.0, 0.0, 0.0]), + model_type(id="second", content="unrelated second", vector=[0.0, 1.0, 0.0, 0.0, 0.0]), + ] + + keys = await collection.upsert(records if batch else records[0]) + assert keys == (["first", "second"] if batch else "first") + assert set(collection.inner_storage) == ({"first", "second"} if batch else {"first"}) + + await collection.delete(["first", "second"] if batch else "first") + assert collection.inner_storage == {} + + +@mark.parametrize("batch", [False, True]) +async def test_upsert_with_pydantic_serialization_aliases(definition, record_type, batch): + definition.key_field.storage_name = "stored_id" + definition.fields[1].storage_name = "id" + + class AliasedRecord(record_type): + model_config = ConfigDict(serialize_by_alias=True) + id: str = Field(serialization_alias="stored_id") + content: str = Field(serialization_alias="id") + + collection = InMemoryCollection(collection_name="test", record_type=AliasedRecord, definition=definition) + records = [ + AliasedRecord(id="first", content="unrelated first", vector=[1.0, 0.0, 0.0, 0.0, 0.0]), + AliasedRecord(id="second", content="unrelated second", vector=[0.0, 1.0, 0.0, 0.0, 0.0]), + ] + + keys = await collection.upsert(records if batch else records[0]) + assert keys == (["first", "second"] if batch else "first") + assert set(collection.inner_storage) == ({"first", "second"} if batch else {"first"}) + + await collection.delete(["first", "second"] if batch else "first") + assert collection.inner_storage == {} + + +@mark.parametrize("batch", [False, True]) +@mark.parametrize("serializer", ["model_dump", "model_serializer", "to_dict"]) +async def test_upsert_with_custom_logical_serialization(definition, record_type, batch, serializer): + definition.key_field.storage_name = "content" + definition.fields[1].storage_name = "id" + if serializer == "model_dump": + + class CustomRecord(record_type): + def model_dump(self, **kwargs): + return super().model_dump(**kwargs) + + model_type = CustomRecord + elif serializer == "model_serializer": + + class CustomRecord(record_type): + @model_serializer(mode="wrap") + def serialize_model(self, handler): + return handler(self) + + model_type = CustomRecord + else: + model_type = record_type + definition.to_dict = lambda record, **kwargs: record.model_dump() + + collection = InMemoryCollection(collection_name="test", record_type=model_type, definition=definition) + records = [ + model_type(id="first", content="unrelated first", vector=[1.0, 0.0, 0.0, 0.0, 0.0]), + model_type(id="second", content="unrelated second", vector=[0.0, 1.0, 0.0, 0.0, 0.0]), + ] + + keys = await collection.upsert(records if batch else records[0]) + assert keys == (["first", "second"] if batch else "first") + assert set(collection.inner_storage) == ({"first", "second"} if batch else {"first"}) + + await collection.delete(["first", "second"] if batch else "first") + assert collection.inner_storage == {} + + async def test_get(collection): record = {"id": "testid", "content": "test content", "vector": [0.1, 0.2, 0.3, 0.4, 0.5]} await collection.upsert(record)