diff --git a/python/semantic_kernel/connectors/in_memory.py b/python/semantic_kernel/connectors/in_memory.py index f2e1dc30fbe1..7f6a6daba04b 100644 --- a/python/semantic_kernel/connectors/in_memory.py +++ b/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]]: diff --git a/python/tests/unit/connectors/memory/test_in_memory.py b/python/tests/unit/connectors/memory/test_in_memory.py index 7955b7696632..b3239d1df60e 100644 --- a/python/tests/unit/connectors/memory/test_in_memory.py +++ b/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"