Implement prebuilt ToolNode, toolCondition, MessagesState

langgraph/src/graph/message.ts

@@ -21,3 +21,7 @@ export class MessageGraph extends StateGraph<BaseMessage[], Messages> {
     });
   }
 }
+
+export interface MessagesState {
+  messages: BaseMessage[];
+}

langgraph/src/graph/state.ts

@@ -254,15 +254,15 @@ export class CompiledStateGraph<
 
     // add node and output channel
     if (key === START) {
-      this.nodes[key] = new PregelNode({
+      this.nodes[key] = new PregelNode<State, Update>({
         tags: [TAG_HIDDEN],
         triggers: [START],
         channels: [START],
         writers: [new ChannelWrite(stateWriteEntries, [TAG_HIDDEN])],
       });
     } else {
       this.channels[key] = new EphemeralValue();
-      this.nodes[key] = new PregelNode({
+      this.nodes[key] = new PregelNode<State, Update>({
         triggers: [],
         // read state keys
         channels:

langgraph/src/prebuilt/index.ts

@@ -11,3 +11,4 @@ export {
   type ToolInvocationInterface,
   ToolExecutor,
 } from "./tool_executor.js";
+export { ToolNode, toolsCondition } from "./tool_node.js";

langgraph/src/prebuilt/tool_node.ts

@@ -0,0 +1,72 @@
+import { BaseMessage, ToolMessage, AIMessage } from "@langchain/core/messages";
+import { RunnableConfig } from "@langchain/core/runnables";
+import { Tool } from "@langchain/core/tools";
+import { RunnableCallable } from "../utils.js";
+import { END } from "../graph/graph.js";
+import { MessagesState } from "../graph/message.js";
+
+export class ToolNode extends RunnableCallable<
+  BaseMessage[] | MessagesState,
+  BaseMessage[] | MessagesState
+> {
+  /**
+  A node that runs the tools requested in the last AIMessage. It can be used
+  either in StateGraph with a "messages" key or in MessageGraph. If multiple
+  tool calls are requested, they will be run in parallel. The output will be
+  a list of ToolMessages, one for each tool call.
+  */
+
+  tools: Tool[];
+
+  constructor(tools: Tool[], name: string = "tools", tags: string[] = []) {
+    super({ name, tags, func: (input, config) => this.run(input, config) });
+    this.tools = tools;
+  }
+
+  private async run(
+    input: BaseMessage[] | MessagesState,
+    config: RunnableConfig
+  ): Promise<BaseMessage[] | MessagesState> {
+    const message = Array.isArray(input)
+      ? input[input.length - 1]
+      : input.messages[input.messages.length - 1];
+
+    if (message._getType() !== "ai") {
+      throw new Error("ToolNode only accepts AIMessages as input.");
+    }
+
+    const outputs = await Promise.all(
+      (message as AIMessage).tool_calls?.map(async (call) => {
+        const tool = this.tools.find((tool) => tool.name === call.name);
+        if (tool === undefined) {
+          throw new Error(`Tool ${call.name} not found.`);
+        }
+        const output = await tool.invoke(call.args, config);
+        return new ToolMessage({
+          name: tool.name,
+          content: typeof output === "string" ? output : JSON.stringify(output),
+          tool_call_id: call.id!,
+        });
+      }) ?? []
+    );
+
+    return Array.isArray(input) ? outputs : { messages: outputs };
+  }
+}
+
+export function toolsCondition(
+  state: BaseMessage[] | MessagesState
+): "tools" | typeof END {
+  const message = Array.isArray(state)
+    ? state[state.length - 1]
+    : state.messages[state.messages.length - 1];
+
+  if (
+    "tool_calls" in message &&
+    ((message as AIMessage).tool_calls?.length ?? 0) > 0
+  ) {
+    return "tools";
+  } else {
+    return END;
+  }
+}

langgraph/src/pregel/read.ts

@@ -90,7 +90,7 @@ interface PregelNodeArgs<RunInput, RunOutput>
   triggers: Array<string>;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   mapper?: (args: any) => any;
-  writers?: Runnable<RunOutput, RunOutput>[];
+  writers?: Runnable<RunOutput, unknown>[];
   tags?: string[];
   bound?: Runnable<RunInput, RunOutput>;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any

langgraph/src/utils.ts

@@ -14,7 +14,7 @@ export interface RunnableCallableArgs extends Partial<any> {
   recurse?: boolean;
 }
 
-export class RunnableCallable extends Runnable {
+export class RunnableCallable<I = unknown, O = unknown> extends Runnable<I, O> {
   lc_namespace: string[] = ["langgraph"];
 
   // eslint-disable-next-line @typescript-eslint/no-explicit-any