fix(langchain): emit tool observations for LangChain tool runs (#787)

Author: hassiebp

packages/langchain/src/CallbackHandler.ts

@@ -17,6 +17,8 @@ import {
   startActiveObservation,
   LangfuseGeneration,
   LangfuseSpan,
+  LangfuseTool,
+  LangfuseObservationAttributes,
   LangfuseGenerationAttributes,
   LangfuseSpanAttributes,
   propagateAttributes,
@@ -49,6 +51,8 @@ type ConstructorParams = {
   traceMetadata?: Record<string, unknown>; // added to all traces
 };
 
+type TrackedObservation = LangfuseSpan | LangfuseGeneration | LangfuseTool;
+
 export class CallbackHandler extends BaseCallbackHandler {
   name = "LangfuseCallbackHandler";
 
@@ -60,7 +64,7 @@ export class CallbackHandler extends BaseCallbackHandler {
 
   private completionStartTimes: Record<string, Date> = {};
   private promptToParentRunMap;
-  private runMap: Map<string, LangfuseSpan | LangfuseGeneration> = new Map();
+  private runMap: Map<string, TrackedObservation> = new Map();
 
   public last_trace_id: string | null = null;
 
@@ -457,6 +461,7 @@ export class CallbackHandler extends BaseCallbackHandler {
       this.logger.debug(`Tool start with ID: ${runId}`);
 
       this.startAndRegisterOtelSpan({
+        type: "tool",
         runId,
         parentRunId,
         runName: name ?? tool.id.at(-1)?.toString() ?? "Tool execution",
@@ -716,7 +721,7 @@ export class CallbackHandler extends BaseCallbackHandler {
     runName: string;
     runId: string;
     parentRunId?: string;
-    attributes: LangfuseGenerationAttributes;
+    attributes: LangfuseSpanAttributes;
     metadata?: Record<string, unknown>;
     tags?: string[];
   }): LangfuseSpan;
@@ -730,14 +735,23 @@ export class CallbackHandler extends BaseCallbackHandler {
     tags?: string[];
   }): LangfuseGeneration;
   private startAndRegisterOtelSpan(params: {
-    type?: "span" | "generation";
+    type: "tool";
     runName: string;
     runId: string;
     parentRunId?: string;
-    attributes: LangfuseGenerationAttributes;
+    attributes: LangfuseSpanAttributes;
+    metadata?: Record<string, unknown>;
+    tags?: string[];
+  }): LangfuseTool;
+  private startAndRegisterOtelSpan(params: {
+    type?: "span" | "generation" | "tool";
+    runName: string;
+    runId: string;
+    parentRunId?: string;
+    attributes: LangfuseObservationAttributes;
     metadata?: Record<string, unknown>;
     tags?: string[];
-  }): LangfuseSpan | LangfuseGeneration {
+  }): TrackedObservation {
     const { type, runName, runId, parentRunId, attributes, metadata, tags } =
       params;
 
@@ -746,37 +760,50 @@ export class CallbackHandler extends BaseCallbackHandler {
       metadata: this.joinTagsAndMetaData(tags, metadata),
       level: tags && tags.includes(LANGSMITH_HIDDEN_TAG) ? "DEBUG" : undefined,
       ...attributes,
-    } as LangfuseGenerationAttributes;
+    } as LangfuseObservationAttributes;
+
+    const parentSpanContext = parentRunId
+      ? this.runMap.get(parentRunId)?.otelSpan.spanContext()
+      : undefined;
 
     const observation =
       type === "generation"
         ? startActiveObservation(
             runName,
             (gen) => {
-              gen.update(observationAttributes);
+              gen.update(observationAttributes as LangfuseGenerationAttributes);
               return gen;
             },
             {
               asType: "generation",
-              parentSpanContext: parentRunId
-                ? this.runMap.get(parentRunId)?.otelSpan.spanContext()
-                : undefined,
+              parentSpanContext,
               endOnExit: false,
             },
           )
-        : startActiveObservation(
-            runName,
-            (span) => {
-              span.update(observationAttributes);
-              return span;
-            },
-            {
-              parentSpanContext: parentRunId
-                ? this.runMap.get(parentRunId)?.otelSpan.spanContext()
-                : undefined,
-              endOnExit: false,
-            },
-          );
+        : type === "tool"
+          ? startActiveObservation(
+              runName,
+              (tool) => {
+                tool.update(observationAttributes as LangfuseSpanAttributes);
+                return tool;
+              },
+              {
+                asType: "tool",
+                parentSpanContext,
+                endOnExit: false,
+              },
+            )
+          : startActiveObservation(
+              runName,
+              (span) => {
+                span.update(observationAttributes as LangfuseSpanAttributes);
+                return span;
+              },
+              {
+                parentSpanContext,
+                endOnExit: false,
+              },
+            );
     this.runMap.set(runId, observation);
 
     return observation;

tests/integration/langchain.integration.test.ts

@@ -0,0 +1,67 @@
+import { DynamicTool } from "@langchain/core/tools";
+import { CallbackHandler } from "@langfuse/langchain";
+import { LangfuseOtelSpanAttributes } from "@langfuse/tracing";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+
+import { SpanAssertions } from "./helpers/assertions.js";
+import {
+  setupTestEnvironment,
+  teardownTestEnvironment,
+  waitForSpanExport,
+  type TestEnvironment,
+} from "./helpers/testSetup.js";
+
+describe("LangChain callback handler integration tests", () => {
+  let testEnv: TestEnvironment;
+  let assertions: SpanAssertions;
+
+  beforeEach(async () => {
+    testEnv = await setupTestEnvironment();
+    assertions = new SpanAssertions(testEnv.mockExporter);
+  });
+
+  afterEach(async () => {
+    await teardownTestEnvironment(testEnv);
+  });
+
+  it("should mark LangChain tool runs as tool observations", async () => {
+    const calculatorTool = new DynamicTool({
+      name: "calculator",
+      description:
+        "Perform basic arithmetic operations. Input should be a mathematical expression.",
+      func: async (input: string) => {
+        const sanitizedInput = input.replace(/[^0-9+\-*/().]/g, "");
+        const result = eval(sanitizedInput);
+        return `The result is: ${result}`;
+      },
+    });
+
+    const handler = new CallbackHandler();
+
+    const result = await calculatorTool.invoke("25*4", {
+      callbacks: [handler],
+    });
+
+    expect(result).toBe("The result is: 100");
+
+    await waitForSpanExport(testEnv.mockExporter, 1);
+
+    assertions.expectSpanCount(1);
+    assertions.expectSpanWithName("calculator");
+    assertions.expectSpanAttribute(
+      "calculator",
+      LangfuseOtelSpanAttributes.OBSERVATION_TYPE,
+      "tool",
+    );
+    assertions.expectSpanAttributeContains(
+      "calculator",
+      LangfuseOtelSpanAttributes.OBSERVATION_INPUT,
+      "25*4",
+    );
+    assertions.expectSpanAttribute(
+      "calculator",
+      LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT,
+      "The result is: 100",
+    );
+  });
+});