fix(langgraph): Fix behavior around array return values from nodes (#760

)
langchain-ai · Dec 24, 2024 · 3bce2cb · 3bce2cb
1 parent 1a7d4cd
commit 3bce2cb
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 5 deletions.
diff --git a/libs/langgraph/src/graph/state.ts b/libs/langgraph/src/graph/state.ts
@@ -582,7 +582,7 @@ export class CompiledStateGraph<
       } else {
         const typeofInput = Array.isArray(input) ? "array" : typeof input;
         throw new InvalidUpdateError(
-          `Expected node "${nodeKey.toString()}" to return an object, received ${typeofInput}`,
+          `Expected node "${nodeKey.toString()}" to return an object or an array containing at least one Command object, received ${typeofInput}`,
           {
             lc_error_code: "INVALID_GRAPH_NODE_RETURN_VALUE",
           }
@@ -809,13 +809,26 @@ function _controlBranch(value: any): (string | Send)[] {
   if (_isSend(value)) {
     return [value];
   }
-  if (!isCommand(value)) {
+  const commands = [];
+  if (isCommand(value)) {
+    commands.push(value);
+  } else if (Array.isArray(value) && value.every(isCommand)) {
+    commands.push(...value);
+  } else {
     return [];
   }
-  if (value.graph === Command.PARENT) {
-    throw new ParentCommand(value);
+  const destinations: (string | Send)[] = [];
+  for (const command of commands) {
+    if (command.graph === Command.PARENT) {
+      throw new ParentCommand(command);
+    }
+    if (_isSend(command.goto) || typeof command.goto === "string") {
+      destinations.push(command.goto);
+    } else {
+      destinations.push(...command.goto);
+    }
   }
-  return Array.isArray(value.goto) ? value.goto : [value.goto];
+  return destinations;
 }
 
 function _getControlBranch() {

diff --git a/libs/langgraph/src/tests/pregel.test.ts b/libs/langgraph/src/tests/pregel.test.ts
@@ -2061,6 +2061,59 @@ graph TD;
     expect(await graph.invoke({ foo: "" })).toEqual({ foo: "a|c" });
   });
 
+  it("should support a simple edgeless graph", async () => {
+    const StateAnnotation = Annotation.Root({
+      foo: Annotation<string>,
+      bar: Annotation<string>,
+    });
+
+    const nodeA = async (state: typeof StateAnnotation.State) => {
+      const goto = state.foo === "foo" ? "nodeB" : "nodeC";
+      return [
+        new Command({
+          update: {
+            foo: "a",
+          },
+          goto,
+        }),
+      ];
+    };
+
+    const nodeB = async (state: typeof StateAnnotation.State) => {
+      return [
+        {
+          foo: state.foo + "|b",
+        },
+        new Command({
+          update: {
+            bar: "test",
+          },
+        }),
+      ];
+    };
+
+    const nodeC = async (state: typeof StateAnnotation.State) => {
+      return {
+        foo: state.foo + "|c",
+      };
+    };
+
+    const graph = new StateGraph(StateAnnotation)
+      .addNode("nodeA", nodeA, {
+        ends: ["nodeB", "nodeC"],
+      })
+      .addNode("nodeB", nodeB)
+      .addNode("nodeC", nodeC)
+      .addEdge("__start__", "nodeA")
+      .compile();
+
+    expect(await graph.invoke({ foo: "foo" })).toEqual({
+      foo: "a|b",
+      bar: "test",
+    });
+    expect(await graph.invoke({ foo: "" })).toEqual({ foo: "a|c" });
+  });
+
   it("should handle send sequences correctly", async () => {
     const StateAnnotation = Annotation.Root({
       items: Annotation<any[]>({