fix(js): LangChain / traceable handoff

js/.gitignore

@@ -51,6 +51,10 @@ Chinook_Sqlite.sql
 /schemas.js
 /schemas.d.ts
 /schemas.d.cts
+/langchain.cjs
+/langchain.js
+/langchain.d.ts
+/langchain.d.cts
 /wrappers.cjs
 /wrappers.js
 /wrappers.d.ts
@@ -59,6 +63,10 @@ Chinook_Sqlite.sql
 /wrappers/openai.js
 /wrappers/openai.d.ts
 /wrappers/openai.d.cts
+/singletons/traceable.cjs
+/singletons/traceable.js
+/singletons/traceable.d.ts
+/singletons/traceable.d.cts
 /index.cjs
 /index.js
 /index.d.ts

js/package.json

@@ -25,6 +25,10 @@
     "schemas.js",
     "schemas.d.ts",
     "schemas.d.cts",
+    "langchain.cjs",
+    "langchain.js",
+    "langchain.d.ts",
+    "langchain.d.cts",
     "wrappers.cjs",
     "wrappers.js",
     "wrappers.d.ts",
@@ -33,6 +37,10 @@
     "wrappers/openai.js",
     "wrappers/openai.d.ts",
     "wrappers/openai.d.cts",
+    "singletons/traceable.cjs",
+    "singletons/traceable.js",
+    "singletons/traceable.d.ts",
+    "singletons/traceable.d.cts",
     "index.cjs",
     "index.js",
     "index.d.ts",
@@ -109,11 +117,15 @@
     "typescript": "^5.4.5"
   },
   "peerDependencies": {
-    "openai": "*"
+    "openai": "*",
+    "@langchain/core": "*"
   },
   "peerDependenciesMeta": {
     "openai": {
       "optional": true
+    },
+    "@langchain/core": {
+      "optional": true
     }
   },
   "lint-staged": {
@@ -177,6 +189,15 @@
       "import": "./schemas.js",
       "require": "./schemas.cjs"
     },
+    "./langchain": {
+      "types": {
+        "import": "./langchain.d.ts",
+        "require": "./langchain.d.cts",
+        "default": "./langchain.d.ts"
+      },
+      "import": "./langchain.js",
+      "require": "./langchain.cjs"
+    },
     "./wrappers": {
       "types": {
         "import": "./wrappers.d.ts",
@@ -195,6 +216,15 @@
       "import": "./wrappers/openai.js",
       "require": "./wrappers/openai.cjs"
     },
+    "./singletons/traceable": {
+      "types": {
+        "import": "./singletons/traceable.d.ts",
+        "require": "./singletons/traceable.d.cts",
+        "default": "./singletons/traceable.d.ts"
+      },
+      "import": "./singletons/traceable.js",
+      "require": "./singletons/traceable.cjs"
+    },
     "./package.json": "./package.json"
   }
-}
+}

js/scripts/create-entrypoints.js

@@ -12,8 +12,10 @@ const entrypoints = {
   traceable: "traceable",
   evaluation: "evaluation/index",
   schemas: "schemas",
+  langchain: "langchain",
   wrappers: "wrappers/index",
   "wrappers/openai": "wrappers/openai",
+  "singletons/traceable": "singletons/traceable",
 };
 const updateJsonFile = (relativePath, updateFunction) => {
   const contents = fs.readFileSync(relativePath).toString();

js/src/env.ts

@@ -0,0 +1,14 @@
+import { getEnvironmentVariable } from "./utils/env.js";
+
+export const isTracingEnabled = (tracingEnabled?: boolean): boolean => {
+  if (tracingEnabled !== undefined) {
+    return tracingEnabled;
+  }
+  const envVars = [
+    "LANGSMITH_TRACING_V2",
+    "LANGCHAIN_TRACING_V2",
+    "LANGSMITH_TRACING",
+    "LANGCHAIN_TRACING",
+  ];
+  return !!envVars.find((envVar) => getEnvironmentVariable(envVar) === "true");
+};

js/src/langchain.ts

@@ -0,0 +1,122 @@
+import { CallbackManager } from "@langchain/core/callbacks/manager";
+import { LangChainTracer } from "@langchain/core/tracers/tracer_langchain";
+import { Runnable, RunnableConfig } from "@langchain/core/runnables";
+
+import { RunTree } from "./run_trees.js";
+import { Run } from "./schemas.js";
+import {
+  TraceableFunction,
+  getCurrentRunTree,
+  isTraceableFunction,
+} from "./traceable.js";
+
+/**
+ * Converts the current run tree active within a traceable-wrapped function
+ * into a LangChain compatible callback manager. This is useful to handoff tracing
+ * from LangSmith to LangChain Runnables and LLMs.
+ *
+ * @param {RunTree | undefined} currentRunTree Current RunTree from within a traceable-wrapped function. If not provided, the current run tree will be inferred from AsyncLocalStorage.
+ * @returns {CallbackManager | undefined} Callback manager used by LangChain Runnable objects.
+ */
+export async function getLangchainCallbacks(
+  currentRunTree?: RunTree | undefined
+) {
+  const runTree: RunTree | undefined = currentRunTree ?? getCurrentRunTree();
+  if (!runTree) return undefined;
+
+  // TODO: CallbackManager.configure() is only async due to LangChainTracer
+  // factory being unnecessarily async.
+  let callbacks = await CallbackManager.configure();
+  if (!callbacks && runTree.tracingEnabled) {
+    callbacks = new CallbackManager();
+  }
+
+  let langChainTracer = callbacks?.handlers.find(
+    (handler): handler is LangChainTracer =>
+      handler?.name === "langchain_tracer"
+  );
+
+  if (!langChainTracer && runTree.tracingEnabled) {
+    langChainTracer = new LangChainTracer();
+    callbacks?.addHandler(langChainTracer);
+  }
+
+  const runMap = new Map<string, Run>();
+
+  // find upward root run
+  let rootRun = runTree;
+  const rootVisited = new Set<string>();
+  while (rootRun.parent_run) {
+    if (rootVisited.has(rootRun.id)) break;
+    rootVisited.add(rootRun.id);
+    rootRun = rootRun.parent_run;
+  }
+
+  const queue = [rootRun];
+  const visited = new Set<string>();
+
+  while (queue.length > 0) {
+    const current = queue.shift();
+    if (!current || visited.has(current.id)) continue;
+    visited.add(current.id);
+
+    runMap.set(current.id, current);
+    if (current.child_runs) {
+      queue.push(...current.child_runs);
+    }
+  }
+
+  if (callbacks != null) {
+    Object.assign(callbacks, { _parentRunId: runTree.id });
+  }
+
+  if (langChainTracer != null) {
+    Object.assign(langChainTracer, {
+      runMap,
+      client: runTree.client,
+      projectName: runTree.project_name || langChainTracer.projectName,
+      exampleId: runTree.reference_example_id || langChainTracer.exampleId,
+    });
+  }
+
+  return callbacks;
+}
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type AnyTraceableFunction = TraceableFunction<(...any: any[]) => any>;
+
+/**
+ * RunnableTraceable is a Runnable that wraps a traceable function.
+ * This allows adding Langsmith traced functions into LangChain sequences.
+ */
+export class RunnableTraceable<RunInput, RunOutput> extends Runnable<
+  RunInput,
+  RunOutput
+> {
+  lc_serializable = false;
+
+  lc_namespace = ["langchain_core", "runnables"];
+
+  protected func: AnyTraceableFunction;
+
+  constructor(fields: { func: AnyTraceableFunction }) {
+    super(fields);
+
+    if (!isTraceableFunction(fields.func)) {
+      throw new Error(
+        "RunnableTraceable requires a function that is wrapped in traceable higher-order function"
+      );
+    }
+
+    this.func = fields.func;
+  }
+
+  async invoke(input: RunInput, options?: Partial<RunnableConfig>) {
+    const [config] = this._getOptionsList(options ?? {}, 1);
+    return (await this.func(config, input)) as RunOutput;
+  }
+
+  static from(func: AnyTraceableFunction) {
+    return new RunnableTraceable({ func });
+  }
+}

js/src/run_trees.ts

@@ -6,6 +6,7 @@ import {
   getRuntimeEnvironment,
 } from "./utils/env.js";
 import { Client } from "./client.js";
+import { isTracingEnabled } from "./env.js";
 
 const warnedMessages: Record<string, boolean> = {};
 
@@ -93,6 +94,7 @@ interface LangChainTracerLike extends TracerLike {
   name: "langchain_tracer";
   projectName: string;
   getRun?: (id: string) => RunTree | undefined;
+  client: Client;
 }
 
 export class RunTree implements BaseRun {
@@ -157,44 +159,57 @@ export class RunTree implements BaseRun {
   }
 
   static fromRunnableConfig(
-    config: RunnableConfigLike,
+    parentConfig: RunnableConfigLike,
     props: {
       name: string;
       tags?: string[];
       metadata?: KVMap;
     }
   ): RunTree {
     // We only handle the callback manager case for now
-    const callbackManager = config?.callbacks as
+    const callbackManager = parentConfig?.callbacks as
       | CallbackManagerLike
       | undefined;
     let parentRun: RunTree | undefined;
     let projectName: string | undefined;
+    let client: Client | undefined;
+
+    let tracingEnabled = isTracingEnabled();
+
     if (callbackManager) {
       const parentRunId = callbackManager?.getParentRunId?.() ?? "";
       const langChainTracer = callbackManager?.handlers?.find(
         (handler: TracerLike) => handler?.name == "langchain_tracer"
       ) as LangChainTracerLike | undefined;
+
       parentRun = langChainTracer?.getRun?.(parentRunId);
       projectName = langChainTracer?.projectName;
+      client = langChainTracer?.client;
+      tracingEnabled = tracingEnabled || !!langChainTracer;
     }
-    const dedupedTags = [
-      ...new Set((parentRun?.tags ?? []).concat(config?.tags ?? [])),
-    ];
-    const dedupedMetadata = {
-      ...parentRun?.extra?.metadata,
-      ...config?.metadata,
-    };
-    const rt = new RunTree({
-      name: props?.name ?? "<lambda>",
-      parent_run: parentRun,
-      tags: dedupedTags,
+
+    const parentRunTree = new RunTree({
+      name: parentRun?.name ?? "<parent>",
+      id: parentRun?.id,
+      client,
+      tracingEnabled,
+      project_name: projectName,
+      tags: [
+        ...new Set((parentRun?.tags ?? []).concat(parentConfig?.tags ?? [])),
+      ],
       extra: {
-        metadata: dedupedMetadata,
+        metadata: {
+          ...parentRun?.extra?.metadata,
+          ...parentConfig?.metadata,
+        },
       },
-      project_name: projectName,
     });
-    return rt;
+
+    return parentRunTree.createChild({
+      name: props?.name ?? "<lambda>",
+      tags: props.tags,
+      metadata: props.metadata,
+    });
   }
 
   private static getDefaultConfig(): object {