diff --git a/python/packages/core/agent_framework/_vectors.py b/python/packages/core/agent_framework/_vectors.py index 1993d4dc75..142a709b78 100644 --- a/python/packages/core/agent_framework/_vectors.py +++ b/python/packages/core/agent_framework/_vectors.py @@ -572,16 +572,33 @@ def encode(value: Any) -> Mapping[str, Any]: return encode +def _pydantic_validation_key(name: str, field_info: Any) -> str: + """Return the key that feeds one model field from a mapping keyed by field name. + + A string validation alias (which ``alias=`` also installs) or a string model alias is the key + Pydantic accepts, so it is used unchanged. ``AliasChoices`` declares several accepted aliases and + validates on any one of them, so the first string choice is used; keying such a field by its own + name only validates when the model also allows population by name, which declaring an alias does + not require. Nested ``AliasPath`` aliases keep the previous field-name fallback. + """ + validation_alias = getattr(field_info, "validation_alias", None) + if isinstance(validation_alias, str): + return validation_alias + choices = getattr(validation_alias, "choices", None) + if isinstance(choices, Sequence) and not isinstance(choices, str): + for choice in cast("Sequence[Any]", choices): + if isinstance(choice, str): + return choice + alias = getattr(field_info, "alias", None) + return alias if isinstance(alias, str) else name + + def _default_vector_model_decoder(record_type: type[Any]) -> VectorModelDecoder: if issubclass(record_type, BaseModel): def decode_pydantic(value: Mapping[str, Any]) -> Any: validation_value = { - field.validation_alias - if isinstance(field.validation_alias, str) - else field.alias - if isinstance(field.alias, str) - else name: value[name] + _pydantic_validation_key(name, field): value[name] for name, field in record_type.model_fields.items() if name in value } diff --git a/python/packages/core/tests/core/test_vectors.py b/python/packages/core/tests/core/test_vectors.py index 34d0cd8793..a79ca1cf01 100644 --- a/python/packages/core/tests/core/test_vectors.py +++ b/python/packages/core/tests/core/test_vectors.py @@ -15,7 +15,7 @@ import msgspec import pytest -from pydantic import BaseModel +from pydantic import AliasChoices, BaseModel from pydantic import Field as PydanticField from agent_framework import ( @@ -721,6 +721,38 @@ class AliasedRecord(BaseModel): assert restored.id == "one" +async def test_pydantic_alias_choices_records_round_trip_through_a_collection() -> None: + # The encoder stores field names, so the decoder feeds the first declared alias choice back to + # Pydantic. Falling back to the field name only validates when the model also populates by name, + # which an alias-only field does not, so the record read back failed after a successful upsert. + @vectorstoremodel(collection_name="alias_choices_records") + class AliasChoiceRecord(BaseModel): + id: Annotated[ + str, + PydanticField(validation_alias=AliasChoices("record_id", "document_id")), + VectorStoreField("key"), + ] + text: Annotated[str, VectorStoreField("data")] + embedding: Annotated[list[float], VectorStoreField("vector", dimensions=2)] = PydanticField( + default_factory=list + ) + + collection: InMemoryCollection[str, AliasChoiceRecord] = InMemoryCollection(AliasChoiceRecord) + await collection.ensure_collection_exists() + + keys = await collection.upsert( + [AliasChoiceRecord(record_id="one", text="plain sample", embedding=[0.25, 0.75])], + generate_vectors=False, + ) + restored = await collection.get(keys, include_vectors=True) + + assert keys == ["one"] + assert len(restored) == 1 + assert restored[0].id == "one" + assert restored[0].text == "plain sample" + assert restored[0].embedding == [0.25, 0.75] + + async def test_collection_serializes_records_and_generates_vectors() -> None: embedding_client = MockEmbeddingClient() collection = MockCollection(embedding_generator=embedding_client)