Store events on each run (langchain-ai#1548)

* Store events on each run * Fix test
ZoVoS · Jun 13, 2023 · 8652e9c · 8652e9c
1 parent 6fc517c
commit 8652e9c
Show file tree

Hide file tree

Showing 3 changed files with 171 additions and 7 deletions.
diff --git a/langchain/src/callbacks/handlers/tracer.ts b/langchain/src/callbacks/handlers/tracer.ts
@@ -19,6 +19,11 @@ export interface Run extends BaseRun {
   // some additional fields that don't exist in sdk runs
   child_runs: this[];
   child_execution_order: number;
+  events: Array<{
+    name: string;
+    time: number;
+    kwargs?: Record<string, unknown>;
+  }>;
 }
 
 export interface AgentRun extends Run {
@@ -85,12 +90,19 @@ export abstract class BaseTracer extends BaseCallbackHandler {
     tags?: string[]
   ): Promise<void> {
     const execution_order = this._getExecutionOrder(parentRunId);
+    const start_time = Date.now();
     const run: Run = {
       id: runId,
       name: llm.id[llm.id.length - 1],
       parent_run_id: parentRunId,
-      start_time: Date.now(),
+      start_time,
       serialized: llm,
+      events: [
+        {
+          name: "start",
+          time: start_time,
+        },
+      ],
       inputs: { prompts },
       execution_order,
       child_runs: [],
@@ -113,12 +125,19 @@ export abstract class BaseTracer extends BaseCallbackHandler {
     tags?: string[]
   ): Promise<void> {
     const execution_order = this._getExecutionOrder(parentRunId);
+    const start_time = Date.now();
     const run: Run = {
       id: runId,
       name: llm.id[llm.id.length - 1],
       parent_run_id: parentRunId,
-      start_time: Date.now(),
+      start_time,
       serialized: llm,
+      events: [
+        {
+          name: "start",
+          time: start_time,
+        },
+      ],
       inputs: { messages },
       execution_order,
       child_runs: [],
@@ -139,6 +158,10 @@ export abstract class BaseTracer extends BaseCallbackHandler {
     }
     run.end_time = Date.now();
     run.outputs = output;
+    run.events.push({
+      name: "end",
+      time: run.end_time,
+    });
     await this.onLLMEnd?.(run);
     await this._endTrace(run);
   }
@@ -150,6 +173,10 @@ export abstract class BaseTracer extends BaseCallbackHandler {
     }
     run.end_time = Date.now();
     run.error = error.message;
+    run.events.push({
+      name: "error",
+      time: run.end_time,
+    });
     await this.onLLMError?.(run);
     await this._endTrace(run);
   }
@@ -162,12 +189,19 @@ export abstract class BaseTracer extends BaseCallbackHandler {
     tags?: string[]
   ): Promise<void> {
     const execution_order = this._getExecutionOrder(parentRunId);
+    const start_time = Date.now();
     const run: Run = {
       id: runId,
       name: chain.id[chain.id.length - 1],
       parent_run_id: parentRunId,
-      start_time: Date.now(),
+      start_time,
       serialized: chain,
+      events: [
+        {
+          name: "start",
+          time: start_time,
+        },
+      ],
       inputs,
       execution_order,
       child_execution_order: execution_order,
@@ -188,6 +222,10 @@ export abstract class BaseTracer extends BaseCallbackHandler {
     }
     run.end_time = Date.now();
     run.outputs = outputs;
+    run.events.push({
+      name: "end",
+      time: run.end_time,
+    });
     await this.onChainEnd?.(run);
     await this._endTrace(run);
   }
@@ -199,6 +237,10 @@ export abstract class BaseTracer extends BaseCallbackHandler {
     }
     run.end_time = Date.now();
     run.error = error.message;
+    run.events.push({
+      name: "error",
+      time: run.end_time,
+    });
     await this.onChainError?.(run);
     await this._endTrace(run);
   }
@@ -211,12 +253,19 @@ export abstract class BaseTracer extends BaseCallbackHandler {
     tags?: string[]
   ): Promise<void> {
     const execution_order = this._getExecutionOrder(parentRunId);
+    const start_time = Date.now();
     const run: Run = {
       id: runId,
       name: tool.id[tool.id.length - 1],
       parent_run_id: parentRunId,
-      start_time: Date.now(),
+      start_time,
       serialized: tool,
+      events: [
+        {
+          name: "start",
+          time: start_time,
+        },
+      ],
       inputs: { input },
       execution_order,
       child_execution_order: execution_order,
@@ -237,6 +286,10 @@ export abstract class BaseTracer extends BaseCallbackHandler {
     }
     run.end_time = Date.now();
     run.outputs = { output };
+    run.events.push({
+      name: "end",
+      time: run.end_time,
+    });
     await this.onToolEnd?.(run);
     await this._endTrace(run);
   }
@@ -248,6 +301,10 @@ export abstract class BaseTracer extends BaseCallbackHandler {
     }
     run.end_time = Date.now();
     run.error = error.message;
+    run.events.push({
+      name: "error",
+      time: run.end_time,
+    });
     await this.onToolError?.(run);
     await this._endTrace(run);
   }
@@ -260,9 +317,27 @@ export abstract class BaseTracer extends BaseCallbackHandler {
     const agentRun = run as AgentRun;
     agentRun.actions = agentRun.actions || [];
     agentRun.actions.push(action);
+    agentRun.events.push({
+      name: "agent_action",
+      time: Date.now(),
+      kwargs: { action },
+    });
     await this.onAgentAction?.(run as AgentRun);
   }
 
+  async handleText(text: string, runId: string): Promise<void> {
+    const run = this.runMap.get(runId);
+    if (!run || run?.run_type !== "chain") {
+      return;
+    }
+    run.events.push({
+      name: "text",
+      time: Date.now(),
+      kwargs: { text },
+    });
+    await this.onText?.(run);
+  }
+
   // custom event handlers
 
   onLLMStart?(run: Run): void | Promise<void>;
@@ -289,5 +364,5 @@ export abstract class BaseTracer extends BaseCallbackHandler {
 
   // onAgentEnd?(run: ChainRun): void | Promise<void>;
 
-  // onText?(run: Run): void | Promise<void>;
+  onText?(run: Run): void | Promise<void>;
 }
diff --git a/langchain/src/callbacks/handlers/tracer_langchain.ts b/langchain/src/callbacks/handlers/tracer_langchain.ts
@@ -1,5 +1,9 @@
 import { LangChainPlusClient } from "langchainplus-sdk";
-import { BaseRun, RunCreate, RunUpdate } from "langchainplus-sdk/schemas";
+import {
+  BaseRun,
+  RunCreate,
+  RunUpdate as BaseRunUpdate,
+} from "langchainplus-sdk/schemas";
 import {
   getEnvironmentVariable,
   getRuntimeEnvironment,
@@ -13,6 +17,10 @@ export interface Run extends BaseRun {
   child_execution_order: number;
 }
 
+export interface RunUpdate extends BaseRunUpdate {
+  events: BaseRun["events"];
+}
+
 export interface LangChainTracerFields extends BaseCallbackHandlerInput {
   exampleId?: string;
   sessionName?: string;
@@ -72,6 +80,7 @@ export class LangChainTracer
       end_time: run.end_time,
       error: run.error,
       outputs: run.outputs,
+      events: run.events,
     };
     await this.client.updateRun(run.id, runUpdate);
   }

diff --git a/langchain/src/callbacks/tests/tracer.test.ts b/langchain/src/callbacks/tests/tracer.test.ts
@@ -45,6 +45,16 @@ test("Test LLMRun", async () => {
     execution_order: 1,
     child_execution_order: 1,
     serialized,
+    events: [
+      {
+        name: "start",
+        time: 1620000000000,
+      },
+      {
+        name: "end",
+        time: 1620000000000,
+      },
+    ],
     inputs: { prompts: ["test"] },
     run_type: "llm",
     outputs: { generations: [] },
@@ -72,6 +82,16 @@ test("Test Chat Model Run", async () => {
       "child_execution_order": 1,
       "child_runs": [],
       "end_time": 1620000000000,
+      "events": [
+        {
+          "name": "start",
+          "time": 1620000000000,
+        },
+        {
+          "name": "end",
+          "time": 1620000000000,
+        },
+      ],
       "execution_order": 1,
       "extra": {},
       "id": Any<String>,
@@ -128,6 +148,16 @@ test("Test Chain Run", async () => {
     execution_order: 1,
     child_execution_order: 1,
     serialized,
+    events: [
+      {
+        name: "start",
+        time: 1620000000000,
+      },
+      {
+        name: "end",
+        time: 1620000000000,
+      },
+    ],
     inputs: { foo: "bar" },
     outputs: { foo: "bar" },
     run_type: "chain",
@@ -153,6 +183,16 @@ test("Test Tool Run", async () => {
     execution_order: 1,
     child_execution_order: 1,
     serialized,
+    events: [
+      {
+        name: "start",
+        time: 1620000000000,
+      },
+      {
+        name: "end",
+        time: 1620000000000,
+      },
+    ],
     inputs: { input: "test" },
     outputs: { output: "output" },
     run_type: "tool",
@@ -215,6 +255,16 @@ test("Test nested runs", async () => {
               generations: [[]],
             },
             serialized: { ...serialized, id: ["test_llm_child_run"] },
+            events: [
+              {
+                name: "start",
+                time: 1620000000000,
+              },
+              {
+                name: "end",
+                time: 1620000000000,
+              },
+            ],
             start_time: 1620000000000,
             run_type: "llm",
             child_runs: [],
@@ -227,6 +277,16 @@ test("Test nested runs", async () => {
         child_execution_order: 3,
         outputs: { output: "output" },
         serialized: { ...serialized, id: ["test_tool"] },
+        events: [
+          {
+            name: "start",
+            time: 1620000000000,
+          },
+          {
+            name: "end",
+            time: 1620000000000,
+          },
+        ],
         start_time: 1620000000000,
         inputs: { input: "test" },
         run_type: "tool",
@@ -245,6 +305,16 @@ test("Test nested runs", async () => {
           generations: [[]],
         },
         serialized: { ...serialized, id: ["test_llm2"] },
+        events: [
+          {
+            name: "start",
+            time: 1620000000000,
+          },
+          {
+            name: "end",
+            time: 1620000000000,
+          },
+        ],
         start_time: 1620000000000,
         run_type: "llm",
         child_runs: [],
@@ -262,8 +332,18 @@ test("Test nested runs", async () => {
     outputs: {
       foo: "bar",
     },
-    serialized: { ...serialized, id: ["test"] },
+    events: [
+      {
+        name: "start",
+        time: 1620000000000,
+      },
+      {
+        name: "end",
+        time: 1620000000000,
+      },
+    ],
     name: "test",
+    serialized,
     start_time: 1620000000000,
     run_type: "chain",
     extra: {},