Python: accept nullable auto-generated keys in vector upsert tools

python/packages/core/agent_framework/_vectors.py

@@ -2060,13 +2060,35 @@ 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: (
             dict(collection.key_json_schema()) if field.field_type == "key" else _vector_tool_field_schema(field)
         )
         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))

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:

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)