Custom Tool Node.

[mgathena]

I used to create a custom tool node like this.

def db_tool_node(state: dict, config: RunnableConfig) -> dict[str, Any]:
    """Initialise a tool node that executes the tools and writes
    all returned keys and values to the state variable.

    TODO: This should be a modular class.
    """
    messages = []
    out = {}
    for tool_call in state["messages"][-1].tool_calls:
        tool = tools_by_name[tool_call["name"]]
        state_args = {var: state for var in _get_state_args(tool)}
        observation = tool.invoke({**tool_call["args"], **state_args}, config=config)

        # observation = observation['score']
        # print(f"observation: {observation}")
        message_content = observation

        messages.append(
            ToolMessage(content=message_content, tool_call_id=tool_call["id"])
        )
        # out = {**out, **observation}
        out = {**out, "observation": observation}

    return {"messages": messages, **out}

But I can't import _get_state_args via from langgraph.prebuilt.tool_node import _get_state_args after upgrading to langgraph > 0.3. What is the workaround? I found some discussion here but I couldn't get it resolved. Also, it is found here, but I can't figure a way to import it.