core[patch]: RunnableWithFallbacks tracing fix, refactor (#5995)

* RunnableWithFallbacks tracing fix, refactor * Add comment * Format
langchain-ai · Jul 7, 2024 · 484191a · 484191a
1 parent afe5572
commit 484191a
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 29 deletions.
diff --git a/langchain-core/src/callbacks/manager.ts b/langchain-core/src/callbacks/manager.ts
@@ -9,7 +9,6 @@ import {
   NewTokenIndices,
 } from "./base.js";
 import { ConsoleCallbackHandler } from "../tracers/console.js";
-import { getTracingV2CallbackHandler } from "../tracers/initialize.js";
 import { type BaseMessage } from "../messages/base.js";
 import { getBufferString } from "../messages/utils.js";
 import { getEnvironmentVariable } from "../utils/env.js";
@@ -20,6 +19,7 @@ import {
 import { consumeCallback } from "./promises.js";
 import { Serialized } from "../load/serializable.js";
 import type { DocumentInterface } from "../documents/document.js";
+import { isTracingEnabled } from "../utils/callbacks.js";
 
 if (
   /* #__PURE__ */ getEnvironmentVariable("LANGCHAIN_TRACING_V2") === "true" &&
@@ -111,7 +111,7 @@ export abstract class BaseCallbackManager {
 /**
  * Base class for run manager in LangChain.
  */
-class BaseRunManager {
+export class BaseRunManager {
   constructor(
     public readonly runId: string,
     public readonly handlers: BaseCallbackHandler[],
@@ -123,6 +123,10 @@ class BaseRunManager {
     protected readonly _parentRunId?: string
   ) {}
 
+  get parentRunId() {
+    return this._parentRunId;
+  }
+
   async handleText(text: string): Promise<void> {
     await Promise.all(
       this.handlers.map((handler) =>
@@ -962,6 +966,27 @@ export class CallbackManager
     localMetadata?: Record<string, unknown>,
     options?: CallbackManagerOptions
   ): Promise<CallbackManager | undefined> {
+    return this._configureSync(
+      inheritableHandlers,
+      localHandlers,
+      inheritableTags,
+      localTags,
+      inheritableMetadata,
+      localMetadata,
+      options
+    );
+  }
+
+  // TODO: Deprecate async method in favor of this one.
+  static _configureSync(
+    inheritableHandlers?: Callbacks,
+    localHandlers?: Callbacks,
+    inheritableTags?: string[],
+    localTags?: string[],
+    inheritableMetadata?: Record<string, unknown>,
+    localMetadata?: Record<string, unknown>,
+    options?: CallbackManagerOptions
+  ) {
     let callbackManager: CallbackManager | undefined;
     if (inheritableHandlers || localHandlers) {
       if (Array.isArray(inheritableHandlers) || !inheritableHandlers) {
@@ -984,9 +1009,7 @@ export class CallbackManager
     const verboseEnabled =
       getEnvironmentVariable("LANGCHAIN_VERBOSE") === "true" ||
       options?.verbose;
-    const tracingV2Enabled =
-      getEnvironmentVariable("LANGCHAIN_TRACING_V2") === "true" ||
-      getEnvironmentVariable("LANGSMITH_TRACING") === "true";
+    const tracingV2Enabled = isTracingEnabled();
 
     const tracingEnabled =
       tracingV2Enabled ||
@@ -1011,7 +1034,7 @@ export class CallbackManager
         )
       ) {
         if (tracingV2Enabled) {
-          const tracerV2 = await getTracingV2CallbackHandler();
+          const tracerV2 = new LangChainTracer();
           callbackManager.addHandler(tracerV2, true);
 
           // handoff between langchain and langsmith/traceable

diff --git a/langchain-core/src/runnables/base.ts b/langchain-core/src/runnables/base.ts
@@ -7,10 +7,7 @@ import {
   isTraceableFunction,
 } from "langsmith/singletons/traceable";
 import type { RunnableInterface, RunnableBatchOptions } from "./types.js";
-import {
-  CallbackManager,
-  CallbackManagerForChainRun,
-} from "../callbacks/manager.js";
+import { CallbackManagerForChainRun } from "../callbacks/manager.js";
 import {
   LogStreamCallbackHandler,
   LogStreamCallbackHandlerInput,
@@ -2433,29 +2430,24 @@ export class RunnableWithFallbacks<RunInput, RunOutput> extends Runnable<
     input: RunInput,
     options?: Partial<RunnableConfig>
   ): Promise<RunOutput> {
-    const callbackManager_ = await CallbackManager.configure(
-      options?.callbacks,
-      undefined,
-      options?.tags,
-      undefined,
-      options?.metadata
-    );
-    const { runId, ...otherOptions } = options ?? {};
+    const config = ensureConfig(options);
+    const callbackManager_ = await getCallbackManagerForConfig(options);
+    const { runId, ...otherConfigFields } = config;
     const runManager = await callbackManager_?.handleChainStart(
       this.toJSON(),
       _coerceToDict(input, "input"),
       runId,
       undefined,
       undefined,
       undefined,
-      otherOptions?.runName
+      otherConfigFields?.runName
     );
     let firstError;
     for (const runnable of this.runnables()) {
       try {
         const output = await runnable.invoke(
           input,
-          patchConfig(otherOptions, { callbacks: runManager?.getChild() })
+          patchConfig(otherConfigFields, { callbacks: runManager?.getChild() })
         );
         await runManager?.handleChainEnd(_coerceToDict(output, "output"));
         return output;
@@ -2500,15 +2492,7 @@ export class RunnableWithFallbacks<RunInput, RunOutput> extends Runnable<
     }
     const configList = this._getOptionsList(options ?? {}, inputs.length);
     const callbackManagers = await Promise.all(
-      configList.map((config) =>
-        CallbackManager.configure(
-          config?.callbacks,
-          undefined,
-          config?.tags,
-          undefined,
-          config?.metadata
-        )
-      )
+      configList.map((config) => getCallbackManagerForConfig(config))
     );
     const runManagers = await Promise.all(
       callbackManagers.map(async (callbackManager, i) => {

diff --git a/langchain-core/src/tracers/initialize.ts b/langchain-core/src/tracers/initialize.ts
@@ -23,6 +23,8 @@ export async function getTracingCallbackHandler(
 }
 
 /**
+ * @deprecated Instantiate directly using the LangChainTracer constructor.
+ *
  * Function that returns an instance of `LangChainTracer`. It does not
  * load any session data.
  * @returns An instance of `LangChainTracer`.

diff --git a/langchain-core/src/utils/callbacks.ts b/langchain-core/src/utils/callbacks.ts
@@ -0,0 +1,14 @@
+import { getEnvironmentVariable } from "./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");
+};