diff --git a/python/packages/core/agent_framework/_vectors.py b/python/packages/core/agent_framework/_vectors.py index 54273b5d3b2..c758ecdd31b 100644 --- a/python/packages/core/agent_framework/_vectors.py +++ b/python/packages/core/agent_framework/_vectors.py @@ -2060,6 +2060,25 @@ def _vector_tool_field_schema(field: VectorStoreField) -> dict[str, Any]: return {"type": json_type} if json_type is not None else {} +def _vector_tool_has_nullable_auto_key(collection: BaseVectorCollection[Any, Any]) -> bool: + key_field = collection.definition.key_field + if not key_field.is_auto_generated: + return False + if collection.record_type is dict: + return True + try: + annotations = get_type_hints(collection.record_type, include_extras=True) + if not any( + any(isinstance(metadata, VectorStoreField) for metadata in get_args(annotation)[1:]) + for annotation in annotations.values() + if get_origin(annotation) is Annotated + ): + annotations = get_type_hints(cast(Any, collection.record_type).__init__, include_extras=True) + except (NameError, TypeError): + return False + return type(None) in get_args(_unwrap_annotation(annotations.get(key_field.name))) + + def _vector_tool_record_schema(collection: BaseVectorCollection[Any, Any]) -> dict[str, Any]: properties = { field.name: ( @@ -2067,6 +2086,9 @@ def _vector_tool_record_schema(collection: BaseVectorCollection[Any, Any]) -> di ) for field in collection.definition.fields } + if _vector_tool_has_nullable_auto_key(collection): + key_name = collection.definition.key_name + properties[key_name] = {"anyOf": [properties[key_name], {"type": "null"}]} required = [ field.name for field in collection.definition.fields @@ -2239,12 +2261,13 @@ def _decode_vector_tool_records( ) decoded: list[ModelT] = [] registration = _VECTOR_MODEL_REGISTRY.get(collection.record_type) + nullable_auto_key = _vector_tool_has_nullable_auto_key(collection) for record in raw_records: if not isinstance(record, Mapping): raise TypeError("Each record must be a mapping.") logical_record = dict(cast(Mapping[str, Any], record)) key_name = collection.definition.key_name - if key_name in logical_record: + if key_name in logical_record and not (logical_record[key_name] is None and nullable_auto_key): logical_record[key_name] = collection.key_from_json(logical_record[key_name]) if collection.record_type is dict: decoded.append(cast(ModelT, logical_record)) diff --git a/python/packages/core/tests/core/test_vectors.py b/python/packages/core/tests/core/test_vectors.py index e5e05331a56..0be0ceea6a8 100644 --- a/python/packages/core/tests/core/test_vectors.py +++ b/python/packages/core/tests/core/test_vectors.py @@ -2767,7 +2767,10 @@ async def test_upsert_tool_forwards_provider_embedding_options() -> None: assert embedding_client.options == {"task_type": "RETRIEVAL_DOCUMENT", "dimensions": 2} -async def test_vector_crud_tools_support_auto_generated_keys_when_model_can_omit_them() -> None: +@pytest.mark.parametrize("key_values", [{}, {"id": None}, {"id": "provided"}]) +async def test_vector_crud_tools_support_auto_generated_keys_when_model_can_omit_them( + key_values: dict[str, Any], +) -> None: definition = VectorStoreCollectionDefinition( [ VectorStoreField("key", name="id", type_="str", is_auto_generated=True), @@ -2780,12 +2783,16 @@ async def test_vector_crud_tools_support_auto_generated_keys_when_model_can_omit upsert_tool = create_upsert_tool(collection, generate_vectors=False) record_schema = upsert_tool.parameters()["properties"]["records"]["items"] - assert record_schema["properties"]["id"] == {"type": "string"} + content = await upsert_tool.invoke(arguments={"records": [{**key_values, "text": "generated"}]}) + assert content[0].text is not None + result = msgspec.json.decode(content[0].text) + assert record_schema["properties"]["id"] == {"anyOf": [{"type": "string"}, {"type": "null"}]} assert record_schema["required"] == ["text"] - - result = await upsert_tool.invoke(arguments={"records": [{"text": "generated"}]}, skip_parsing=True) assert len(result["keys"]) == 1 assert isinstance(result["keys"][0], str) + assert await collection.get(result["keys"]) == [{"id": result["keys"][0], "text": "generated"}] + if key_values.get("id") is not None: + assert result["keys"] == [key_values["id"]] def test_vector_upsert_tool_requires_auto_key_when_typed_model_has_no_default() -> None: @@ -2800,9 +2807,11 @@ class RequiredAutoKey: record_schema = create_upsert_tool(collection).parameters()["properties"]["records"]["items"] assert record_schema["required"] == ["id", "text"] + assert record_schema["properties"]["id"] == {"type": "string"} -async def test_vector_upsert_tool_omits_auto_key_when_typed_model_has_default() -> None: +@pytest.mark.parametrize("key_values", [{}, {"id": None}, {"id": "provided"}]) +async def test_vector_upsert_tool_omits_auto_key_when_typed_model_has_default(key_values: dict[str, Any]) -> None: @vectorstoremodel(collection_name="default_auto_key") @dataclass class DefaultAutoKey: @@ -2814,10 +2823,104 @@ class DefaultAutoKey: tool = create_upsert_tool(collection, generate_vectors=False) record_schema = tool.parameters()["properties"]["records"]["items"] - result = await tool.invoke(arguments={"records": [{"text": "generated"}]}, skip_parsing=True) + content = await tool.invoke(arguments={"records": [{**key_values, "text": "generated"}]}) + assert content[0].text is not None + result = msgspec.json.decode(content[0].text) assert record_schema["required"] == ["text"] + assert record_schema["properties"]["id"] == {"anyOf": [{"type": "string"}, {"type": "null"}]} assert len(result["keys"]) == 1 + assert await collection.get(result["keys"]) == [DefaultAutoKey(text="generated", id=result["keys"][0])] + if key_values.get("id") is not None: + assert result["keys"] == [key_values["id"]] + + +@pytest.mark.parametrize("auto_generated", [False, True]) +async def test_vector_upsert_tool_rejects_invalid_non_null_keys(auto_generated: bool) -> None: + definition = VectorStoreCollectionDefinition( + [ + VectorStoreField("key", name="id", type_="str", is_auto_generated=auto_generated), + VectorStoreField("data", name="text", type_="str"), + ], + collection_name="invalid_keys", + ) + collection: InMemoryCollection[str, dict[str, Any]] = InMemoryCollection(dict, definition=definition) + await collection.ensure_collection_exists() + tool = create_upsert_tool(collection, generate_vectors=False) + + with pytest.raises(TypeError, match="Key value must match"): + await tool.invoke(arguments={"records": [{"id": 1, "text": "invalid"}]}) + if not auto_generated: + with pytest.raises(TypeError, match="Key value must match"): + await tool.invoke(arguments={"records": [{"id": None, "text": "invalid"}]}) + assert tool.parameters()["properties"]["records"]["items"]["properties"]["id"] == {"type": "string"} + assert await collection.get() == [] + + +async def test_vector_upsert_tool_preserves_non_nullable_model_key_validation() -> None: + @vectorstoremodel(collection_name="non_nullable_auto_key") + @dataclass + class NonNullableAutoKey: + text: Annotated[str, VectorStoreField("data")] + id: Annotated[str, VectorStoreField("key", is_auto_generated=True)] = "provided" + + collection: InMemoryCollection[str, NonNullableAutoKey] = InMemoryCollection(NonNullableAutoKey) + await collection.ensure_collection_exists() + tool = create_upsert_tool(collection, generate_vectors=False) + + assert tool.parameters()["properties"]["records"]["items"]["properties"]["id"] == {"type": "string"} + with pytest.raises((TypeError, msgspec.ValidationError)): + await tool.invoke(arguments={"records": [{"id": None, "text": "invalid"}]}) + assert await collection.get() == [] + + +@pytest.mark.parametrize("nullable", [False, True]) +@pytest.mark.parametrize("class_annotations", [False, True]) +async def test_vector_upsert_tool_respects_constructor_key_nullability(nullable: bool, class_annotations: bool) -> None: + @vectorstoremodel(collection_name="constructor_auto_key", encoder=vars) + class PlainRecord: + __annotations__ = {"id": str | None} if class_annotations else {} + + def __init__( + self, + text: Annotated[str, VectorStoreField("data")], + id: Annotated[str, VectorStoreField("key", is_auto_generated=True)] = "provided", + ) -> None: + self.text = text + self.id = id + + @vectorstoremodel(collection_name="nullable_constructor_auto_key", encoder=vars) + class NullablePlainRecord: + __annotations__ = {"id": str} if class_annotations else {} + + def __init__( + self, + text: Annotated[str, VectorStoreField("data")], + id: Annotated[str | None, VectorStoreField("key", is_auto_generated=True)] = None, + ) -> None: + self.text = text + self.id = id + + record_type = NullablePlainRecord if nullable else PlainRecord + collection = InMemoryCollection(record_type) + await collection.ensure_collection_exists() + tool = create_upsert_tool(collection, generate_vectors=False) + + key_schema = tool.parameters()["properties"]["records"]["items"]["properties"]["id"] + if nullable: + assert key_schema == {"anyOf": [{"type": "string"}, {"type": "null"}]} + content = await tool.invoke(arguments={"records": [{"id": None, "text": "generated"}]}) + assert content[0].text is not None + key = msgspec.json.decode(content[0].text)["keys"][0] + assert isinstance(key, str) + record = (await collection.get([key]))[0] + assert isinstance(record, NullablePlainRecord) + assert record.text == "generated" + else: + assert key_schema == {"type": "string"} + with pytest.raises(TypeError, match="Key value must match"): + await tool.invoke(arguments={"records": [{"id": None, "text": "invalid"}]}) + assert await collection.get() == [] async def test_vector_crud_tools_round_trip_uuid_keys_as_json_strings() -> None: diff --git a/python/packages/mongodb/tests/mongodb/test_vector_store.py b/python/packages/mongodb/tests/mongodb/test_vector_store.py index 331d662b9b0..1f29917ec6d 100644 --- a/python/packages/mongodb/tests/mongodb/test_vector_store.py +++ b/python/packages/mongodb/tests/mongodb/test_vector_store.py @@ -378,7 +378,8 @@ async def test_generated_object_id_dict_crud_and_projection(mongo_mocks, cursor_ native_collection.delete_many.assert_awaited_once_with({"_id": {"$in": keys}}) -async def test_object_id_keys_round_trip_through_crud_tools(mongo_mocks, cursor_factory): +@pytest.mark.parametrize("key_values", [{}, {"id": None}]) +async def test_object_id_keys_round_trip_through_crud_tools(mongo_mocks, cursor_factory, key_values): client, _, native_collection = mongo_mocks definition = VectorStoreCollectionDefinition([ VectorStoreField("key", name="id", type_="ObjectId", is_auto_generated=True), @@ -396,10 +397,12 @@ async def test_object_id_keys_round_trip_through_crud_tools(mongo_mocks, cursor_ delete_tool = create_delete_tool(collection) schema = upsert_tool.parameters()["properties"]["records"]["items"] - assert schema["properties"]["id"] == {"type": "string", "pattern": "^[0-9a-fA-F]{24}$"} + assert schema["properties"]["id"] == { + "anyOf": [{"type": "string", "pattern": "^[0-9a-fA-F]{24}$"}, {"type": "null"}] + } assert schema["required"] == ["text"] - upserted = await upsert_tool.invoke(arguments={"records": [{"text": "hello"}]}, skip_parsing=True) + upserted = await upsert_tool.invoke(arguments={"records": [{**key_values, "text": "hello"}]}, skip_parsing=True) key = upserted["keys"][0] assert isinstance(key, str) and ObjectId.is_valid(key)