fix: preserve discriminated-union schema in tool conversion

[planetf1]

Bug Fix: discriminated-union tool parameters

Type of PR

• Bug Fix
• New Feature
• Documentation
• Other

Description

• Link to Issue: Fix: Discriminated-union tool parameters lose their schema and fail validation #989

A tool parameter typed as a Pydantic discriminated union — Annotated[A | B, Field(discriminator="kind")], with or without | None — currently collapses to {"type": "string"} in the schema emitted by convert_function_to_ollama_tool.

Despite the function name, that schema is consumed by every backend (as_json_tool is read by Ollama, OpenAI, LiteLLM, Watsonx, and HuggingFace), so discriminated-union tool parameters are silently broken across the board: the model sees a string and hallucinates a payload, and validate_tool_arguments rejects valid dicts.

Root cause

Pydantic emits the union as oneOf plus an OAS-3 discriminator keyword:

// Required parameter:
{"discriminator": {...}, "oneOf": [{"$ref": "#/$defs/Cat"}, {"$ref": "#/$defs/Dog"}]}

// Optional parameter:
{"anyOf": [{"discriminator": {...}, "oneOf": [...]}, {"type": "null"}]}

Neither oneOf nor discriminator is in the JSON Schema subset accepted by tool-calling APIs. The existing inliner only descends into anyOf and $ref, so the structure falls through to the primitive-flattening branch and emerges as {"type": "string"}.

Fix

Adds a pre-pass that flattens both shapes to plain anyOf of inlined object branches, with the OAS discriminator keyword stripped. The Literal constraints on the tag field already carry the discriminator signal, so dropping the OAS keyword is a no-op semantically but makes the schema acceptable to tool-calling APIs.

_flatten_discriminated_union is non-mutating (returns a new schema) and defensively merges into any pre-existing top-level anyOf rather than overwriting. Flattening is single-level; nested discriminated unions are tracked alongside the recursive $ref resolution work in #911 (follow-up filed as #1105).

Before / after

class Cat(BaseModel):
    kind: Literal["cat"]
    name: str

class Dog(BaseModel):
    kind: Literal["dog"]
    name: str
    breed: str

def act(pet: Annotated[Cat | Dog, Field(discriminator="kind")]) -> str:
    """Act on a pet.

    Args:
        pet: the pet
    """
    return "ok"

Before ({"type": "string"} is wrong):

"pet": {"type": "string", "description": "the pet"}

After (preserved union, discriminator stripped, refs inlined):

"pet": {
  "anyOf": [
    {"type": "object", "properties": {"kind": {"const": "cat", ...}, "name": {...}}, "required": ["kind", "name"], "title": "Cat"},
    {"type": "object", "properties": {"kind": {"const": "dog", ...}, "name": {...}, "breed": {...}}, "required": ["kind", "name", "breed"], "title": "Dog"}
  ],
  "title": "Pet",
  "description": "the pet"
}

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)

Test additions:

• test/backends/test_schema_helpers.py — extends _is_complex_anyof coverage for oneOf branches (with and without discriminator metadata).
• test/backends/test_discriminated_union_tools.py — new file:
  • Schema-shape assertions (required + optional variants): no collapse to {"type":"string"}, both Cat and Dog branches survive as fully-inlined object schemas, OAS discriminator keyword is stripped (asserted on both required and optional paths), no dangling $ref leaks, optional variant is removed from required.
  • Three-arm union: Annotated[Cat | Dog | Fish, Field(discriminator="kind")] preserves all three branches.
  • Non-discriminated Optional[Email] regression guard: the new pre-pass must be a no-op for plain $ref + | None shapes that the existing inliner already handles.
  • Validation round-trip via validate_tool_arguments: accepts valid {"kind":"dog", ...} and {"kind":"cat", ...} payloads, rejects bare strings, rejects dicts missing the discriminator, accepts both omitted and supplied for the optional variant.

Local regression: uv run pytest test/ -m "not qualitative" --ignore=test/stdlib/tools/test_mcp.py → 1847 passed, 1 pre-existing failure unrelated to this change (test_example_collection_sanity — depends on optional extras like mcp being installed).

Notes for reviewers

• The fix lives in one place (mellea/backends/tools.py) but applies to every tool-calling backend because they all consume MelleaTool.as_json_tool. No backend-specific changes are needed.
• oneOf → anyOf rather than the reverse: OpenAI strict-mode tool schema rejects oneOf; the Literal tag enforces uniqueness in practice, so the choice is purely about which keyword the consumers accept.
• The discriminator-stripping is a behaviour change for any consumer that may have been relying on it. Grep finds no internal consumer.
• Out of scope: the function name convert_function_to_ollama_tool is misleading (it is the canonical OpenAI-tool-format converter for all backends), but renaming is a public-API change worth a separate PR.

Follow-ups filed during review

• Recursively flatten nested discriminated unions in tool schemas #1105 — recursively flatten nested discriminated unions. This PR handles single-level only; a discriminated-union member whose own field is another discriminated union still ships oneOf + discriminator + dangling $refs on the inner field. Related to Support recursive nested-model $ref resolution in Pydantic tool parameter schemas #911.
• validate_tool_arguments does not enforce Literal/const constraints on tool fields #1106 — validate_tool_arguments does not enforce Literal / const / enum constraints. Pre-existing gap, surfaced more visibly by this PR. Affects every Literal[...] tool field, not just discriminated unions.

Attribution

• AI coding assistants used