fix: add pick method to base asynclocalstorage method

langchain-ai · Dec 5, 2024 · d9109f7 · d9109f7
1 parent 6c12c7b
commit d9109f7
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 20 deletions.
diff --git a/langchain-core/src/runnables/base.ts b/langchain-core/src/runnables/base.ts
@@ -41,7 +41,6 @@ import {
   getCallbackManagerForConfig,
   mergeConfigs,
   patchConfig,
-  pickRunnableConfigKeys,
 } from "./config.js";
 import { AsyncCaller } from "../utils/async_caller.js";
 import { Run } from "../tracers/base.js";
@@ -2533,7 +2532,7 @@ export class RunnableLambda<
         recursionLimit: (config?.recursionLimit ?? DEFAULT_RECURSION_LIMIT) - 1,
       });
       void AsyncLocalStorageProviderSingleton.runWithConfig(
-        pickRunnableConfigKeys(childConfig),
+        childConfig,
         async () => {
           try {
             let output = await this.func(input, {
@@ -2631,7 +2630,7 @@ export class RunnableLambda<
     const output = await new Promise<RunOutput | Runnable>(
       (resolve, reject) => {
         void AsyncLocalStorageProviderSingleton.runWithConfig(
-          pickRunnableConfigKeys(childConfig),
+          childConfig,
           async () => {
             try {
               const res = await this.func(finalChunk as RunInput, {

diff --git a/langchain-core/src/runnables/iter.ts b/langchain-core/src/runnables/iter.ts
@@ -1,6 +1,5 @@
 import type { RunnableConfig } from "../runnables/types.js";
 import { AsyncLocalStorageProviderSingleton } from "../singletons/index.js";
-import { pickRunnableConfigKeys } from "./config.js";
 
 export function isIterableIterator(
   thing: unknown
@@ -37,7 +36,7 @@ export function* consumeIteratorInContext<T>(
 ): IterableIterator<T> {
   while (true) {
     const { value, done } = AsyncLocalStorageProviderSingleton.runWithConfig(
-      pickRunnableConfigKeys(context),
+      context,
       iter.next.bind(iter),
       true
     );
@@ -57,7 +56,7 @@ export async function* consumeAsyncIterableInContext<T>(
   while (true) {
     const { value, done } =
       await AsyncLocalStorageProviderSingleton.runWithConfig(
-        pickRunnableConfigKeys(context),
+        context,
         iterator.next.bind(iter),
         true
       );

diff --git a/langchain-core/src/singletons/async_local_storage/index.ts b/langchain-core/src/singletons/async_local_storage/index.ts
@@ -7,6 +7,7 @@ import {
 } from "./globals.js";
 import { CallbackManager } from "../../callbacks/manager.js";
 import { LangChainTracer } from "../../tracers/tracer_langchain.js";
+import { pickRunnableConfigKeys } from "../../runnables/config.js";
 
 export class MockAsyncLocalStorage implements AsyncLocalStorageInterface {
   getStore(): any {
@@ -46,12 +47,13 @@ class AsyncLocalStorageProvider {
     callback: () => T,
     avoidCreatingRootRunTree?: boolean
   ): T {
+    const cleanedConfig = pickRunnableConfigKeys(config);
     const callbackManager = CallbackManager._configureSync(
-      config?.callbacks,
+      cleanedConfig?.callbacks,
       undefined,
-      config?.tags,
+      cleanedConfig?.tags,
       undefined,
-      config?.metadata
+      cleanedConfig?.metadata
     );
     const storage = this.getInstance();
     const previousValue = storage.getStore();
@@ -72,7 +74,7 @@ class AsyncLocalStorageProvider {
     }
 
     if (runTree) {
-      runTree.extra = { ...runTree.extra, [LC_CHILD_KEY]: config };
+      runTree.extra = { ...runTree.extra, [LC_CHILD_KEY]: cleanedConfig };
     }
 
     if (

diff --git a/langchain-core/src/tools/index.ts b/langchain-core/src/tools/index.ts
@@ -12,7 +12,6 @@ import {
 import {
   ensureConfig,
   patchConfig,
-  pickRunnableConfigKeys,
   type RunnableConfig,
 } from "../runnables/config.js";
 import type { RunnableFunc, RunnableInterface } from "../runnables/base.js";
@@ -595,7 +594,7 @@ export function tool<
             callbacks: runManager?.getChild(),
           });
           void AsyncLocalStorageProviderSingleton.runWithConfig(
-            pickRunnableConfigKeys(childConfig),
+            childConfig,
             async () => {
               try {
                 // TS doesn't restrict the type here based on the guard above
@@ -626,7 +625,7 @@ export function tool<
           callbacks: runManager?.getChild(),
         });
         void AsyncLocalStorageProviderSingleton.runWithConfig(
-          pickRunnableConfigKeys(childConfig),
+          childConfig,
           async () => {
             try {
               // TS doesn't restrict the type here based on the guard above

diff --git a/langchain-core/src/utils/stream.ts b/langchain-core/src/utils/stream.ts
@@ -1,4 +1,3 @@
-import { pickRunnableConfigKeys } from "../runnables/config.js";
 import { AsyncLocalStorageProviderSingleton } from "../singletons/index.js";
 import type { IterableReadableStreamInterface } from "../types/stream.js";
 import { raceWithSignal } from "./signal.js";
@@ -215,9 +214,7 @@ export class AsyncGeneratorWithSetup<
     // to each generator is available.
     this.setup = new Promise((resolve, reject) => {
       void AsyncLocalStorageProviderSingleton.runWithConfig(
-        pickRunnableConfigKeys(
-          params.config as Record<string, unknown> | undefined
-        ),
+        params.config,
         async () => {
           this.firstResult = params.generator.next();
           if (params.startSetup) {
@@ -240,9 +237,7 @@ export class AsyncGeneratorWithSetup<
     }
 
     return AsyncLocalStorageProviderSingleton.runWithConfig(
-      pickRunnableConfigKeys(
-        this.config as Record<string, unknown> | undefined
-      ),
+      this.config,
       this.signal
         ? async () => {
             return raceWithSignal(this.generator.next(...args), this.signal);