From 9bbb1d2507f41aa1dde2b993fe14992e6875d0ba Mon Sep 17 00:00:00 2001 From: fhgffy <102001626+fhgffy@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:18:52 +0800 Subject: [PATCH 1/4] fix(python): honor in-memory key storage names --- .../semantic_kernel/connectors/in_memory.py | 8 ++++-- .../unit/connectors/memory/test_in_memory.py | 27 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/python/semantic_kernel/connectors/in_memory.py b/python/semantic_kernel/connectors/in_memory.py index f2e1dc30fbe1..19d57cf72973 100644 --- a/python/semantic_kernel/connectors/in_memory.py +++ b/python/semantic_kernel/connectors/in_memory.py @@ -633,8 +633,12 @@ async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence updated_keys = [] for record in records: record = AttributeDict(record) - self.inner_storage[record[self._key_field_name]] = record - updated_keys.append(record[self._key_field_name]) + 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]]: diff --git a/python/tests/unit/connectors/memory/test_in_memory.py b/python/tests/unit/connectors/memory/test_in_memory.py index 7955b7696632..67f8e20d6459 100644 --- a/python/tests/unit/connectors/memory/test_in_memory.py +++ b/python/tests/unit/connectors/memory/test_in_memory.py @@ -35,6 +35,33 @@ 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] + + 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) From 4f8b092b5cf836502996bac3ae89e7add3b384c6 Mon Sep 17 00:00:00 2001 From: fhgffy <102001626+fhgffy@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:48:08 +0800 Subject: [PATCH 2/4] Preserve logical keys when storage names collide --- python/semantic_kernel/connectors/in_memory.py | 4 +--- .../unit/connectors/memory/test_in_memory.py | 15 +++++++++++++++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/python/semantic_kernel/connectors/in_memory.py b/python/semantic_kernel/connectors/in_memory.py index 19d57cf72973..4ea031631778 100644 --- a/python/semantic_kernel/connectors/in_memory.py +++ b/python/semantic_kernel/connectors/in_memory.py @@ -633,9 +633,7 @@ async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence updated_keys = [] for record in records: record = AttributeDict(record) - key_field_name = ( - self._key_field_storage_name if self._key_field_storage_name in record else self._key_field_name - ) + key_field_name = self._key_field_name if self._key_field_name in record else self._key_field_storage_name key = record[key_field_name] self.inner_storage[key] = record updated_keys.append(key) diff --git a/python/tests/unit/connectors/memory/test_in_memory.py b/python/tests/unit/connectors/memory/test_in_memory.py index 67f8e20d6459..790621ba40a8 100644 --- a/python/tests/unit/connectors/memory/test_in_memory.py +++ b/python/tests/unit/connectors/memory/test_in_memory.py @@ -62,6 +62,21 @@ async def test_upsert_with_key_storage_name(definition, dataclass_vector_data_mo 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 == {} + + 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) From d228c860869c4699a3fc9fadd5513e6d7ba93c3e Mon Sep 17 00:00:00 2001 From: fhgffy <102001626+fhgffy@users.noreply.github.com> Date: Mon, 5 Oct 2026 21:01:50 +0800 Subject: [PATCH 3/4] fix(python): track serialized key names for in-memory indexing --- .../semantic_kernel/connectors/in_memory.py | 37 ++++- .../unit/connectors/memory/test_in_memory.py | 127 ++++++++++++++++++ 2 files changed, 162 insertions(+), 2 deletions(-) diff --git a/python/semantic_kernel/connectors/in_memory.py b/python/semantic_kernel/connectors/in_memory.py index 4ea031631778..42dfb09d0cc8 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 # 2026-10-05 from scipy.spatial.distance import cityblock, cosine, euclidean, hamming, sqeuclidean from typing_extensions import override @@ -81,6 +81,16 @@ def __delattr__(self, name) -> None: raise AttributeError(name) +# 2026-10-05: 携带默认 Pydantic 输出的实际键名, 字段和值保持原样. +class _InMemoryRecord(dict[str, Any]): + """携带已知键字段名称的序列化记录.""" + + def __init__(self, record: dict[str, Any], key_field_name: str): + """保存序列化字段和对应的键名称.""" + 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,8 +642,15 @@ async def _inner_get( async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence[TKey]: updated_keys = [] for record in records: + # 2026-10-05: 按实际序列化格式选择键, 避免两种方向的字段名碰撞. + key_field_name = ( + record.key_field_name if isinstance(record, _InMemoryRecord) else self._key_field_storage_name + ) record = AttributeDict(record) - key_field_name = self._key_field_name if self._key_field_name in record else self._key_field_storage_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) @@ -643,6 +660,22 @@ def _deserialize_store_models_to_dicts(self, records: Sequence[Any], **kwargs: A return records def _serialize_dicts_to_store_models(self, records: Sequence[dict[str, Any]], **kwargs: Any) -> Sequence[Any]: + # 2026-10-05: 此边界仅处理字典序列化结果, 成功的自定义 serialize 会绕过它. + if ( + issubclass(self.record_type, BaseModel) + and not self.definition.to_dict + and self.record_type.model_dump is BaseModel.model_dump + and not self.record_type.__pydantic_decorators__.model_serializers + ): + key_field_name = self._key_field_name + # 2026-10-05: 仅在当前 Pydantic 支持且启用别名输出时采用已解析的序列化别名. + if ( + "serialize_by_alias" in ConfigDict.__annotations__ + and self.record_type.model_config.get("serialize_by_alias") + and (key_field := self.record_type.model_fields.get(key_field_name)) + ): + key_field_name = key_field.serialization_alias or key_field_name + return [_InMemoryRecord(record, key_field_name) for record in records] return records @override diff --git a/python/tests/unit/connectors/memory/test_in_memory.py b/python/tests/unit/connectors/memory/test_in_memory.py index 790621ba40a8..8832d34f5bdc 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 # 2026-10-05 from pytest import fixture, mark, raises from semantic_kernel.connectors.in_memory import InMemoryCollection, InMemoryStore @@ -77,6 +78,132 @@ async def test_upsert_when_storage_name_matches_logical_field(definition, record assert collection.inner_storage == {} +# 2026-10-05: 覆盖 Pydantic 逻辑字段与键存储名相撞时的单条和批量索引. +@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 == {} + + +# 2026-10-05: 覆盖其他字段的存储名与逻辑键相撞时的字典和 dataclass 序列化. +@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 == {} + + +# 2026-10-05: 自定义 Pydantic 序列化使用存储名时保留其实际键和值. +@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" + definition.fields[1].storage_name = "id" + if serializer == "model_serialize": + + class CustomRecord(record_type): + def serialize(self, **kwargs): + return {"stored_id": self.id, "id": 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, "id": 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, "id": self.content, "vector": self.vector} + + model_type = CustomSerializedRecord + else: + model_type = record_type + setattr( + definition, + serializer, + lambda record, **kwargs: {"stored_id": record.id, "id": 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 == {} + + +# 2026-10-05: 默认按 Pydantic 序列化别名输出时使用实际键名, 避免反向别名碰撞. +@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 == {} + + 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) From 2c5fb848b8ae6b2d1adf6cec94aaf34faa575441 Mon Sep 17 00:00:00 2001 From: fhgffy <102001626+fhgffy@users.noreply.github.com> Date: Mon, 5 Oct 2026 22:23:15 +0800 Subject: [PATCH 4/4] Preserve key provenance through custom record serialization --- .../semantic_kernel/connectors/in_memory.py | 43 ++++++++------- .../unit/connectors/memory/test_in_memory.py | 55 +++++++++++++++---- 2 files changed, 67 insertions(+), 31 deletions(-) diff --git a/python/semantic_kernel/connectors/in_memory.py b/python/semantic_kernel/connectors/in_memory.py index 42dfb09d0cc8..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 BaseModel, ConfigDict, Field # 2026-10-05 +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,12 +81,11 @@ def __delattr__(self, name) -> None: raise AttributeError(name) -# 2026-10-05: 携带默认 Pydantic 输出的实际键名, 字段和值保持原样. 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 @@ -642,7 +641,6 @@ async def _inner_get( async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence[TKey]: updated_keys = [] for record in records: - # 2026-10-05: 按实际序列化格式选择键, 避免两种方向的字段名碰撞. key_field_name = ( record.key_field_name if isinstance(record, _InMemoryRecord) else self._key_field_storage_name ) @@ -659,23 +657,26 @@ async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence def _deserialize_store_models_to_dicts(self, records: Sequence[Any], **kwargs: Any) -> Sequence[dict[str, Any]]: return records - def _serialize_dicts_to_store_models(self, records: Sequence[dict[str, Any]], **kwargs: Any) -> Sequence[Any]: - # 2026-10-05: 此边界仅处理字典序列化结果, 成功的自定义 serialize 会绕过它. - if ( - issubclass(self.record_type, BaseModel) - and not self.definition.to_dict - and self.record_type.model_dump is BaseModel.model_dump - and not self.record_type.__pydantic_decorators__.model_serializers - ): - key_field_name = self._key_field_name - # 2026-10-05: 仅在当前 Pydantic 支持且启用别名输出时采用已解析的序列化别名. - if ( - "serialize_by_alias" in ConfigDict.__annotations__ - and self.record_type.model_config.get("serialize_by_alias") - and (key_field := self.record_type.model_fields.get(key_field_name)) + 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 - return [_InMemoryRecord(record, key_field_name) for record in records] + 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 @override diff --git a/python/tests/unit/connectors/memory/test_in_memory.py b/python/tests/unit/connectors/memory/test_in_memory.py index 8832d34f5bdc..5447f267a68f 100644 --- a/python/tests/unit/connectors/memory/test_in_memory.py +++ b/python/tests/unit/connectors/memory/test_in_memory.py @@ -2,7 +2,7 @@ import ast -from pydantic import ConfigDict, Field, model_serializer # 2026-10-05 +from pydantic import ConfigDict, Field, model_serializer from pytest import fixture, mark, raises from semantic_kernel.connectors.in_memory import InMemoryCollection, InMemoryStore @@ -78,7 +78,6 @@ async def test_upsert_when_storage_name_matches_logical_field(definition, record assert collection.inner_storage == {} -# 2026-10-05: 覆盖 Pydantic 逻辑字段与键存储名相撞时的单条和批量索引. @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): @@ -107,7 +106,6 @@ def serialize(self, **kwargs): assert collection.inner_storage == {} -# 2026-10-05: 覆盖其他字段的存储名与逻辑键相撞时的字典和 dataclass 序列化. @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): @@ -130,24 +128,24 @@ async def test_upsert_when_another_storage_name_matches_key(definition, dataclas assert collection.inner_storage == {} -# 2026-10-05: 自定义 Pydantic 序列化使用存储名时保留其实际键和值. @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" - definition.fields[1].storage_name = "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, "id": self.content, "vector": self.vector} + 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, "id": self.content, "vector": self.vector} + return {"stored_id": self.id, content_name: self.content, "vector": self.vector} model_type = CustomDumpRecord elif serializer == "model_serializer": @@ -155,7 +153,7 @@ def model_dump(self, **kwargs): class CustomSerializedRecord(record_type): @model_serializer def serialize_model(self): - return {"stored_id": self.id, "id": self.content, "vector": self.vector} + return {"stored_id": self.id, content_name: self.content, "vector": self.vector} model_type = CustomSerializedRecord else: @@ -163,7 +161,7 @@ def serialize_model(self): setattr( definition, serializer, - lambda record, **kwargs: {"stored_id": record.id, "id": record.content, "vector": record.vector}, + 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 = [ @@ -179,7 +177,6 @@ def serialize_model(self): assert collection.inner_storage == {} -# 2026-10-05: 默认按 Pydantic 序列化别名输出时使用实际键名, 避免反向别名碰撞. @mark.parametrize("batch", [False, True]) async def test_upsert_with_pydantic_serialization_aliases(definition, record_type, batch): definition.key_field.storage_name = "stored_id" @@ -204,6 +201,44 @@ class AliasedRecord(record_type): 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)