core[patch]: Move async generator consumption code into local storage context (#5439)

* Move async generator consumption code into local storage context

* Typing fix, adds test, remove unnecessary void

* lint n stuff

---------

Co-authored-by: bracesproul <braceasproul@gmail.com>

Author: jacoblee93

langchain-core/src/runnables/base.ts

@@ -298,9 +298,11 @@ export abstract class Runnable<
   ): Promise<IterableReadableStream<RunOutput>> {
     // Buffer the first streamed chunk to allow for initial errors
     // to surface immediately.
-    const wrappedGenerator = new AsyncGeneratorWithSetup(
-      this._streamIterator(input, ensureConfig(options))
-    );
+    const config = ensureConfig(options);
+    const wrappedGenerator = new AsyncGeneratorWithSetup({
+      generator: this._streamIterator(input, config),
+      config,
+    });
     await wrappedGenerator.setup;
     return IterableReadableStream.fromAsyncGenerator(wrappedGenerator);
   }
@@ -1945,9 +1947,11 @@ export class RunnableMap<
     async function* generator() {
       yield input;
     }
-    const wrappedGenerator = new AsyncGeneratorWithSetup(
-      this.transform(generator(), options)
-    );
+    const config = ensureConfig(options);
+    const wrappedGenerator = new AsyncGeneratorWithSetup({
+      generator: this.transform(generator(), config),
+      config,
+    });
     await wrappedGenerator.setup;
     return IterableReadableStream.fromAsyncGenerator(wrappedGenerator);
   }
@@ -2151,9 +2155,11 @@ export class RunnableLambda<RunInput, RunOutput> extends Runnable<
     async function* generator() {
       yield input;
     }
-    const wrappedGenerator = new AsyncGeneratorWithSetup(
-      this.transform(generator(), options)
-    );
+    const config = ensureConfig(options);
+    const wrappedGenerator = new AsyncGeneratorWithSetup({
+      generator: this.transform(generator(), config),
+      config,
+    });
     await wrappedGenerator.setup;
     return IterableReadableStream.fromAsyncGenerator(wrappedGenerator);
   }
@@ -2458,9 +2464,11 @@ export class RunnableAssign<
     async function* generator() {
       yield input;
     }
-    const wrappedGenerator = new AsyncGeneratorWithSetup(
-      this.transform(generator(), options)
-    );
+    const config = ensureConfig(options);
+    const wrappedGenerator = new AsyncGeneratorWithSetup({
+      generator: this.transform(generator(), config),
+      config,
+    });
     await wrappedGenerator.setup;
     return IterableReadableStream.fromAsyncGenerator(wrappedGenerator);
   }
@@ -2549,9 +2557,11 @@ export class RunnablePick<
     async function* generator() {
       yield input;
     }
-    const wrappedGenerator = new AsyncGeneratorWithSetup(
-      this.transform(generator(), options)
-    );
+    const config = ensureConfig(options);
+    const wrappedGenerator = new AsyncGeneratorWithSetup({
+      generator: this.transform(generator(), config),
+      config,
+    });
     await wrappedGenerator.setup;
     return IterableReadableStream.fromAsyncGenerator(wrappedGenerator);
   }

langchain-core/src/runnables/passthrough.ts

@@ -40,7 +40,8 @@ type RunnablePassthroughFunc<RunInput = any> =
  * );
  * ```
  */
-export class RunnablePassthrough<RunInput> extends Runnable<
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export class RunnablePassthrough<RunInput = any> extends Runnable<
   RunInput,
   RunInput
 > {

langchain-core/src/runnables/tests/runnable_stream_events.test.ts

@@ -4,7 +4,12 @@
 
 import { test } from "@jest/globals";
 import { z } from "zod";
-import { RunnableLambda } from "../index.js";
+import {
+  RunnableLambda,
+  RunnableMap,
+  RunnablePassthrough,
+  RunnablePick,
+} from "../index.js";
 import { ChatPromptTemplate } from "../../prompts/chat.js";
 import {
   FakeListChatModel,
@@ -381,6 +386,166 @@ test("Runnable streamEvents method with three runnables with filtering", async (
   ]);
 });
 
+test("Runnable streamEvents method with a runnable map", async () => {
+  const r = RunnableLambda.from(reverse);
+
+  const chain = RunnableMap.from({
+    reversed: r,
+    original: new RunnablePassthrough(),
+  }).pipe(new RunnablePick("reversed"));
+
+  const events = [];
+  const eventStream = await chain.streamEvents("hello", { version: "v1" });
+  for await (const event of eventStream) {
+    events.push(event);
+  }
+  console.log(events);
+  expect(events).toEqual([
+    {
+      run_id: expect.any(String),
+      event: "on_chain_start",
+      name: "RunnableSequence",
+      tags: [],
+      metadata: {},
+      data: { input: "hello" },
+    },
+    {
+      event: "on_chain_start",
+      name: "RunnableMap",
+      run_id: expect.any(String),
+      tags: ["seq:step:1"],
+      metadata: {},
+      data: {},
+    },
+    {
+      event: "on_chain_start",
+      name: "RunnableLambda",
+      run_id: expect.any(String),
+      tags: ["map:key:reversed"],
+      metadata: {},
+      data: {},
+    },
+    {
+      event: "on_chain_start",
+      name: "RunnablePassthrough",
+      run_id: expect.any(String),
+      tags: ["map:key:original"],
+      metadata: {},
+      data: {},
+    },
+    {
+      event: "on_chain_stream",
+      name: "RunnablePassthrough",
+      run_id: expect.any(String),
+      tags: ["map:key:original"],
+      metadata: {},
+      data: { chunk: "hello" },
+    },
+    {
+      event: "on_chain_stream",
+      name: "RunnableLambda",
+      run_id: expect.any(String),
+      tags: ["map:key:reversed"],
+      metadata: {},
+      data: { chunk: "olleh" },
+    },
+    {
+      event: "on_chain_stream",
+      name: "RunnableMap",
+      run_id: expect.any(String),
+      tags: ["seq:step:1"],
+      metadata: {},
+      data: {
+        chunk: {
+          original: "hello",
+        },
+      },
+    },
+    {
+      event: "on_chain_start",
+      name: "RunnablePick",
+      run_id: expect.any(String),
+      tags: ["seq:step:2"],
+      metadata: {},
+      data: {},
+    },
+    {
+      event: "on_chain_stream",
+      name: "RunnableMap",
+      run_id: expect.any(String),
+      tags: ["seq:step:1"],
+      metadata: {},
+      data: {
+        chunk: {
+          reversed: "olleh",
+        },
+      },
+    },
+    {
+      event: "on_chain_end",
+      name: "RunnablePassthrough",
+      run_id: expect.any(String),
+      tags: ["map:key:original"],
+      metadata: {},
+      data: { input: "hello", output: "hello" },
+    },
+    {
+      event: "on_chain_stream",
+      name: "RunnablePick",
+      run_id: expect.any(String),
+      tags: ["seq:step:2"],
+      metadata: {},
+      data: { chunk: "olleh" },
+    },
+    {
+      event: "on_chain_stream",
+      run_id: expect.any(String),
+      tags: [],
+      metadata: {},
+      name: "RunnableSequence",
+      data: { chunk: "olleh" },
+    },
+    {
+      event: "on_chain_end",
+      name: "RunnableLambda",
+      run_id: expect.any(String),
+      tags: ["map:key:reversed"],
+      metadata: {},
+      data: { input: "hello", output: "olleh" },
+    },
+    {
+      event: "on_chain_end",
+      name: "RunnableMap",
+      run_id: expect.any(String),
+      tags: ["seq:step:1"],
+      metadata: {},
+      data: {
+        input: "hello",
+        output: {
+          original: "hello",
+          reversed: "olleh",
+        },
+      },
+    },
+    {
+      event: "on_chain_end",
+      name: "RunnablePick",
+      run_id: expect.any(String),
+      tags: ["seq:step:2"],
+      metadata: {},
+      data: { output: "olleh" },
+    },
+    {
+      event: "on_chain_end",
+      name: "RunnableSequence",
+      run_id: expect.any(String),
+      tags: [],
+      metadata: {},
+      data: { output: "olleh" },
+    },
+  ]);
+});
+
 test("Runnable streamEvents method with llm", async () => {
   const model = new FakeStreamingLLM({
     responses: ["hey!"],

langchain-core/src/singletons/tests/async_local_storage.test.ts

@@ -131,3 +131,31 @@ test("Config should be automatically populated after setting global async local
     events.filter((event) => event.event === "on_llm_start").length
   ).toEqual(1);
 });
+
+test("Runnable streamEvents method with streaming nested in a RunnableLambda", async () => {
+  AsyncLocalStorageProviderSingleton.initializeGlobalInstance(
+    new AsyncLocalStorage()
+  );
+  const chat = new FakeListChatModel({
+    responses: ["Hello"],
+  });
+  const myFunc = async (input: string) => {
+    for await (const _ of await chat.stream(input)) {
+      // no-op
+    }
+  };
+
+  const myNestedLambda = RunnableLambda.from(myFunc);
+
+  const events = [];
+  for await (const event of myNestedLambda.streamEvents("hello", {
+    version: "v1",
+  })) {
+    console.log(event);
+    events.push(event);
+  }
+  const chatModelStreamEvent = events.find((event) => {
+    return event.event === "on_llm_stream";
+  });
+  expect(chatModelStreamEvent).toBeDefined();
+});

langchain-core/src/utils/stream.ts

@@ -1,5 +1,7 @@
 // Make this a type to override ReadableStream's async iterator type in case
 // the popular web-streams-polyfill is imported - the supplied types
+import { AsyncLocalStorageProviderSingleton } from "../singletons/index.js";
+
 // in this case don't quite match.
 export type IterableReadableStreamInterface<T> = ReadableStream<T> &
   AsyncIterable<T>;
@@ -182,23 +184,33 @@ export class AsyncGeneratorWithSetup<
 
   public setup: Promise<S>;
 
+  public config?: unknown;
+
   private firstResult: Promise<IteratorResult<T>>;
 
   private firstResultUsed = false;
 
-  constructor(generator: AsyncGenerator<T>, startSetup?: () => Promise<S>) {
-    this.generator = generator;
+  constructor(params: {
+    generator: AsyncGenerator<T>;
+    startSetup?: () => Promise<S>;
+    config?: unknown;
+  }) {
+    this.generator = params.generator;
+    this.config = params.config;
     // setup is a promise that resolves only after the first iterator value
     // is available. this is useful when setup of several piped generators
     // needs to happen in logical order, ie. in the order in which input to
     // to each generator is available.
     this.setup = new Promise((resolve, reject) => {
-      this.firstResult = generator.next();
-      if (startSetup) {
-        this.firstResult.then(startSetup).then(resolve, reject);
-      } else {
-        this.firstResult.then((_result) => resolve(undefined as S), reject);
-      }
+      const storage = AsyncLocalStorageProviderSingleton.getInstance();
+      void storage.run(params.config, async () => {
+        this.firstResult = params.generator.next();
+        if (params.startSetup) {
+          this.firstResult.then(params.startSetup).then(resolve, reject);
+        } else {
+          this.firstResult.then((_result) => resolve(undefined as S), reject);
+        }
+      });
     });
   }
 
@@ -208,7 +220,10 @@ export class AsyncGeneratorWithSetup<
       return this.firstResult;
     }
 
-    return this.generator.next(...args);
+    const storage = AsyncLocalStorageProviderSingleton.getInstance();
+    return storage.run(this.config, async () => {
+      return this.generator.next(...args);
+    });
   }
 
   async return(
@@ -245,7 +260,7 @@ export async function pipeGeneratorWithSetup<
   startSetup: () => Promise<S>,
   ...args: A
 ) {
-  const gen = new AsyncGeneratorWithSetup(generator, startSetup);
+  const gen = new AsyncGeneratorWithSetup({ generator, startSetup });
   const setup = await gen.setup;
   return { output: to(gen, setup, ...args), setup };
 }