From d9109f7a25986ef2363beebde0c339fafec5ae3d Mon Sep 17 00:00:00 2001 From: bracesproul Date: Thu, 5 Dec 2024 14:05:36 -0800 Subject: [PATCH] fix: add pick method to base asynclocalstorage method --- langchain-core/src/runnables/base.ts | 5 ++--- langchain-core/src/runnables/iter.ts | 5 ++--- .../src/singletons/async_local_storage/index.ts | 10 ++++++---- langchain-core/src/tools/index.ts | 5 ++--- langchain-core/src/utils/stream.ts | 9 ++------- 5 files changed, 14 insertions(+), 20 deletions(-) diff --git a/langchain-core/src/runnables/base.ts b/langchain-core/src/runnables/base.ts index d2ba10b6f6bd..b3583bed58c2 100644 --- 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( (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 index 4d7ead6efa60..43ab57dbeaff 100644 --- 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( ): IterableIterator { while (true) { const { value, done } = AsyncLocalStorageProviderSingleton.runWithConfig( - pickRunnableConfigKeys(context), + context, iter.next.bind(iter), true ); @@ -57,7 +56,7 @@ export async function* consumeAsyncIterableInContext( 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 index f89c9fbce50a..4f60edeb2b24 100644 --- 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 index 8ce02d28c935..348e85103904 100644 --- 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 index cd3e592806be..19c974a4ea8b 100644 --- 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 | 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 | undefined - ), + this.config, this.signal ? async () => { return raceWithSignal(this.generator.next(...args), this.signal);