Python: Fix in-memory upserts with key storage names

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
 

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)