Call a tool with both args_schema and InjectedState results in error

[NathanAP]

Hi, I have two different classes in my Tool. One to control my State values and another to control the AI generated values (or args_schema).

Heres a simple idea:

class FindCategory(BaseTool):
    """"""

    name: str = "find_category"
    description: str = "Use this tool to find all available categories when needed."
    args_schema: Type[MyModel] = MyModel

    def _run(
        self,
        category_name: Optional[str],
        state: Annotated[MyStateModel, InjectedState],
    ):

    # my logic

Now when I do:

state = MyStateModel(messages=state.messages, ...)

await action.ainvoke(state)

I get Error: TypeError("FindCategory._run() missing 1 required positional argument: \'state\'").

Now if I do:

await action.ainvoke(input=state.messages, state=state)

I get ToolNode._afunc() got an unexpected keyword argument 'state'.

How am I supposed to do this?
In the first example of this part of documentation, it says I should be able to have some_arg and state at the same time, but it doesn't show how to call it.

[vbarda]

Hi - the state part needs to be defined in the args schema model (e.g. in MyModel in your case)

[NathanAP]

The problem is that my state is quite complex because it has Enum, Dict, other models, etc... is this a problem somehow?
The idea behind using two classes was to prevent useless AI data generation

[vbarda]

anything marked with InjectedState will be ignored by the LLM, so it's fine to put the state into your pydantic model arg schema

e.g.

class MyModel(BaseModel):
    ...
    state: Annotated[MyStateModel, InjectedState],

[NathanAP]

I'll test it tomorrow! 🙏🏻

[NathanAP]

It worked @vbarda
Thanks ❤