From 3dcbdaad7bde0f1077e485fd7e149ddb7dc32034 Mon Sep 17 00:00:00 2001 From: Jessica He Date: Fri, 14 Aug 2026 15:36:03 -0400 Subject: [PATCH 1/2] fix constraints on tool fields --- mellea/backends/tools.py | 24 +++++++++- .../test_discriminated_union_tools.py | 10 ++++ .../backends/test_pydantic_tool_parameters.py | 46 +++++++++++++++++++ 3 files changed, 79 insertions(+), 1 deletion(-) diff --git a/mellea/backends/tools.py b/mellea/backends/tools.py index 12b02493c..29efa0ea2 100644 --- a/mellea/backends/tools.py +++ b/mellea/backends/tools.py @@ -688,6 +688,15 @@ def _build_pydantic_type_from_schema(schema: dict[str, Any]) -> Any: return (result | None) if has_null else result + # Handle enum constraint (e.g. Literal["read", "write"] on a plain field) + if "enum" in schema: + enum_vals = tuple(schema["enum"]) + return Literal[enum_vals] # type: ignore + + # Handle const constraint (e.g. a bare {"type": "string", "const": "cat"}) + if "const" in schema: + return Literal[schema["const"]] # type: ignore + # Simple type mapping return JSON_TYPE_TO_PYTHON.get(json_type, Any) @@ -1413,18 +1422,31 @@ def convert_function_to_ollama_tool( # This now handles Optional[primitive] types correctly if "anyOf" in v: types = {t.get("type", "string") for t in v.get("anyOf")} + # Collect enum values from each non-null anyOf branch + # (handles Optional[Literal[...]] which Pydantic emits as anyOf) + enum_values: list | None = None + for sub in v.get("anyOf", []): + if sub.get("type") == "null": + continue + if "enum" in sub: + enum_values = list(sub["enum"]) + break else: types = {v.get("type", "string")} + enum_values = list(v["enum"]) if "enum" in v else None if "null" in types: if k in schema.get("required", []): schema["required"].remove(k) types.discard("null") - schema["properties"][k] = { + prop: dict = { "description": parsed_docstring.get(k, ""), "type": ", ".join(types), } + if enum_values is not None: + prop["enum"] = enum_values + schema["properties"][k] = prop # Final pass: recursively inline all remaining $refs at any depth. # This catches dangling references in nested model properties that weren't diff --git a/test/backends/test_discriminated_union_tools.py b/test/backends/test_discriminated_union_tools.py index 88fd28abe..3c30660a0 100644 --- a/test/backends/test_discriminated_union_tools.py +++ b/test/backends/test_discriminated_union_tools.py @@ -263,6 +263,16 @@ def test_strict_rejects_missing_discriminator(self): with pytest.raises(ValidationError): validate_tool_arguments(mt, {"pet": {"name": "Rex"}}, strict=True) + def test_strict_rejects_invalid_discriminator_value(self): + """A kind value outside the allowed set must be rejected.""" + mt = MelleaTool.from_callable(act) + with pytest.raises(ValidationError): + validate_tool_arguments( + mt, + {"pet": {"kind": "horse", "name": "Bob", "breed": "lab"}}, + strict=True, + ) + def test_optional_accepts_omitted(self): """The optional variant accepts the parameter being omitted.""" mt = MelleaTool.from_callable(act_optional) diff --git a/test/backends/test_pydantic_tool_parameters.py b/test/backends/test_pydantic_tool_parameters.py index 6b9191fe9..17b61f45b 100644 --- a/test/backends/test_pydantic_tool_parameters.py +++ b/test/backends/test_pydantic_tool_parameters.py @@ -1009,5 +1009,51 @@ def create_person(person: Person) -> str: assert validated_no_addr["person"]["address"] is None +class TestLiteralFields: + """Schema generation and validation for plain Literal[...] tool parameters.""" + + def test_strict_rejects_value_outside_literal(self): + """validate_tool_arguments must reject values not in Literal[...] when strict=True.""" + from typing import Literal + + from pydantic import ValidationError + + def file_op(path: str, mode: Literal["read", "write"]) -> str: + """Perform a file operation. + + Args: + path: the file path + mode: the operation mode + """ + return "ok" + + mt = MelleaTool.from_callable(file_op) + with pytest.raises(ValidationError): + validate_tool_arguments( + mt, + {"path": "/etc/passwd", "mode": "delete"}, + strict=True, + ) + + def test_strict_accepts_value_inside_literal(self): + """validate_tool_arguments must accept values that are in Literal[...].""" + from typing import Literal + + def file_op(path: str, mode: Literal["read", "write"]) -> str: + """Perform a file operation. + + Args: + path: the file path + mode: the operation mode + """ + return "ok" + + mt = MelleaTool.from_callable(file_op) + result = validate_tool_arguments( + mt, {"path": "/tmp/file.txt", "mode": "read"}, strict=True + ) + assert result["mode"] == "read" + + if __name__ == "__main__": pytest.main([__file__, "-v"]) From 500d382f4614a6d8c4529b8ee289af6fda135022 Mon Sep 17 00:00:00 2001 From: Jessica He Date: Mon, 24 Aug 2026 16:19:16 -0400 Subject: [PATCH 2/2] ruff format --- test/backends/test_pydantic_tool_parameters.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/test/backends/test_pydantic_tool_parameters.py b/test/backends/test_pydantic_tool_parameters.py index 17b61f45b..fe0826edd 100644 --- a/test/backends/test_pydantic_tool_parameters.py +++ b/test/backends/test_pydantic_tool_parameters.py @@ -1030,9 +1030,7 @@ def file_op(path: str, mode: Literal["read", "write"]) -> str: mt = MelleaTool.from_callable(file_op) with pytest.raises(ValidationError): validate_tool_arguments( - mt, - {"path": "/etc/passwd", "mode": "delete"}, - strict=True, + mt, {"path": "/etc/passwd", "mode": "delete"}, strict=True ) def test_strict_accepts_value_inside_literal(self):