How to use Union parameter of object as tool input?

OSS Product Help LangGraph
3

regarding that error → How to access tool call id using runtime - #2 by pawel-twardziak

And regarding the main issue, I think it is a mismatch in the type: string instead of dict.

Try this (Pydantic v2 style):

from typing import Annotated, Union
from pydantic import BaseModel, Field, field_validator
from typing import Literal
import json

class Create(BaseModel):
    """Create a TODO list; overwrites if one exists."""
    action: Literal["create"]
    todo_list: list[str]

class Insert(BaseModel):
    """Insert new todos into the existing list."""
    action: Literal["insert"]
    index: int
    todo_list: list[str]

class Modify(BaseModel):
    """Modify todo status."""
    action: Literal["modify"]
    index: int
    modify_type: Literal["stop", "done"]

Command = Annotated[Union[Create, Insert, Modify], Field(discriminator="action")]

class TodoInput(BaseModel):
    command: Command

    @field_validator("command", mode="before")
    @classmethod
    def parse_command_json(cls, v):
        # If the tool adapter / model sent a JSON string, parse it first
        if isinstance(v, str):
            return json.loads(v)
        return v

Then keep your tool like:

from langchain.tools import tool, ToolRuntime

@tool(args_schema=TodoInput)
def set_todo_list(command: Command, runtime: ToolRuntime) -> Command:
    state = runtime.state

    if isinstance(command, Create):
        todo_list = [ {"todo_text": t, "status": "doing"} for t in command.todo_list ]
    elif isinstance(command, Insert):
        todo_list = state.todo_list
        for t in command.todo_list:
            todo_list.insert(command.index, {"todo_text": t, "status": "doing"})
    elif isinstance(command, Modify):
        todo_list = state.todo_list
        todo_list[command.index - 1]["status"] = command.modify_type
    else:
        ...

    ...

Or flatten the schema (no union in tool args)

If you want maximum compatibility and simplicity, avoid a union parameter at the tool boundary and branch on action yourself:

class TodoInput(BaseModel):
    action: Literal["create", "insert", "modify"]
    todo_list: list[str] | None = None
    index: int | None = None
    modify_type: Literal["stop", "done"] | None = None

@tool(args_schema=TodoInput)
def set_todo_list(
    action: Literal["create", "insert", "modify"],
    todo_list: list[str] | None = None,
    index: int | None = None,
    modify_type: Literal["stop", "done"] | None = None,
    runtime: ToolRuntime,
) -> Command:
    state = runtime.state

    if action == "create":
        ...
    elif action == "insert":
        ...
    elif action == "modify":
        ...

The LLM then always sends:

{
  "action": "create",
  "todo_list": ["check loss rate", "check rtt rate", "summarize the data"]
}

which Pydantic can validate without any special handling.

2 Likes