diff --git a/python/semantic_kernel/connectors/in_memory.py b/python/semantic_kernel/connectors/in_memory.py index a1f23ac200dd..2b899d82c84c 100644 --- a/python/semantic_kernel/connectors/in_memory.py +++ b/python/semantic_kernel/connectors/in_memory.py @@ -655,7 +655,31 @@ async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence return updated_keys def _deserialize_store_models_to_dicts(self, records: Sequence[Any], **kwargs: Any) -> Sequence[dict[str, Any]]: - return records + if not (isinstance(self.record_type, type) and issubclass(self.record_type, BaseModel)): + return records + # Pydantic records are stored under their Python field names (they are + # serialized with model_dump()), but the shared deserializer for + # Pydantic models expects store dicts keyed by storage name and maps + # them back to field names. Rename name -> storage name in one pass + # so a data field whose name overlaps the key's storage name (or vice + # versa) keeps its own value instead of overwriting the key's. + converted_records: list[dict[str, Any]] = [] + for record in records: + renames = { + field.storage_name: record[field.name] + for field in self.definition.fields + if field.storage_name and field.storage_name != field.name and field.name in record + } + if not renames: + converted_records.append(record) + continue + # Drop every renamed field's original name, then apply all + # renames at once so overlapping names cannot overwrite each other. + renamed_names = {field.name for field in self.definition.fields if field.storage_name in renames} + converted = {key: value for key, value in record.items() if key not in renamed_names} + converted.update(renames) + converted_records.append(converted) + return converted_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.""" diff --git a/python/tests/unit/connectors/memory/test_in_memory.py b/python/tests/unit/connectors/memory/test_in_memory.py index bdd1d4b90b7f..779f62a6b5a6 100644 --- a/python/tests/unit/connectors/memory/test_in_memory.py +++ b/python/tests/unit/connectors/memory/test_in_memory.py @@ -2,12 +2,12 @@ import ast -from pydantic import ConfigDict, Field, model_serializer +from pydantic import BaseModel, ConfigDict, Field, model_serializer from pytest import fixture, mark, raises 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 @@ -518,3 +518,28 @@ 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"})()) + + +async def test_upsert_pydantic_key_storage_name_overlapping_data_field(): + # The key field's storage name ("stored_id") is also the name of a data + # field. Pydantic records serialize with logical names, so the key must + # come from "id", not from the data field named "stored_id". + definition = VectorStoreCollectionDefinition( + fields=[ + VectorStoreField("key", name="id", type="str", storage_name="stored_id"), + VectorStoreField("data", name="stored_id", type="str", storage_name="payload"), + VectorStoreField("vector", name="vector", type="float", dimensions=2), + ] + ) + + class OverlappingRecord(BaseModel): + id: str + stored_id: str + vector: list[float] + + collection = InMemoryCollection(collection_name="test", record_type=OverlappingRecord, definition=definition) + key = await collection.upsert(OverlappingRecord(id="actual-key", stored_id="payload-value", vector=[1.0, 0.0])) + assert key == "actual-key" + result = await collection.get("actual-key", include_vectors=True) + assert result.id == "actual-key" + assert result.stored_id == "payload-value"