Created
October 19, 2025 07:09
-
-
Save ifedayoprince/791a4cdf8207ca6513c76c8e9fe0c5b6 to your computer and use it in GitHub Desktop.
Passing State to Tools in LanggraphJS
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import { AIMessage, BaseMessage, ToolMessage } from "@langchain/core/messages"; | |
| import { StructuredTool, tool } from "@langchain/core/tools"; | |
| import { Command } from "@langchain/langgraph"; | |
| import z, { ZodType } from "zod"; | |
| type StateToolOptions<TInput extends ZodType> = { | |
| name: string; | |
| description?: string; | |
| schema: TInput; | |
| }; | |
| /** | |
| * A typed abstraction over LangChain `tool` that automatically | |
| * injects the graph state into the tool handler when invoked. | |
| */ | |
| export function stateTool< | |
| TState, | |
| TInput extends ZodType, | |
| TResult = unknown | |
| >( | |
| handler: (state: TState, input: z.infer<TInput>) => Promise<TResult> | TResult, | |
| options: StateToolOptions<TInput> | |
| ) { | |
| const baseTool = tool( | |
| async (input: unknown) => { | |
| throw new Error( | |
| `StateTool "${options.name}" cannot be invoked without state. Use .invoke(state, input).` | |
| ); | |
| }, | |
| options | |
| ); | |
| async function invoke(state: TState, input: z.infer<TInput>): Promise<TResult> { | |
| const parsed = options.schema.parse(input); | |
| return handler(state, parsed); | |
| } | |
| return Object.assign(baseTool, { | |
| invoke, | |
| }) as StatefulTool<TState, z.infer<TInput>, TResult>; | |
| } | |
| /** | |
| * A tool that can be invoked with a state and input. | |
| */ | |
| export type StatefulTool<TState, TInput, TResult> = Omit<StructuredTool, "invoke"> & { | |
| invoke: (state: TInput, input: TState) => Promise<TResult>; | |
| }; | |
| /** | |
| * The equivalent of Langchain's 'ToolNode' class, but for tools. | |
| * @example graph.addNode("tools", StatefulToolNode(tools, "text-renderer")) | |
| * @param tools The tools to use. | |
| * @param goto The state to transition to after the tool has been invoked. | |
| * @returns | |
| */ | |
| export function StatefulToolNode<TState extends { messages: BaseMessage[] }>(tools: StatefulTool<any, any, any>[], goto: string) { | |
| const toolMap = new Map<string, StatefulTool<any, any, any>>(); | |
| tools.forEach(tool => toolMap.set(tool.name, tool)); | |
| return async (state: TState) => { | |
| const lastMessage = state.messages[state.messages.length - 1] as AIMessage; | |
| const toolCallResponses: ToolMessage[] = []; | |
| let update = { | |
| messages: [] | |
| }; | |
| if ('tool_calls' in lastMessage && lastMessage.tool_calls.length > 0) { | |
| for (const toolCall of lastMessage.tool_calls) { | |
| const tool = toolMap.get(toolCall.name); | |
| if (!tool) { | |
| console.error(`tool ${toolCall.name} not found`); | |
| continue; | |
| } | |
| const result = await tool.invoke(state, toolCall.args); | |
| if (result instanceof ToolCommand) { | |
| const toolCallResult = new ToolMessage({ | |
| name: toolCall.name, | |
| tool_call_id: toolCall.id, | |
| content: result.params?.content | |
| }); | |
| toolCallResponses.push(toolCallResult) | |
| update = deepMerge(update, result.params?.update); | |
| } else { | |
| const toolCallResult = new ToolMessage({ | |
| name: toolCall.name, | |
| tool_call_id: toolCall.id, | |
| content: result | |
| }); | |
| toolCallResponses.push(toolCallResult); | |
| } | |
| } | |
| } | |
| return new Command({ | |
| update: { | |
| ...update, | |
| messages: [...toolCallResponses, ...(update.messages ?? [])] | |
| }, | |
| goto | |
| }) | |
| } | |
| } | |
| /** | |
| * The equivalent of Langgraph's 'Command' class, but for tools. | |
| * @param update The update to apply to the state. | |
| * @param content The tool's response | |
| */ | |
| export class ToolCommand<TState extends { messages: BaseMessage[] }> { | |
| constructor( | |
| public readonly params: { | |
| update?: Partial<TState>, | |
| content?: any | |
| } | |
| ) { } | |
| } | |
| function isObject(obj) { | |
| return obj !== null && typeof obj === "object" && !Array.isArray(obj); | |
| } | |
| function deepMerge(target, source) { | |
| if (!isObject(target)) target = {}; | |
| if (!isObject(source)) return target; | |
| for (const key of Object.keys(source)) { | |
| const tVal = target[key]; | |
| const sVal = source[key]; | |
| if (Array.isArray(tVal) && Array.isArray(sVal)) { | |
| target[key] = [...tVal, ...sVal]; | |
| } else if (isObject(tVal) && isObject(sVal)) { | |
| target[key] = deepMerge({ ...tVal }, sVal); | |
| } else { | |
| target[key] = sVal; | |
| } | |
| } | |
| return target; | |
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| // graph.ts | |
| const graph = new StateGraphgraph(AgentState) | |
| .addNode("tools", StatefulToolNode(tools, "goto")) | |
| .compile(); | |
| // tools.ts | |
| export const getWeather = stateTool(async (state: typeof AgentState.State, input) => { | |
| return new ToolCommand({ | |
| content: "The weather is sunny" | |
| }); | |
| }); | |
| // model.ts | |
| model.bindTools([getWeather]); |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment