From 81ebc7cbbe7eba13a8e2bad757600219447e5028 Mon Sep 17 00:00:00 2001
From: Jacob Lee <jacoblee93@gmail.com>
Date: Fri, 1 Mar 2024 08:13:41 -0700
Subject: [PATCH] Add streaming support for tuple chat widget (#499)

CC @eyurtsev
---
 .../ChatMessageTuplesControlRenderer.tsx      | 33 ++++++++++++++-----
 .../ChatMessagesControlRenderer.tsx           |  6 ++++
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/langserve/playground/src/components/ChatMessageTuplesControlRenderer.tsx b/langserve/playground/src/components/ChatMessageTuplesControlRenderer.tsx
index 9ae40253..841bbb76 100644
--- a/langserve/playground/src/components/ChatMessageTuplesControlRenderer.tsx
+++ b/langserve/playground/src/components/ChatMessageTuplesControlRenderer.tsx
@@ -49,32 +49,47 @@ export const ChatMessageTuplesControlRenderer = withJsonFormsControlProps(
   (props) => {
     const data: Array<MessageTuple> = props.data ?? [];
 
-    useStreamCallback("onSuccess", (ctx) => {
+    useStreamCallback("onChunk", (_chunk, aggregatedState) => {
       if (!isJsonSchemaExtra(props.schema)) return;
       const widget = props.schema.extra.widget;
       if (!("input" in widget) && !("output" in widget)) return;
 
-      const inputPath = getNormalizedJsonPath(widget.input ?? "");
       const outputPath = getNormalizedJsonPath(widget.output ?? "");
 
       const isSingleOutputKey =
-        ctx.output != null &&
-        Object.keys(ctx.output).length === 1 &&
-        Object.keys(ctx.output)[0] === "output";
+        aggregatedState?.final_output != null &&
+        Object.keys(aggregatedState?.final_output).length === 1 &&
+        Object.keys(aggregatedState?.final_output)[0] === "output";
 
-      const human = traverseNaiveJsonPath(ctx.input, inputPath);
-      let ai = traverseNaiveJsonPath(ctx.output, outputPath);
+      let ai = traverseNaiveJsonPath(aggregatedState?.final_output, outputPath);
 
       if (isSingleOutputKey) {
         ai = traverseNaiveJsonPath(ai, ["output", ...outputPath]) ?? ai;
       }
 
       ai = getMessageContent(ai);
-      if (typeof human === "string" && typeof ai === "string") {
-        props.handleChange(props.path, [...data, [human, ai]]);
+      if (typeof ai === "string") {
+        props.handleChange(props.path, [...data.slice(0, -1), [data[data.length - 1][0], ai]]);
+      }
+    });
+
+    useStreamCallback("onStart", (ctx) => {
+      if (!isJsonSchemaExtra(props.schema)) return;
+      const widget = props.schema.extra.widget;
+      if (!("input" in widget) && !("output" in widget)) return;
+
+      const inputPath = getNormalizedJsonPath(widget.input ?? "");
+
+      const human = traverseNaiveJsonPath(ctx.input, inputPath);
+      if (typeof human === "string") {
+        props.handleChange(props.path, [...data, [human, ""]]);
       }
     });
 
+    useStreamCallback("onError", () => {
+      props.handleChange(props.path, [...data.slice(0, -1)]);
+    });
+
     return (
       <div className="control">
         <div className="flex items-center justify-between">
diff --git a/langserve/playground/src/components/ChatMessagesControlRenderer.tsx b/langserve/playground/src/components/ChatMessagesControlRenderer.tsx
index 40e93352..b1bc1711 100644
--- a/langserve/playground/src/components/ChatMessagesControlRenderer.tsx
+++ b/langserve/playground/src/components/ChatMessagesControlRenderer.tsx
@@ -101,6 +101,12 @@ export const ChatMessagesControlRenderer = withJsonFormsControlProps(
       props.handleChange(props.path, [...data, { content: "", type: "human" }]);
     });
 
+    useStreamCallback("onError", () => {
+      if (data.length && data[data.length - 1].type === "ai") {
+        props.handleChange(props.path, [...data.slice(0, -1)]);
+      }
+    });
+
     return (
       <div className="control">
         <div className="flex items-center justify-between">