From 1f6424906d7dda9db631c0cf0b329475b565c645 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 15 Apr 2024 17:46:16 -0700
Subject: [PATCH] Get it working

---
 frontend/src/components/Chat.tsx          |  7 +++---
 frontend/src/components/MessageEditor.tsx |  6 +++---
 frontend/src/components/StringEditor.tsx  | 13 +-----------
 frontend/src/hooks/useChatMessages.ts     | 26 +++++++++++------------
 frontend/src/hooks/useMessageEditing.ts   | 24 ++++++++++++++++++---
 5 files changed, 42 insertions(+), 34 deletions(-)

diff --git a/frontend/src/components/Chat.tsx b/frontend/src/components/Chat.tsx
index 46a09da0..576bdce4 100644
--- a/frontend/src/components/Chat.tsx
+++ b/frontend/src/components/Chat.tsx
@@ -63,14 +63,15 @@ function CommitEdits(props: {
 
 export function Chat(props: ChatProps) {
   const { chatId } = useParams();
-  const { messages, next } = useChatMessages(
+  const { messages, next, refreshMessages } = useChatMessages(
     chatId ?? null,
     props.stream,
     props.stopStream,
   );
   const { currentChat, assistantConfig, isLoading } = useThreadAndAssistant();
   const { editing, recordEdits, commitEdits, abandonEdits } = useMessageEditing(
-    currentChat?.thread_id,
+    chatId,
+    refreshMessages,
   );
 
   const prevMessages = usePrevious(messages);
@@ -94,7 +95,7 @@ export function Chat(props: ChatProps) {
         editing[msg.id] ? (
           <MessageEditor
             key={msg.id}
-            message={msg}
+            message={editing[msg.id]}
             onUpdate={recordEdits}
             abandonEdits={() => abandonEdits(msg)}
           />
diff --git a/frontend/src/components/MessageEditor.tsx b/frontend/src/components/MessageEditor.tsx
index 847f689a..6cffb285 100644
--- a/frontend/src/components/MessageEditor.tsx
+++ b/frontend/src/components/MessageEditor.tsx
@@ -8,7 +8,7 @@ import {
   PlusCircleIcon,
   TrashIcon,
 } from "@heroicons/react/24/outline";
-import { AutosizeTextarea } from "./StringEditor";
+import { StringEditor } from "./StringEditor";
 import { JsonEditor } from "./JsonEditor";
 
 // TODO adapt (and use) or remove
@@ -99,7 +99,7 @@ function Function(props: {
                       )}
                     >
                       <div className="flex items-center">
-                        <AutosizeTextarea
+                        <StringEditor
                           className="py-0 px-0 prose text-sm leading-normal bg-white"
                           value={str(value)?.toString()}
                           onChange={(newValue) => {
@@ -184,7 +184,7 @@ export function MessageContentEditor(props: {
       return null;
     }
     return (
-      <AutosizeTextarea
+      <StringEditor
         onChange={(newValue) =>
           props.onUpdate({
             ...props.message,
diff --git a/frontend/src/components/StringEditor.tsx b/frontend/src/components/StringEditor.tsx
index ee0b0b89..92f0319d 100644
--- a/frontend/src/components/StringEditor.tsx
+++ b/frontend/src/components/StringEditor.tsx
@@ -1,20 +1,14 @@
-import { Ref } from "react";
 import { cn } from "../utils/cn";
 
 const COMMON_CLS = cn(
   "text-sm col-[1] row-[1] m-0 resize-none overflow-hidden whitespace-pre-wrap break-words bg-transparent px-2 py-1 rounded shadow-none",
 );
 
-export function AutosizeTextarea(props: {
-  id?: string;
-  inputRef?: Ref<HTMLTextAreaElement>;
+export function StringEditor(props: {
   value?: string | null | undefined;
   placeholder?: string;
   className?: string;
   onChange?: (e: string) => void;
-  onFocus?: () => void;
-  onBlur?: () => void;
-  onKeyDown?: (e: React.KeyboardEvent<HTMLTextAreaElement>) => void;
   autoFocus?: boolean;
   readOnly?: boolean;
   cursorPointer?: boolean;
@@ -29,8 +23,6 @@ export function AutosizeTextarea(props: {
       }
     >
       <textarea
-        ref={props.inputRef}
-        id={props.id}
         className={cn(
           COMMON_CLS,
           "text-transparent caret-black rounded focus:outline-0 focus:ring-0",
@@ -42,12 +34,9 @@ export function AutosizeTextarea(props: {
           const target = e.target as HTMLTextAreaElement;
           props.onChange?.(target.value);
         }}
-        onFocus={props.onFocus}
-        onBlur={props.onBlur}
         placeholder={props.placeholder}
         readOnly={props.readOnly}
         autoFocus={props.autoFocus && !props.readOnly}
-        onKeyDown={props.onKeyDown}
       />
       <div
         aria-hidden
diff --git a/frontend/src/hooks/useChatMessages.ts b/frontend/src/hooks/useChatMessages.ts
index c6c2ce25..85c99a57 100644
--- a/frontend/src/hooks/useChatMessages.ts
+++ b/frontend/src/hooks/useChatMessages.ts
@@ -1,4 +1,4 @@
-import { useEffect, useMemo, useRef, useState } from "react";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { Message } from "../types";
 import { StreamState, mergeMessagesById } from "./useStreamState";
 
@@ -23,26 +23,25 @@ export function useChatMessages(
   threadId: string | null,
   stream: StreamState | null,
   stopStream?: (clear?: boolean) => void,
-): { messages: Message[] | null; next: string[] } {
+) {
   const [messages, setMessages] = useState<Message[] | null>(null);
   const [next, setNext] = useState<string[]>([]);
   const prevStreamStatus = usePrevious(stream?.status);
 
-  useEffect(() => {
-    async function fetchMessages() {
-      if (threadId) {
-        const { values, next } = await getState(threadId);
-        setMessages(values);
-        setNext(next);
-      }
+  const refreshMessages = useCallback(async () => {
+    if (threadId) {
+      const { values, next } = await getState(threadId);
+      setMessages(values);
+      setNext(next);
     }
+  }, [threadId]);
 
-    fetchMessages();
-
+  useEffect(() => {
+    refreshMessages();
     return () => {
       setMessages(null);
     };
-  }, [threadId]);
+  }, [threadId, refreshMessages]);
 
   useEffect(() => {
     async function fetchMessages() {
@@ -64,9 +63,10 @@ export function useChatMessages(
 
   return useMemo(
     () => ({
+      refreshMessages,
       messages: mergeMessagesById(messages, stream?.messages),
       next,
     }),
-    [messages, stream?.messages, next],
+    [messages, stream?.messages, next, refreshMessages],
   );
 }
diff --git a/frontend/src/hooks/useMessageEditing.ts b/frontend/src/hooks/useMessageEditing.ts
index 2f65fce7..91ba3ba9 100644
--- a/frontend/src/hooks/useMessageEditing.ts
+++ b/frontend/src/hooks/useMessageEditing.ts
@@ -2,7 +2,10 @@ import { useCallback, useState } from "react";
 import { Message } from "../types";
 import { omit } from "lodash";
 
-export function useMessageEditing(threadId?: string) {
+export function useMessageEditing(
+  threadId: string | undefined,
+  onSuccess: () => void,
+) {
   const [editing, setEditing] = useState<Record<string, Message>>({});
 
   const recordEdits = useCallback((msg: Message) => {
@@ -10,8 +13,23 @@ export function useMessageEditing(threadId?: string) {
   }, []);
   const commitEdits = useCallback(async () => {
     if (!threadId) return;
-    // TODO
-  }, [threadId]);
+    fetch(`/threads/${threadId}/state`, {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify({ values: Object.values(editing) }),
+    })
+      .then((res) => {
+        if (res.ok) {
+          setEditing({});
+          onSuccess();
+        } else {
+          return Promise.reject(res.statusText);
+        }
+      })
+      .catch((e) => {
+        console.error(e);
+      });
+  }, [threadId, editing, onSuccess]);
   const abandonEdits = useCallback((msg?: Message) => {
     if (msg) {
       setEditing((current) => {