feat(core): add registerConfigureHook (#7330)

Co-authored-by: jacoblee93 <[email protected]>

langchain-core/src/callbacks/base.ts

@@ -394,3 +394,13 @@ export abstract class BaseCallbackHandler
     return new Handler();
   }
 }
+
+export const isBaseCallbackHandler = (x: unknown) => {
+  const callbackHandler = x as BaseCallbackHandler;
+  return (
+    callbackHandler !== undefined &&
+    typeof callbackHandler.copy === "function" &&
+    typeof callbackHandler.name === "string" &&
+    typeof callbackHandler.awaitHandlers === "boolean"
+  );
+};

langchain-core/src/callbacks/manager.ts

@@ -6,6 +6,7 @@ import {
   BaseCallbackHandler,
   CallbackHandlerMethods,
   HandleLLMNewTokenCallbackFields,
+  isBaseCallbackHandler,
   NewTokenIndices,
 } from "./base.js";
 import { ConsoleCallbackHandler } from "../tracers/console.js";
@@ -21,6 +22,10 @@ import { Serialized } from "../load/serializable.js";
 import type { DocumentInterface } from "../documents/document.js";
 import { isTracingEnabled } from "../utils/callbacks.js";
 import { isBaseTracer } from "../tracers/base.js";
+import {
+  getContextVariable,
+  _getConfigureHooks,
+} from "../singletons/async_local_storage/context.js";
 
 type BaseCallbackManagerMethods = {
   [K in keyof CallbackHandlerMethods]?: (
@@ -1252,6 +1257,31 @@ export class CallbackManager
         callbackManager.addMetadata(localMetadata ?? {}, false);
       }
     }
+
+    for (const {
+      contextVar,
+      inheritable = true,
+      handlerClass,
+      envVar,
+    } of _getConfigureHooks()) {
+      const createIfNotInContext =
+        envVar && getEnvironmentVariable(envVar) === "true" && handlerClass;
+      let handler: BaseCallbackHandler | undefined;
+      const contextVarValue =
+        contextVar !== undefined ? getContextVariable(contextVar) : undefined;
+      if (contextVarValue && isBaseCallbackHandler(contextVarValue)) {
+        handler = contextVarValue;
+      } else if (createIfNotInContext) {
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        handler = new (handlerClass as any)({});
+      }
+      if (handler !== undefined) {
+        if (!callbackManager?.handlers.some((h) => h.name === handler!.name)) {
+          callbackManager?.addHandler(handler, inheritable);
+        }
+      }
+    }
+
     return callbackManager;
   }
 }

langchain-core/src/callbacks/tests/manager.test.ts

@@ -0,0 +1,73 @@
+/* eslint-disable no-process-env */
+import { expect, test, beforeAll, afterEach } from "@jest/globals";
+
+import { setContextVariable, registerConfigureHook } from "../../context.js";
+import { BaseCallbackHandler } from "../base.js";
+import { CallbackManager } from "../manager.js";
+
+class TestHandler extends BaseCallbackHandler {
+  name = "TestHandler";
+}
+
+const handlerInstance = new TestHandler();
+
+beforeAll(() => {
+  process.env.LANGCHAIN_TRACING_V2 = "false";
+  process.env.LANGSMITH_TRACING_V2 = "false";
+  process.env.__TEST_VAR = "false";
+});
+
+afterEach(() => {
+  setContextVariable("my_test_handler", undefined);
+});
+
+test("configure with empty array", async () => {
+  const manager = CallbackManager.configure([]);
+  expect(manager?.handlers.length).toBe(0);
+});
+
+test("configure with one handler", async () => {
+  const manager = CallbackManager.configure([handlerInstance]);
+  expect(manager?.handlers[0]).toBe(handlerInstance);
+});
+
+test("registerConfigureHook with contextVar", async () => {
+  setContextVariable("my_test_handler", handlerInstance);
+  registerConfigureHook({
+    contextVar: "my_test_handler",
+  });
+  const manager = CallbackManager.configure([]);
+  expect(manager?.handlers[0]).toBe(handlerInstance);
+});
+
+test("registerConfigureHook with env", async () => {
+  process.env.__TEST_VAR = "true";
+  registerConfigureHook({
+    handlerClass: TestHandler,
+    envVar: "__TEST_VAR",
+  });
+  const manager = CallbackManager.configure([]);
+  expect(manager?.handlers[0].name).toBe("TestHandler");
+});
+
+test("registerConfigureHook doesn't add with env false", async () => {
+  process.env.__TEST_VAR = "false";
+  registerConfigureHook({
+    handlerClass: TestHandler,
+    envVar: "__TEST_VAR",
+  });
+  const manager = CallbackManager.configure([]);
+  expect(manager?.handlers.length).toBe(0);
+});
+
+test("registerConfigureHook avoids multiple", async () => {
+  process.env.__TEST_VAR = "true";
+  registerConfigureHook({
+    contextVar: "my_test_handler",
+    handlerClass: TestHandler,
+    envVar: "__TEST_VAR",
+  });
+  const manager = CallbackManager.configure([handlerInstance]);
+  expect(manager?.handlers[0]).toBe(handlerInstance);
+  expect(manager?.handlers[1]).toBe(undefined);
+});

langchain-core/src/context.ts

@@ -1,131 +1,29 @@
 /* __LC_ALLOW_ENTRYPOINT_SIDE_EFFECTS__ */
+
+/**
+ * This file exists as a convenient public entrypoint for functionality
+ * related to context variables.
+ *
+ * Because it automatically initializes AsyncLocalStorage, internal
+ * functionality SHOULD NEVER import from this file outside of tests.
+ */
+
 import { AsyncLocalStorage } from "node:async_hooks";
-import { RunTree } from "langsmith";
-import { isRunTree } from "langsmith/run_trees";
+import { AsyncLocalStorageProviderSingleton } from "./singletons/index.js";
 import {
-  _CONTEXT_VARIABLES_KEY,
-  AsyncLocalStorageProviderSingleton,
-} from "./singletons/index.js";
+  getContextVariable,
+  setContextVariable,
+  type ConfigureHook,
+  registerConfigureHook,
+} from "./singletons/async_local_storage/context.js";
 
 AsyncLocalStorageProviderSingleton.initializeGlobalInstance(
   new AsyncLocalStorage()
 );
 
-/**
- * Set a context variable. Context variables are scoped to any
- * child runnables called by the current runnable, or globally if set outside
- * of any runnable.
- *
- * @remarks
- * This function is only supported in environments that support AsyncLocalStorage,
- * including Node.js, Deno, and Cloudflare Workers.
- *
- * @example
- * ```ts
- * import { RunnableLambda } from "@langchain/core/runnables";
- * import {
- *   getContextVariable,
- *   setContextVariable
- * } from "@langchain/core/context";
- *
- * const nested = RunnableLambda.from(() => {
- *   // "bar" because it was set by a parent
- *   console.log(getContextVariable("foo"));
- *
- *   // Override to "baz", but only for child runnables
- *   setContextVariable("foo", "baz");
- *
- *   // Now "baz", but only for child runnables
- *   return getContextVariable("foo");
- * });
- *
- * const runnable = RunnableLambda.from(async () => {
- *   // Set a context variable named "foo"
- *   setContextVariable("foo", "bar");
- *
- *   const res = await nested.invoke({});
- *
- *   // Still "bar" since child changes do not affect parents
- *   console.log(getContextVariable("foo"));
- *
- *   return res;
- * });
- *
- * // undefined, because context variable has not been set yet
- * console.log(getContextVariable("foo"));
- *
- * // Final return value is "baz"
- * const result = await runnable.invoke({});
- * ```
- *
- * @param name The name of the context variable.
- * @param value The value to set.
- */
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-export function setContextVariable<T>(name: PropertyKey, value: T): void {
-  const runTree = AsyncLocalStorageProviderSingleton.getInstance().getStore();
-  const contextVars = { ...runTree?.[_CONTEXT_VARIABLES_KEY] };
-  contextVars[name] = value;
-  let newValue = {};
-  if (isRunTree(runTree)) {
-    newValue = new RunTree(runTree);
-  }
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  (newValue as any)[_CONTEXT_VARIABLES_KEY] = contextVars;
-  AsyncLocalStorageProviderSingleton.getInstance().enterWith(newValue);
-}
-
-/**
- * Get the value of a previously set context variable. Context variables
- * are scoped to any child runnables called by the current runnable,
- * or globally if set outside of any runnable.
- *
- * @remarks
- * This function is only supported in environments that support AsyncLocalStorage,
- * including Node.js, Deno, and Cloudflare Workers.
- *
- * @example
- * ```ts
- * import { RunnableLambda } from "@langchain/core/runnables";
- * import {
- *   getContextVariable,
- *   setContextVariable
- * } from "@langchain/core/context";
- *
- * const nested = RunnableLambda.from(() => {
- *   // "bar" because it was set by a parent
- *   console.log(getContextVariable("foo"));
- *
- *   // Override to "baz", but only for child runnables
- *   setContextVariable("foo", "baz");
- *
- *   // Now "baz", but only for child runnables
- *   return getContextVariable("foo");
- * });
- *
- * const runnable = RunnableLambda.from(async () => {
- *   // Set a context variable named "foo"
- *   setContextVariable("foo", "bar");
- *
- *   const res = await nested.invoke({});
- *
- *   // Still "bar" since child changes do not affect parents
- *   console.log(getContextVariable("foo"));
- *
- *   return res;
- * });
- *
- * // undefined, because context variable has not been set yet
- * console.log(getContextVariable("foo"));
- *
- * // Final return value is "baz"
- * const result = await runnable.invoke({});
- * ```
- *
- * @param name The name of the context variable.
- */
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-export function getContextVariable<T = any>(name: PropertyKey): T | undefined {
-  const runTree = AsyncLocalStorageProviderSingleton.getInstance().getStore();
-  return runTree?.[_CONTEXT_VARIABLES_KEY]?.[name];
-}
+export {
+  getContextVariable,
+  setContextVariable,
+  registerConfigureHook,
+  type ConfigureHook,
+};