From c5728a8d41446120baa969a98068ef4a927da188 Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Tue, 13 Feb 2024 16:52:34 -0800
Subject: [PATCH] Pass vals in update

---
 .../tests/langchain_tracer.int.test.ts        | 35 ++++++++++++++++++-
 .../src/tracers/tracer_langchain.ts           | 14 +++++++-
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/langchain-core/src/tracers/tests/langchain_tracer.int.test.ts b/langchain-core/src/tracers/tests/langchain_tracer.int.test.ts
index c377355fe8a..51750b3b7f5 100644
--- a/langchain-core/src/tracers/tests/langchain_tracer.int.test.ts
+++ b/langchain-core/src/tracers/tests/langchain_tracer.int.test.ts
@@ -2,7 +2,7 @@
 import * as uuid from "uuid";
 import { test } from "@jest/globals";
 
-import { LangChainTracer } from "../tracer_langchain.js";
+import { LangChainTracer, Run } from "../tracer_langchain.js";
 import { Serialized } from "../../load/serializable.js";
 import { HumanMessage } from "../../messages/index.js";
 
@@ -41,3 +41,36 @@ test("LangChain V2 tracer does not throw errors for its methods", async () => {
   await tracer.handleLLMStart(serialized, ["test"], llmRunId3);
   await tracer.handleLLMEnd({ generations: [[]] }, llmRunId3);
 });
+
+class FakeTracer extends LangChainTracer {
+  createOperations: { [id: string]: Run } = {};
+  updateOperations: { [id: string]: Run } = {};
+
+  async onRunCreate(run: Run): Promise<void> {
+    this.createOperations[run.id] = run;
+  }
+
+  async onRunUpdate(run: Run): Promise<void> {
+    this.updateOperations[run.id] = run;
+  }
+}
+
+test("LangChain V2 tracer creates and updates runs with trace_id and dotted_order", async () => {
+  const tracer = new FakeTracer({
+    projectName: `JS Int Test - ${uuid.v4()}`,
+  });
+  const chainRunId = uuid.v4();
+  const llmRunId = uuid.v4();
+  await tracer.handleChainStart(serialized, { foo: "bar" }, chainRunId);
+
+  await tracer.handleLLMStart(serialized, ["test"], llmRunId, chainRunId);
+  await tracer.handleLLMEnd({ generations: [[]] }, llmRunId);
+  await tracer.handleChainEnd({ foo: "bar" }, chainRunId);
+
+  expect(tracer.createOperations[chainRunId].trace_id).toBeDefined();
+  expect(tracer.createOperations[chainRunId].trace_id).toEqual(chainRunId);
+  expect(tracer.createOperations[chainRunId].dotted_order).toBeDefined();
+  expect(tracer.updateOperations[llmRunId].trace_id).toBeDefined();
+  expect(tracer.updateOperations[llmRunId].trace_id).toEqual(chainRunId);
+  expect(tracer.updateOperations[llmRunId].dotted_order).toBeDefined();
+});
diff --git a/langchain-core/src/tracers/tracer_langchain.ts b/langchain-core/src/tracers/tracer_langchain.ts
index 3c656f318e3..a1bbc7f24a4 100644
--- a/langchain-core/src/tracers/tracer_langchain.ts
+++ b/langchain-core/src/tracers/tracer_langchain.ts
@@ -13,11 +13,20 @@ export interface Run extends BaseRun {
   id: string;
   child_runs: this[];
   child_execution_order: number;
+  dotted_order?: string;
+  trace_id?: string;
+}
+
+export interface RunCreate2 extends RunCreate {
+  trace_id?: string;
+  dotted_order?: string;
 }
 
 export interface RunUpdate extends BaseRunUpdate {
   events: BaseRun["events"];
   inputs: KVMap;
+  trace_id?: string;
+  dotted_order?: string;
 }
 
 export interface LangChainTracerFields extends BaseCallbackHandlerInput {
@@ -69,7 +78,7 @@ export class LangChainTracer
   protected async persistRun(_run: Run): Promise<void> {}
 
   async onRunCreate(run: Run): Promise<void> {
-    const persistedRun: RunCreate = await this._convertToCreate(
+    const persistedRun: RunCreate2 = await this._convertToCreate(
       run,
       this.exampleId
     );
@@ -83,6 +92,9 @@ export class LangChainTracer
       outputs: run.outputs,
       events: run.events,
       inputs: run.inputs,
+      trace_id: run.trace_id,
+      dotted_order: run.dotted_order,
+      parent_run_id: run.parent_run_id,
     };
     await this.client.updateRun(run.id, runUpdate);
   }