How to use Union parameter of object as tool input?

OSS Product Help LangGraph
1

I am building a tool using langgraph. This tool is a todolist that includes three commands: create, insert, and modify. I want the input parameter to be a union object, so I used TodoInput as the args_schema parameter, but it throws an error at runtime. How can I use Union object to cover this problem?

here are the error msg

1 validation error for TodoInput
command
  Input should be a valid dictionary or object to extract fields from [type=model_attributes_type, input_value='{"action": "create", "to...over"]}', input_type=str]
    For further information visit https://errors.pydantic.dev/2.11/v/model_attributes_type

the llm inputs is:

{
    "role": "assistant",
    "content": "",
    "tool_calls": [
        {
            "name": "set_todo_list",
            "args": {
                "command": {
                    "action": "create",
                    "todo_list": [
                        "check loss rate",
                        "check rtt rate",
                        "summarize the data"
                    ]
                }
            },
            "id": "call_e38267c977544c06b5526f49",
            "type": "tool_call"
        }
    ]
}

The tool code is here

class create(BaseModel):
    """ Create a TODO list, if one already exists, the creation will overwrite it"""
    action: Literal["create"] = Field(description="Action type: create")
    todo_list: List[str] = Field(..., description="List of todos to create, do not include numeric prefix labels")


class insert(BaseModel):
    """ Insert new todos into the currently existing TODO list """
    action: Literal["insert"] = Field(description="Action type: insert")
    index: int = Field(..., description="Index where the todo should be inserted, the inserted todo will be placed after this index, indices start from 1, for example: if there are three todos currently, index=2, then the inserted todo will be inserted after the second todo")
    todo_list: List[str] = Field(..., description="List of todos to insert, do not include numeric prefix labels")

class modify(BaseModel):
    """ Modify todo current status """
    action: Literal["modify"] = Field(description="Action type: modify status")
    index: int = Field(..., description="Index of the todo to modify, indices start from 1")
    modify_type: Literal['stop', 'done'] = Field(..., description="Status to modify the todo to, 'stop' means cancel, representing no need to complete or skip, 'done' means complete, representing this todo is finished")

class TodoInput(BaseModel):
    command: Union[create, insert, modify] = Field(..., discriminator='action', description="Your todolist command, must follow the provided format")

class Todo(TypedDict):
    todo_text: str
    status: Literal['stop', 'done', 'doing']

@tool(args_schema=TodoInput)
def set_todo_list(command: Union[create, insert, modify], runtime: ToolRuntime) -> Command:
    """
    This is a tool for creating TODO lists, you need to adjust the type of input parameter 'command' according to different purposes. Create, add, or modify status of your todo list.
    """
    state = runtime.state
    todo_list = None
    if isinstance(command, create):
        todo_list = []
        for todo_text in command.todo_list:
            todo_list.append(Todo(todo_text=todo_text, status='doing'))
    elif isinstance(command, insert):
        todo_list = state.todo_list
        # Fix: Handle command.todo_list properly (it should be a single item for each insert)
        for todo_text in command.todo_list:
            todo_list.insert(command.index, Todo(todo_text=todo_text, status='doing'))
    elif isinstance(command, modify):
        todo_list = state.todo_list
        todo_list[command.index - 1]['status'] = command.modify_type
    else:
        return Command(update={
            "messages": [
                ToolMessage(
                    content="Failed to set todolist",
                    tool_call_id=runtime.tool_call_id
                )
            ]
        })
    runtime.stream_writer(TodoListMsg(todo_list))
    return Command(update={
        "messages": [
            ToolMessage(
                content="Successfully set todolist",
                tool_call_id=runtime.tool_call_id
            )
        ],
        "todo_list": todo_list
    })
2

another error found when I use this code. The runtime parameter is error

 set_todo_list() missing 1 required positional argument: 'runtime'

@tool(args_schema=TodoInput)
def set_todo_list(action:Literal["create", "insert", "modify"], create_param: create, insert_param: insert, modify_param: modify
                  ,runtime: ToolRuntime) -> Command:
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
4

Thanks a lot!

1 Like