feat(langgraph): add typedNode utility

Author: dqbd

libs/langgraph/src/graph/graph.ts

@@ -167,11 +167,11 @@ export type NodeSpec<RunInput, RunOutput> = {
   ends?: string[];
 };
 
-export type AddNodeOptions = {
+export type AddNodeOptions<Nodes extends string = string> = {
   metadata?: Record<string, unknown>;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   subgraphs?: Pregel<any, any>[];
-  ends?: string[];
+  ends?: Nodes[];
 };
 
 export class Graph<

libs/langgraph/src/graph/state.ts

@@ -82,12 +82,12 @@ export type StateGraphNodeSpec<RunInput, RunOutput> = NodeSpec<
   retryPolicy?: RetryPolicy;
 };
 
-export type StateGraphAddNodeOptions = {
+export type StateGraphAddNodeOptions<Nodes extends string = string> = {
   retryPolicy?: RetryPolicy;
   // TODO: Fix generic typing for annotations
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   input?: AnnotationRoot<any> | AnyZodObject;
-} & AddNodeOptions;
+} & AddNodeOptions<Nodes>;
 
 export type StateGraphArgsWithStateSchema<
   SD extends StateDefinition,
@@ -237,7 +237,8 @@ export class StateGraph<
     fields: SD extends StateDefinition
       ? StateGraphArgsWithInputOutputSchemas<SD, ToStateDefinition<O>>
       : never,
-    configSchema?: C | AnnotationRoot<ToStateDefinition<C>>
+    configSchema?: C | AnnotationRoot<ToStateDefinition<C>>,
+    options?: { nodes?: N[] }
   );
 
   constructor(
@@ -252,14 +253,16 @@ export class StateGraph<
               ToStateDefinition<O>
             >
       : StateGraphArgs<S>,
-    configSchema?: C | AnnotationRoot<ToStateDefinition<C>>
+    configSchema?: C | AnnotationRoot<ToStateDefinition<C>>,
+    options?: { nodes?: N[] }
   );
 
   constructor(
     fields: SD extends AnyZodObject
       ? SD | ZodStateGraphArgsWithStateSchema<SD, I, O>
       : never,
-    configSchema?: C | AnnotationRoot<ToStateDefinition<C>>
+    configSchema?: C | AnnotationRoot<ToStateDefinition<C>>,
+    options?: { nodes?: N[] }
   );
 
   constructor(
@@ -277,7 +280,8 @@ export class StateGraph<
             >
           | StateGraphArgsWithInputOutputSchemas<SD, ToStateDefinition<O>>
       : StateGraphArgs<S>,
-    configSchema?: C | AnnotationRoot<ToStateDefinition<C>>
+    configSchema?: C | AnnotationRoot<ToStateDefinition<C>>,
+    _options?: { nodes?: N[] }
   ) {
     super();
 
@@ -434,7 +438,12 @@ export class StateGraph<
       isMultipleNodes(args) // eslint-disable-line no-nested-ternary
         ? Array.isArray(args[0])
           ? args[0]
-          : Object.entries(args[0])
+          : Object.entries(args[0]).map(([key, action]) => [
+              key,
+              action,
+              // eslint-disable-next-line @typescript-eslint/no-explicit-any
+              (action as any)[Symbol.for("langgraph.state.node")] ?? undefined,
+            ])
         : [[args[0], args[1], args[2]]]
     ) as [
       K,
@@ -574,7 +583,12 @@ export class StateGraph<
   ): StateGraph<SD, S, U, N | K, I, O, C> {
     const parsedNodes = Array.isArray(nodes)
       ? nodes
-      : (Object.entries(nodes) as [K, NodeAction<S, U, C>][]);
+      : (Object.entries(nodes).map(([key, action]) => [
+          key,
+          action,
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+          (action as any)[Symbol.for("langgraph.state.node")] ?? undefined,
+        ]) as [K, NodeAction<S, U, C>, StateGraphAddNodeOptions | undefined][]);
 
     if (parsedNodes.length === 0) {
       throw new Error("Sequence requires at least one node.");
@@ -1091,3 +1105,41 @@ function _getControlBranch() {
     path: CONTROL_BRANCH_PATH,
   });
 }
+
+type TypedNodeAction<SD extends StateDefinition, Nodes extends string> = (
+  state: StateType<SD>,
+  config: LangGraphRunnableConfig
+) => UpdateType<SD> | Command<unknown, UpdateType<SD>, Nodes>;
+
+export function typedNode<SD extends SDZod, Nodes extends string>(
+  _state: SD extends StateDefinition ? AnnotationRoot<SD> : never,
+  _options?: { nodes?: Nodes[] }
+): (
+  func: TypedNodeAction<ToStateDefinition<SD>, Nodes>,
+  options?: StateGraphAddNodeOptions<Nodes>
+) => TypedNodeAction<ToStateDefinition<SD>, Nodes>;
+
+export function typedNode<SD extends SDZod, Nodes extends string>(
+  _state: SD extends AnyZodObject ? SD : never,
+  _options?: { nodes?: Nodes[] }
+): (
+  func: TypedNodeAction<ToStateDefinition<SD>, Nodes>,
+  options?: StateGraphAddNodeOptions<Nodes>
+) => TypedNodeAction<ToStateDefinition<SD>, Nodes>;
+
+export function typedNode<SD extends SDZod, Nodes extends string>(
+  _state: SD extends AnyZodObject
+    ? SD
+    : SD extends StateDefinition
+    ? AnnotationRoot<SD>
+    : never,
+  _options?: { nodes?: Nodes[] }
+) {
+  return (
+    func: TypedNodeAction<ToStateDefinition<SD>, Nodes>,
+    options?: StateGraphAddNodeOptions<Nodes>
+  ) => {
+    Object.assign(func, { [Symbol.for("langgraph.state.node")]: options });
+    return func;
+  };
+}

libs/langgraph/src/tests/typed_state.test.ts

@@ -0,0 +1,89 @@
+import { z } from "zod";
+import { Command } from "../constants.js";
+import { Annotation } from "../graph/annotation.js";
+import {
+  MessagesAnnotation,
+  MessagesZodState,
+} from "../graph/messages_annotation.js";
+import { StateGraph, typedNode } from "../graph/state.js";
+import { _AnyIdHumanMessage } from "./utils.js";
+
+it("Annotation.Root", async () => {
+  const StateAnnotation = Annotation.Root({
+    messages: MessagesAnnotation.spec.messages,
+    foo: Annotation<string>,
+  });
+
+  const node = typedNode(StateAnnotation, {
+    nodes: ["nodeA", "nodeB", "nodeC"],
+  });
+
+  const nodeA = node(
+    (state) => {
+      const goto = state.foo === "foo" ? "nodeB" : "nodeC";
+      return new Command({
+        update: { messages: [{ type: "user", content: "a" }], foo: "a" },
+        goto,
+      });
+    },
+    { ends: ["nodeB", "nodeC"] }
+  );
+
+  const nodeB = node(() => {
+    return new Command({
+      goto: "nodeC",
+      update: { foo: "123" },
+    });
+  });
+  const nodeC = node((state) => ({ foo: `${state.foo}|c` }));
+
+  const graph = new StateGraph(StateAnnotation)
+    .addNode({ nodeA, nodeB, nodeC })
+    .addEdge("__start__", "nodeA")
+    .compile();
+
+  expect(await graph.invoke({ foo: "foo" })).toEqual({
+    messages: [new _AnyIdHumanMessage("a")],
+    foo: "123|c",
+  });
+});
+
+it("Zod", async () => {
+  const StateAnnotation = MessagesZodState.extend({
+    foo: z.string(),
+  });
+
+  const node = typedNode(StateAnnotation, {
+    nodes: ["nodeA", "nodeB", "nodeC"],
+  });
+
+  const nodeA = node(
+    (state) => {
+      const goto = state.foo === "foo" ? "nodeB" : "nodeC";
+      return new Command({
+        update: { messages: [{ type: "user", content: "a" }], foo: "a" },
+        goto,
+      });
+    },
+    { ends: ["nodeB", "nodeC"] }
+  );
+
+  const nodeB = node(() => {
+    return new Command({
+      goto: "nodeC",
+      update: { foo: "123" },
+    });
+  });
+
+  const nodeC = node((state) => ({ foo: `${state.foo}|c` }));
+
+  const graph = new StateGraph(StateAnnotation)
+    .addNode({ nodeA, nodeB, nodeC })
+    .addEdge("__start__", "nodeA")
+    .compile();
+
+  expect(await graph.invoke({ foo: "foo" })).toEqual({
+    messages: [new _AnyIdHumanMessage("a")],
+    foo: "123|c",
+  });
+});