Python: fix(python): Fix InMemoryCollection upsert with key storage_name alias

python/semantic_kernel/connectors/in_memory.py

@@ -631,10 +631,16 @@ async def _inner_get(
     @override
     async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence[TKey]:
         updated_keys = []
+        key_storage_name = self._key_field_storage_name
+        key_name = self._key_field_name
         for record in records:
             record = AttributeDict(record)
-            self.inner_storage[record[self._key_field_name]] = record
-            updated_keys.append(record[self._key_field_name])
+            # Serialized records index the key by its storage name (e.g. dict and
+            # dataclass records), while some record types (e.g. Pydantic models)
+            # serialize with the Python field name, so fall back to that.
+            key = record[key_storage_name] if key_storage_name in record else record[key_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]]:

python/tests/unit/connectors/memory/test_in_memory.py

@@ -1,12 +1,14 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import ast
+from dataclasses import dataclass
 
 from pytest import fixture, mark, raises
+from pydantic import BaseModel
 
 from semantic_kernel.connectors.in_memory import InMemoryCollection, InMemoryStore
 from semantic_kernel.data._shared import default_dynamic_filter_function
-from semantic_kernel.data.vector import DistanceFunction
+from semantic_kernel.data.vector import DistanceFunction, VectorStoreCollectionDefinition, VectorStoreField
 from semantic_kernel.exceptions.vector_store_exceptions import VectorStoreOperationException
 
 
@@ -292,3 +294,50 @@ async def test_large_sequence_repeat_filter_is_blocked(collection):
 
     with raises(VectorStoreOperationException, match="Sequence repetition in filter expressions exceeds the maximum"):
         collection._get_filtered_records(type("opt", (), {"filter": "lambda x: [0] * 2000000000"})())
+
+
+@fixture
+def storage_name_definition():
+    return VectorStoreCollectionDefinition(
+        fields=[
+            VectorStoreField("key", name="id", type="str", storage_name="stored_id"),
+            VectorStoreField("vector", name="vector", type="float", dimensions=2),
+        ]
+    )
+
+
+async def test_upsert_with_key_storage_name_dict(storage_name_definition):
+    collection = InMemoryCollection(collection_name="test", record_type=dict, definition=storage_name_definition)
+    key = await collection.upsert({"id": "first", "vector": [1.0, 0.0]})
+    assert key == "first"
+    result = await collection.get("first")
+    assert result["id"] == "first"
+    await collection.delete("first")
+    assert collection.inner_storage == {}
+
+
+async def test_upsert_with_key_storage_name_dataclass(storage_name_definition):
+    @dataclass
+    class DataRecord:
+        id: str
+        vector: list[float]
+
+    collection = InMemoryCollection(collection_name="test", record_type=DataRecord, definition=storage_name_definition)
+    key = await collection.upsert(DataRecord(id="first", vector=[1.0, 0.0]))
+    assert key == "first"
+    result = await collection.get("first", include_vectors=True)
+    assert result.id == "first"
+
+
+async def test_upsert_with_key_storage_name_pydantic(storage_name_definition):
+    class PydanticRecord(BaseModel):
+        id: str
+        vector: list[float]
+
+    collection = InMemoryCollection(
+        collection_name="test", record_type=PydanticRecord, definition=storage_name_definition
+    )
+    key = await collection.upsert(PydanticRecord(id="first", vector=[1.0, 0.0]))
+    assert key == "first"
+    result = await collection.get("first", include_vectors=True)
+    assert result.id == "first"