fix: process nested discriminated unions in tool schemas

[akihikokuroda]

Pull Request

Issue

Fixes #1105

Description

Testing

• Tests added to the respective file if code was changed
• New code has 100% coverage if code was added
• Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

Attribution

• AI coding assistants used

---

Adding a new component, requirement, sampling strategy, or tool?

If your PR adds or modifies one of the types below, check the matching box. A checklist of type-specific review items will be posted as a comment.

• Component
• Requirement
• Sampling Strategy
• Tool

NOTE: Please ensure you have an issue that has been acknowledged by a core contributor and routed you to open a pull request against this repository. Otherwise, please open an issue before continuing with this pull request.

[AngeloDanducci]

The issue mentions

The reproducer above should round-trip a valid payload through
validate_tool_arguments and through a real Ollama tool call without errors.

Looking at the tests I don't think I see any that do a round trip.

Probably worth adding something similar to:

def test_nested_union_payload_round_trips():

    class FullUser(BaseModel):
        type: Literal["full"]
        name: str
        email: str

    class StubUser(BaseModel):
        type: Literal["stub"]
        user_id: str

    class CreateUserFull(BaseModel):
        op: Literal["create_full"]
        user_data: Annotated[FullUser | StubUser, Field(discriminator="type")]

    class DeleteUser(BaseModel):
        op: Literal["delete"]
        user_id: str

    def execute(
        cmd: Annotated[CreateUserFull | DeleteUser, Field(discriminator="op")],
    ) -> str:
        """Execute a command.

        Args:
            cmd: the command to execute
        """
        return "ok"

mt = MelleaTool.from_callable(execute)
    validate_tool_arguments(
        mt,
        {
            "cmd": {
                "op": "create_full",
                "user_data": {
                    "type": "full",
                    "name": "Ada",
                    "email": "ada@example.com",
                },
            }
        },
        strict=True,
    )

[akihikokuroda]

@AngeloDanducci Thanks for review. New tests using "validate_tool_arguments" and code changes fixing issues in validate_tool_arguments are added.

[AngeloDanducci]

After doing some local testing I think acceptance criteria #3 is not 100% met.

The PR's _recursively_flatten_in_properties walks a schema tree and flattens discriminated unions it can see — but it never resolves a plain $ref. So when a union branch has a field like wrapper: Wrapper (rendered as {"$ref": "#/$defs/Wrapper"}), the recursion stops at the $ref boundary. The discriminated union living inside Wrapper is never reached, and the $ref itself survives into the final schema.

[AngeloDanducci]

Suggested change

	# Process additionalProperties
	if "additionalProperties" in schema and isinstance(
	schema["additionalProperties"], dict
	):
	if "oneOf" in schema["additionalProperties"] or (
	"anyOf" in schema["additionalProperties"]
	and any(
	"oneOf" in s for s in schema["additionalProperties"].get("anyOf", [])
	)
	):
	schema["additionalProperties"] = _flatten_discriminated_union(
	schema["additionalProperties"], defs
	)
	# Process additionalProperties
	if "additionalProperties" in schema and isinstance(
	schema["additionalProperties"], dict
	):
	child_seen = _seen_refs
	if "$ref" in schema["additionalProperties"]:
	child_seen = _seen_refs | {schema["additionalProperties"]["$ref"]}
	schema["additionalProperties"] = _inline_ref(schema["additionalProperties"])
	if "oneOf" in schema["additionalProperties"] or (
	"anyOf" in schema["additionalProperties"]
	and any(
	"oneOf" in s for s in schema["additionalProperties"].get("anyOf", [])
	)
	):
	schema["additionalProperties"] = _flatten_discriminated_union(
	schema["additionalProperties"], defs
	)
	_recursively_flatten_in_properties(
	schema["additionalProperties"], defs, child_seen
	)

[AngeloDanducci]

Suggested change

	# Process array items
	if "items" in schema and isinstance(schema["items"], dict):
	if "oneOf" in schema["items"] or (
	"anyOf" in schema["items"]
	and any("oneOf" in s for s in schema["items"].get("anyOf", []))
	):
	schema["items"] = _flatten_discriminated_union(schema["items"], defs)
	_recursively_flatten_in_properties(schema["items"], defs)
	# Process array items
	if "items" in schema and isinstance(schema["items"], dict):
	child_seen = _seen_refs
	if "$ref" in schema["items"]:
	child_seen = _seen_refs | {schema["items"]["$ref"]}
	schema["items"] = _inline_ref(schema["items"])
	if "oneOf" in schema["items"] or (
	"anyOf" in schema["items"]
	and any("oneOf" in s for s in schema["items"].get("anyOf", []))
	):
	schema["items"] = _flatten_discriminated_union(schema["items"], defs)
	_recursively_flatten_in_properties(schema["items"], defs, child_seen)

[AngeloDanducci]

Suggested change

	# Process properties of an object
	if "properties" in schema and isinstance(schema["properties"], dict):
	for prop_name, prop_schema in schema["properties"].items():
	if not isinstance(prop_schema, dict):
	continue
	# Check if this property is a discriminated union (has oneOf + discriminator)
	if "oneOf" in prop_schema or (
	"anyOf" in prop_schema
	and any("oneOf" in s for s in prop_schema.get("anyOf", []))
	):
	schema["properties"][prop_name] = _flatten_discriminated_union(
	prop_schema, defs
	)
	prop_schema = schema["properties"][prop_name]
	# Recurse into the (possibly flattened) property
	# Process properties of an object
	if "properties" in schema and isinstance(schema["properties"], dict):
	for prop_name, prop_schema in schema["properties"].items():
	if not isinstance(prop_schema, dict):
	continue
	# Inline a plain $ref so nested discriminated unions inside the
	# referenced model are reachable (issue #1105 criterion 3: no $ref
	# may survive anywhere in the parameter schema).
	child_seen = _seen_refs
	if "$ref" in prop_schema:
	child_seen = _seen_refs | {prop_schema["$ref"]}
	prop_schema = _inline_ref(prop_schema)
	schema["properties"][prop_name] = prop_schema
	# Check if this property is a discriminated union (has oneOf + discriminator)
	if "oneOf" in prop_schema or (
	"anyOf" in prop_schema
	and any("oneOf" in s for s in prop_schema.get("anyOf", []))
	):
	schema["properties"][prop_name] = _flatten_discriminated_union(
	prop_schema, defs
	)
	prop_schema = schema["properties"][prop_name]
	# Recurse into the (possibly flattened) property
	_recursively_flatten_in_properties(prop_schema, defs, child_seen)

[AngeloDanducci]

Suggested change

	def _recursively_flatten_in_properties(schema: dict, defs: dict) -> None:
	"""Recursively flatten discriminated unions found in nested properties.
	Mutates schema in-place, walking the tree to find and flatten any
	discriminated unions within properties, items, or anyOf branches.
	This is called after branch inlining to ensure nested discriminated unions
	(e.g., a property field that is itself a discriminated union) are also
	flattened.
	Args:
	schema: Schema object to process (mutated in-place).
	defs: Global definitions dict for $ref resolution.
	"""
	if not isinstance(schema, dict):
	return
	def _recursively_flatten_in_properties(
	schema: dict, defs: dict, _seen_refs: frozenset[str] = frozenset()
	) -> None:
	"""Recursively flatten discriminated unions found in nested properties.
	Mutates schema in-place, walking the tree to find and flatten any
	discriminated unions within properties, items, or anyOf branches. Plain
	`$ref` properties are resolved and inlined first so a discriminated union
	reachable only through a `$ref` (a branch field typed as a model that
	itself contains a discriminated union) is still flattened and no `$ref`
	survives.
	This is called after branch inlining to ensure nested discriminated unions
	(e.g., a property field that is itself a discriminated union) are also
	flattened.
	Args:
	schema: Schema object to process (mutated in-place).
	defs: Global definitions dict for $ref resolution.
	_seen_refs: Internal guard tracking $ref paths already inlined on the
	current branch, preventing infinite recursion on self-referential
	models. Callers should not set this.
	"""
	if not isinstance(schema, dict):
	return
	def _inline_ref(node: dict) -> dict:
	"""Resolve a plain `$ref` node to a deep-copied inlined schema.
	Returns the node unchanged if it is not a resolvable `$ref` or if the
	target is already on the current recursion
	"""
	ref_path = node.get("$ref")
	if not ref_path or ref_path in _seen_refs:
	return node
	resolved = _resolve_ref(ref_path, defs)
	return copy.deepcopy(resolved) if resolved

[AngeloDanducci]

Suggested change

	# Test Level1B branch
	validate_tool_arguments(
	mt, {"param": {"l1_type": "b", "name": "direct"}}, strict=True
	)
	def test_nested_union_via_plain_ref_field_flattens_and_validates(self):
	"""A discriminated union reachable only through a plain $ref field must
	be flattened (no surviving $ref) and must still validate strictly."""
	class Leaf1(BaseModel):
	lt: Literal["l1"]
	a: str
	class Leaf2(BaseModel):
	lt: Literal["l2"]
	b: int
	class Wrapper(BaseModel):
	# Wrapper holds the discriminated union...
	leaf: Annotated[Leaf1 | Leaf2, Field(discriminator="lt")]
	class Outer(BaseModel):
	ot: Literal["outer"]
	# ...but Outer references Wrapper as a plain field -> {"$ref": ...}
	wrapper: Wrapper
	def f(cmd: Annotated[Outer, Field(discriminator="ot")]) -> str:
	"""Command.
	Args:
	cmd: the command
	"""
	return "ok"
	tool = convert_function_to_ollama_tool(f)
	assert tool.function is not None
	assert tool.function.parameters is not None
	rendered = json.dumps(
	tool.function.parameters.model_dump(exclude_none=True)["properties"]["cmd"]
	)
	assert "$ref" not in rendered, f"no $ref may survive: {rendered}"
	assert "oneOf" not in rendered
	assert "discriminator" not in rendered
	mt = MelleaTool.from_callable(f)
	validate_tool_arguments(
	mt,
	{"cmd": {"ot": "outer", "wrapper": {"leaf": {"lt": "l1", "a": "x"}}}},
	strict=True,
	)
	with pytest.raises(ValidationError):
	validate_tool_arguments(
	mt,
	{"cmd": {"ot": "outer", "wrapper": {"leaf": {"lt": "GARBAGE"}}}},
	strict=True,
	)

[akihikokuroda]

@AngeloDanducci The "$ref" issue is handled in the follow on PR #1416. It is stacked on this PR. I wanted to separate them to simplify the review process. If it's good to review them together. I can close this PR and make #1416 as ready to review.

[akihikokuroda]

@AngeloDanducci I close this. Please review #1416 instead. Thanks!

[AngeloDanducci]

Ah sounds good, I'll review that one tomorrow morning.