feat(js): Adds further required support for context vars (#1093)

js/package.json

@@ -1,6 +1,6 @@
 {
   "name": "langsmith",
-  "version": "0.1.64",
+  "version": "0.1.65",
   "description": "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform.",
   "packageManager": "[email protected]",
   "files": [

js/src/index.ts

@@ -14,4 +14,4 @@ export { RunTree, type RunTreeConfig } from "./run_trees.js";
 export { overrideFetchImplementation } from "./singletons/fetch.js";
 
 // Update using yarn bump-version
-export const __version__ = "0.1.64";
+export const __version__ = "0.1.65";

js/src/run_trees.ts

@@ -8,6 +8,7 @@ import {
 import { Client } from "./client.js";
 import { isTracingEnabled } from "./env.js";
 import { warnOnce } from "./utils/warn.js";
+import { _LC_CONTEXT_VARIABLES_KEY } from "./singletons/constants.js";
 
 function stripNonAlphanumeric(input: string) {
   return input.replace(/[-:.]/g, "");
@@ -172,7 +173,12 @@ export class RunTree implements BaseRun {
   execution_order: number;
   child_execution_order: number;
 
-  constructor(originalConfig: RunTreeConfig) {
+  constructor(originalConfig: RunTreeConfig | RunTree) {
+    // If you pass in a run tree directly, return a shallow clone
+    if (isRunTree(originalConfig)) {
+      Object.assign(this, { ...originalConfig });
+      return;
+    }
     const defaultConfig = RunTree.getDefaultConfig();
     const { metadata, ...config } = originalConfig;
     const client = config.client ?? RunTree.getSharedClient();
@@ -248,6 +254,13 @@ export class RunTree implements BaseRun {
       child_execution_order: child_execution_order,
     });
 
+    // Copy context vars over into the new run tree.
+    if (_LC_CONTEXT_VARIABLES_KEY in this) {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      (child as any)[_LC_CONTEXT_VARIABLES_KEY] =
+        this[_LC_CONTEXT_VARIABLES_KEY];
+    }
+
     type ExtraWithSymbol = Record<string | symbol, unknown>;
     const LC_CHILD = Symbol.for("lc:child_config");
 

js/src/singletons/constants.ts

@@ -0,0 +1 @@
+export const _LC_CONTEXT_VARIABLES_KEY = Symbol.for("lc:context_variables");

js/src/tests/traceable.test.ts

@@ -1,9 +1,11 @@
 import { jest } from "@jest/globals";
 import { RunTree, RunTreeConfig } from "../run_trees.js";
+import { _LC_CONTEXT_VARIABLES_KEY } from "../singletons/constants.js";
 import { ROOT, traceable, withRunTree } from "../traceable.js";
 import { getAssumedTreeFromCalls } from "./utils/tree.js";
 import { mockClient } from "./utils/mock_client.js";
 import { Client, overrideFetchImplementation } from "../index.js";
+import { AsyncLocalStorageProviderSingleton } from "../singletons/traceable.js";
 
 test("basic traceable implementation", async () => {
   const { client, callSpy } = mockClient();
@@ -103,6 +105,80 @@ test("nested traceable implementation", async () => {
   });
 });
 
+test("nested traceable passes through LangChain context vars", (done) => {
+  const alsInstance = AsyncLocalStorageProviderSingleton.getInstance();
+
+  alsInstance.run(
+    {
+      [_LC_CONTEXT_VARIABLES_KEY]: { foo: "bar" },
+    } as any,
+    // eslint-disable-next-line @typescript-eslint/no-misused-promises
+    async () => {
+      try {
+        expect(
+          (alsInstance.getStore() as any)?.[_LC_CONTEXT_VARIABLES_KEY]?.foo
+        ).toEqual("bar");
+        const { client, callSpy } = mockClient();
+
+        const llm = traceable(async function llm(input: string) {
+          expect(
+            (alsInstance.getStore() as any)?.[_LC_CONTEXT_VARIABLES_KEY]?.foo
+          ).toEqual("bar");
+          return input.repeat(2);
+        });
+
+        const str = traceable(async function* str(input: string) {
+          const response = input.split("").reverse();
+          for (const char of response) {
+            yield char;
+          }
+          expect(
+            (alsInstance.getStore() as any)?.[_LC_CONTEXT_VARIABLES_KEY]?.foo
+          ).toEqual("bar");
+        });
+
+        const chain = traceable(
+          async function chain(input: string) {
+            expect(
+              (alsInstance.getStore() as any)?.[_LC_CONTEXT_VARIABLES_KEY]?.foo
+            ).toEqual("bar");
+            const question = await llm(input);
+
+            let answer = "";
+            for await (const char of str(question)) {
+              answer += char;
+            }
+
+            return { question, answer };
+          },
+          { client, tracingEnabled: true }
+        );
+
+        const result = await chain("Hello world");
+
+        expect(result).toEqual({
+          question: "Hello worldHello world",
+          answer: "dlrow olleHdlrow olleH",
+        });
+
+        expect(getAssumedTreeFromCalls(callSpy.mock.calls)).toMatchObject({
+          nodes: ["chain:0", "llm:1", "str:2"],
+          edges: [
+            ["chain:0", "llm:1"],
+            ["chain:0", "str:2"],
+          ],
+        });
+        expect(
+          (alsInstance.getStore() as any)?.[_LC_CONTEXT_VARIABLES_KEY]?.foo
+        ).toEqual("bar");
+        done();
+      } catch (e) {
+        done(e);
+      }
+    }
+  );
+});
+
 test("trace circular input and output objects", async () => {
   const { client, callSpy } = mockClient();
   const a: Record<string, any> = {};

js/src/traceable.ts

@@ -13,6 +13,7 @@ import {
   ROOT,
   AsyncLocalStorageProviderSingleton,
 } from "./singletons/traceable.js";
+import { _LC_CONTEXT_VARIABLES_KEY } from "./singletons/constants.js";
 import { TraceableFunction } from "./singletons/types.js";
 import {
   isKVMap,
@@ -422,6 +423,17 @@ export function traceable<Func extends (...args: any[]) => any>(
         processedArgs,
         config?.getInvocationParams
       );
+      // If a context var is set by LangChain outside of a traceable,
+      // it will be an object with a single property and we should copy
+      // context vars over into the new run tree.
+      if (
+        prevRunFromStore !== undefined &&
+        _LC_CONTEXT_VARIABLES_KEY in prevRunFromStore
+      ) {
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        (currentRunTree as any)[_LC_CONTEXT_VARIABLES_KEY] =
+          prevRunFromStore[_LC_CONTEXT_VARIABLES_KEY];
+      }
       return [currentRunTree, processedArgs as Inputs];
     })();