fix: Return input in write method (langchain-ai#626)

benjamincburns · Oct 24, 2024 · 21d7992 · 21d7992
1 parent 70453f2
commit 21d7992
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/libs/langgraph/src/pregel/write.ts b/libs/langgraph/src/pregel/write.ts
@@ -109,9 +109,10 @@ export class ChannelWrite<
     return [...writes, ...values];
   }
 
-  async _write(input: unknown, config: RunnableConfig): Promise<void> {
+  async _write(input: unknown, config: RunnableConfig): Promise<unknown> {
     const values = await this._getWriteValues(input, config);
     ChannelWrite.doWrite(config, values);
+    return input;
   }
 
   // TODO: Support requireAtLeastOneOf

diff --git a/libs/langgraph/src/tests/tracing.test.ts b/libs/langgraph/src/tests/tracing.test.ts
@@ -72,7 +72,7 @@ it("stream events for a multi-node graph", async () => {
     {
       event: "on_chain_end",
       data: {
-        output: { output: undefined },
+        output: { messages: [] },
         input: {
           messages: [],
         },
@@ -174,7 +174,7 @@ it("stream events for a multi-node graph", async () => {
     {
       event: "on_chain_end",
       data: {
-        output: { output: undefined },
+        output: { messages: [new _AnyIdAIMessage("hey!")] },
         input: {
           messages: [new _AnyIdAIMessage("hey!")],
         },
@@ -413,7 +413,7 @@ it("Should respect .withConfig", async () => {
     {
       event: "on_chain_end",
       data: {
-        output: { output: undefined },
+        output: { messages: [] },
         input: {
           messages: [],
         },
@@ -515,7 +515,7 @@ it("Should respect .withConfig", async () => {
     {
       event: "on_chain_end",
       data: {
-        output: { output: undefined },
+        output: { messages: [new _AnyIdAIMessage("hey!")] },
         input: {
           messages: [new _AnyIdAIMessage("hey!")],
         },