feat(core): Adds support for context variables (#6967)

langchain-core/.gitignore

@@ -30,6 +30,10 @@ chat_history.cjs
 chat_history.js
 chat_history.d.ts
 chat_history.d.cts
+context.cjs
+context.js
+context.d.ts
+context.d.cts
 documents.cjs
 documents.js
 documents.d.ts

langchain-core/langchain.config.js

@@ -20,6 +20,7 @@ export const config = {
     "callbacks/manager": "callbacks/manager",
     "callbacks/promises": "callbacks/promises",
     chat_history: "chat_history",
+    context: "context",
     documents: "documents/index",
     "document_loaders/base": "document_loaders/base",
     "document_loaders/langsmith": "document_loaders/langsmith",

langchain-core/package.json

@@ -37,7 +37,7 @@
     "camelcase": "6",
     "decamelize": "1.2.0",
     "js-tiktoken": "^1.0.12",
-    "langsmith": "^0.1.56",
+    "langsmith": "^0.1.65",
     "mustache": "^4.2.0",
     "p-queue": "^6.6.2",
     "p-retry": "4",
@@ -160,6 +160,15 @@
       "import": "./chat_history.js",
       "require": "./chat_history.cjs"
     },
+    "./context": {
+      "types": {
+        "import": "./context.d.ts",
+        "require": "./context.d.cts",
+        "default": "./context.d.ts"
+      },
+      "import": "./context.js",
+      "require": "./context.cjs"
+    },
     "./documents": {
       "types": {
         "import": "./documents.d.ts",
@@ -646,6 +655,10 @@
     "chat_history.js",
     "chat_history.d.ts",
     "chat_history.d.cts",
+    "context.cjs",
+    "context.js",
+    "context.d.ts",
+    "context.d.cts",
     "documents.cjs",
     "documents.js",
     "documents.d.ts",

langchain-core/src/callbacks/dispatch/index.ts

@@ -1,10 +1,12 @@
+/* __LC_ALLOW_ENTRYPOINT_SIDE_EFFECTS__ */
+
 import { AsyncLocalStorage } from "node:async_hooks";
 import { dispatchCustomEvent as dispatchCustomEventWeb } from "./web.js";
 import { type RunnableConfig, ensureConfig } from "../../runnables/config.js";
 import { AsyncLocalStorageProviderSingleton } from "../../singletons/index.js";
 
-/* #__PURE__ */ AsyncLocalStorageProviderSingleton.initializeGlobalInstance(
-  /* #__PURE__ */ new AsyncLocalStorage()
+AsyncLocalStorageProviderSingleton.initializeGlobalInstance(
+  new AsyncLocalStorage()
 );
 
 /**

langchain-core/src/context.ts

@@ -0,0 +1,131 @@
+/* __LC_ALLOW_ENTRYPOINT_SIDE_EFFECTS__ */
+import { AsyncLocalStorage } from "node:async_hooks";
+import { RunTree } from "langsmith";
+import { isRunTree } from "langsmith/run_trees";
+import {
+  _CONTEXT_VARIABLES_KEY,
+  AsyncLocalStorageProviderSingleton,
+} from "./singletons/index.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(name: PropertyKey, value: any): 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(name: PropertyKey): any {
+  const runTree = AsyncLocalStorageProviderSingleton.getInstance().getStore();
+  return runTree?.[_CONTEXT_VARIABLES_KEY]?.[name];
+}

langchain-core/src/singletons/index.ts

@@ -7,6 +7,8 @@ export interface AsyncLocalStorageInterface {
   getStore: () => any | undefined;
 
   run: <T>(store: any, callback: () => T) => T;
+
+  enterWith: (store: any) => void;
 }
 
 export class MockAsyncLocalStorage implements AsyncLocalStorageInterface {
@@ -17,13 +19,19 @@ export class MockAsyncLocalStorage implements AsyncLocalStorageInterface {
   run<T>(_store: any, callback: () => T): T {
     return callback();
   }
+
+  enterWith(_store: any) {
+    return undefined;
+  }
 }
 
 const mockAsyncLocalStorage = new MockAsyncLocalStorage();
 
 const TRACING_ALS_KEY = Symbol.for("ls:tracing_async_local_storage");
 const LC_CHILD_KEY = Symbol.for("lc:child_config");
 
+export const _CONTEXT_VARIABLES_KEY = Symbol.for("lc:context_variables");
+
 class AsyncLocalStorageProvider {
   getInstance(): AsyncLocalStorageInterface {
     return (globalThis as any)[TRACING_ALS_KEY] ?? mockAsyncLocalStorage;
@@ -50,6 +58,7 @@ class AsyncLocalStorageProvider {
       config?.metadata
     );
     const storage = this.getInstance();
+    const previousValue = storage.getStore();
     const parentRunId = callbackManager?.getParentRunId();
 
     const langChainTracer = callbackManager?.handlers?.find(
@@ -70,6 +79,14 @@ class AsyncLocalStorageProvider {
       runTree.extra = { ...runTree.extra, [LC_CHILD_KEY]: config };
     }
 
+    if (
+      previousValue !== undefined &&
+      previousValue[_CONTEXT_VARIABLES_KEY] !== undefined
+    ) {
+      (runTree as any)[_CONTEXT_VARIABLES_KEY] =
+        previousValue[_CONTEXT_VARIABLES_KEY];
+    }
+
     return storage.run(runTree, callback);
   }
 

langchain-core/src/tests/context.test.ts

@@ -0,0 +1,26 @@
+import { test, expect } from "@jest/globals";
+import { RunnableLambda } from "../runnables/base.js";
+import { getContextVariable, setContextVariable } from "../context.js";
+
+test("Getting and setting context variables within nested runnables", async () => {
+  const nested = RunnableLambda.from(() => {
+    expect(getContextVariable("foo")).toEqual("bar");
+    expect(getContextVariable("toplevel")).toEqual(9);
+    setContextVariable("foo", "baz");
+    return getContextVariable("foo");
+  });
+  const runnable = RunnableLambda.from(async () => {
+    setContextVariable("foo", "bar");
+    expect(getContextVariable("foo")).toEqual("bar");
+    expect(getContextVariable("toplevel")).toEqual(9);
+    const res = await nested.invoke({});
+    expect(getContextVariable("foo")).toEqual("bar");
+    return res;
+  });
+  expect(getContextVariable("foo")).toEqual(undefined);
+  setContextVariable("toplevel", 9);
+  expect(getContextVariable("toplevel")).toEqual(9);
+  const result = await runnable.invoke({});
+  expect(getContextVariable("toplevel")).toEqual(9);
+  expect(result).toEqual("baz");
+});