Fix format

langchain-ai · Jul 19, 2024 · 909d611 · 909d611
1 parent 8d23eee
commit 909d611
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 50 deletions.
diff --git a/langchain-core/src/runnables/tests/runnable_stream_events_v2.test.ts b/langchain-core/src/runnables/tests/runnable_stream_events_v2.test.ts
@@ -653,18 +653,6 @@ test("Runnable streamEvents method with llm", async () => {
         a: "b",
       },
     },
-    {
-      event: "on_llm_stream",
-      run_id: expect.any(String),
-      name: "my_model",
-      tags: ["my_model"],
-      metadata: {
-        a: "b",
-      },
-      data: {
-        chunk: "h",
-      },
-    },
     {
       event: "on_llm_stream",
       data: {
@@ -679,18 +667,6 @@ test("Runnable streamEvents method with llm", async () => {
         a: "b",
       },
     },
-    {
-      event: "on_llm_stream",
-      run_id: expect.any(String),
-      name: "my_model",
-      tags: ["my_model"],
-      metadata: {
-        a: "b",
-      },
-      data: {
-        chunk: "e",
-      },
-    },
     {
       event: "on_llm_stream",
       data: {
@@ -705,18 +681,6 @@ test("Runnable streamEvents method with llm", async () => {
         a: "b",
       },
     },
-    {
-      event: "on_llm_stream",
-      run_id: expect.any(String),
-      name: "my_model",
-      tags: ["my_model"],
-      metadata: {
-        a: "b",
-      },
-      data: {
-        chunk: "y",
-      },
-    },
     {
       event: "on_llm_stream",
       data: {
@@ -731,18 +695,6 @@ test("Runnable streamEvents method with llm", async () => {
         a: "b",
       },
     },
-    {
-      event: "on_llm_stream",
-      run_id: expect.any(String),
-      name: "my_model",
-      tags: ["my_model"],
-      metadata: {
-        a: "b",
-      },
-      data: {
-        chunk: "!",
-      },
-    },
     {
       event: "on_llm_end",
       data: {

diff --git a/langchain-core/src/tracers/event_stream.ts b/langchain-core/src/tracers/event_stream.ts
@@ -244,6 +244,13 @@ export class EventStreamCallbackHandler extends BaseTracer {
       yield firstChunk.value;
       return;
     }
+    // Match format from handlers below
+    function _formatOutputChunk(eventType: string, data: unknown) {
+      if (eventType === "llm" && typeof data === "string") {
+        return new GenerationChunk({ text: data });
+      }
+      return data;
+    }
     let tappedPromise = this.tappedPromises.get(runId);
     // if we are the first to tap, issue stream events
     if (tappedPromise === undefined) {
@@ -264,7 +271,9 @@ export class EventStreamCallbackHandler extends BaseTracer {
         await this.send(
           {
             ...event,
-            data: { chunk: firstChunk.value },
+            data: {
+              chunk: _formatOutputChunk(runInfo.runType, firstChunk.value),
+            },
           },
           runInfo
         );
@@ -276,7 +285,7 @@ export class EventStreamCallbackHandler extends BaseTracer {
               {
                 ...event,
                 data: {
-                  chunk,
+                  chunk: _formatOutputChunk(runInfo.runType, chunk),
                 },
               },
               runInfo