Skip to content

Recursively flatten nested discriminated unions in tool schemas #1105

@planetf1

Description

@planetf1

Background

PR #1104 fixed the case where a tool parameter is a Pydantic discriminated
union, but the flattening only runs at the top level. If one of the union's
branches contains another discriminated union as a field, the inner union
slips through unchanged.

End-user impact

Take a command-pattern tool like this:

from typing import Annotated, Literal
from pydantic import BaseModel, Field

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

class StubUser(BaseModel):
    kind: Literal["stub"]
    name: str

class CreateUser(BaseModel):
    action: Literal["create"]
    payload: Annotated[FullUser | StubUser, Field(discriminator="kind")]

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

def admin_op(
    op: Annotated[CreateUser | DeleteUser, Field(discriminator="action")],
) -> str:
    """Admin operation.

    Args:
        op: the operation to run
    """
    ...

The outer op parameter renders correctly. The inner payload field, however,
arrives at the backend looking like:

{
  "discriminator": {"propertyName": "kind", "mapping": {...}},
  "oneOf": [{"$ref": "#/$defs/FullUser"}, {"$ref": "#/$defs/StubUser"}]
}

Ollama rejects that shape outright. OpenAI's strict tool-schema mode rejects
it too. On more permissive backends, the model is left staring at unresolved
$ref placeholders and produces a malformed payload, or the tool call
silently goes wrong. This leaves the motivating case from #989 — command-pattern
tools and polymorphic message/event parameters — partly unaddressed when the
nesting goes more than one level deep.

Workaround

Restructure the model to keep discriminated unions flat, or wrap the inner
union in a custom validator that emits a plain object schema. Both are
awkward and lose the type clarity discriminated unions exist to provide.

Success criteria

A tool whose parameters contain a discriminated union at any nesting depth
should produce a schema where:

  1. Every oneOf has been rewritten to anyOf of fully inlined object branches.
  2. The OAS-3 discriminator keyword has been stripped at every level.
  3. No $ref survives anywhere in the parameter schema.

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

Suggested direction

Have _flatten_discriminated_union (or a recursive companion) walk the
properties, items, and additionalProperties of each inlined branch and
re-apply itself. The recursion machinery overlaps heavily with #911's
recursive $ref resolution work, so the two are best tackled together.

Related

Metadata

Metadata

Assignees

No one assigned

    Labels

    area/toolsTool framework, Bash/Python tools, tool call lifecyclebugSomething isn't workingenhancementNew feature or request

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions