Route each tool from the ToolNode to a different graph chain.

[shlok-ansys]

Hello, I want to route the Tools node to different chain of the graph.

For example, A graph having 2 tools, I the next Stage of both the tools to be controlled individually in the graph, Rather than them being controlled by the 'tools' node

tools = [multiply, get_more_context_information]
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))

builder.add_edge(START, "assistant")
builder.add_conditional_edges(
    "assistant",
    tools_condition,
)
builder.add_edge("tools", END)

---

But i want something like this

What is the best way to achieve this ?

As the prebuilt tools_condition, ToolNode are amazing already, How can i utilize them to do something like this ?

Thanks !

[vbarda]

you could just add two separate tool nodes for each tool, e.g.

builder.add_node("get_more_context_information", ToolNode([context_tool]))
builder.add_node("mul", ToolNode([mul]))

and then add conditional edges to route to the correct node accordingly based on the tool_calls in the last AI message (similar to the tools_condition)

[shlok-ansys]

Thanks for the quick response @vbarda,
This is what i have come up to, Is this a good way ? Or there is a better way ?

class MyState(MessagesState):
    multiplication_result: int = 0
    final_result: any = ""

class MultiplyModelSchema(BaseModel):
    a: int = Field(..., title="First number")
    b: int = Field(..., title="Second number")

def multiply(numbers: MultiplyModelSchema, tool_call_id: Annotated[str, InjectedToolCallId]) -> Command:
    """Multiply a and b"""
    return Command(
        update={
            "multiplication_result": numbers.a * numbers.b,
            "messages": [ToolMessage("Successfully looked up user information", tool_call_id=tool_call_id)]

        },
        # goto="do_something_with_the_answer"
    )


def do_something_with_the_answer(state: MyState) -> MyState:
    """This function does something with the answer."""
    return {
        "final_result": "The result of the multiplication is " + str(state["multiplication_result"] + 10),
    }

def get_more_user_context_info():
    """This is a helper function that returns more context info."""
    return "This is a more context info. My name is shlok"

tools = [multiply, get_more_user_context_info]
llm_with_tools = llm.bind_tools(tools)

# System message
sys_msg = SystemMessage(content="You are a helpful assistant tasked with performing arithmetic on a set of inputs.")

# Node
def assistant(state: MyState):
   return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}

def conditional_tool_calling(state: MyState):
    # Check if the latest message is a tool call
    ai_message = state["messages"][-1]
    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
        return ai_message.tool_calls[0]["name"]

    return END

builder = StateGraph(MyState)

# Define nodes: these do the work
builder.add_node("assistant", assistant)
builder.add_node("get_more_user_context_info", ToolNode([get_more_user_context_info]))
builder.add_node("multiply", ToolNode([multiply]))
builder.add_node("do_something_with_the_answer", do_something_with_the_answer)

# Define edges: these determine the control flow
builder.add_edge(START, "assistant")
builder.add_conditional_edges(
    "assistant",
    conditional_tool_calling,
    ["multiply", "get_more_user_context_info", END],
)
builder.add_edge("get_more_user_context_info", "assistant")
builder.add_edge("multiply", "do_something_with_the_answer")
builder.add_edge("do_something_with_the_answer", END)
builder.add_edge("assistant", END)

graph = builder.compile()