From aca53b79a0280a85837502ae4c505cbbc75fee5e Mon Sep 17 00:00:00 2001
From: Cogent Apps <127109874+cogentapps@users.noreply.github.com>
Date: Thu, 9 Mar 2023 01:38:59 -0800
Subject: [PATCH] edit and regenerate messages

---
 src/chat-manager.ts        | 60 +++++++++++++++++++++-------
 src/components/message.tsx | 43 +++++++++++++++++++-
 src/context.tsx            | 81 ++++++++++++++++++++++++++++++++++++--
 3 files changed, 163 insertions(+), 21 deletions(-)

diff --git a/src/chat-manager.ts b/src/chat-manager.ts
index f04720d..a4b96e2 100644
--- a/src/chat-manager.ts
+++ b/src/chat-manager.ts
@@ -2,7 +2,7 @@ import { BroadcastChannel } from 'broadcast-channel';
 import EventEmitter from 'events';
 import MiniSearch, { SearchResult } from 'minisearch'
 import { v4 as uuidv4 } from 'uuid';
-import { Chat, getOpenAIMessageFromMessage, Message, UserSubmittedMessage } from './types';
+import { Chat, getOpenAIMessageFromMessage, Message, Parameters, UserSubmittedMessage } from './types';
 import { MessageTree } from './message-tree';
 import { createStreamingChatCompletion } from './openai';
 import { createTitle } from './titles';
@@ -81,18 +81,7 @@ export class ChatManager extends EventEmitter {
             done: true,
         };
 
-        const reply: Message = {
-            id: uuidv4(),
-            parentID: newMessage.id,
-            chatID: chat.id,
-            timestamp: Date.now(),
-            role: 'assistant',
-            content: '',
-            done: false,
-        };
-
         chat.messages.addMessage(newMessage);
-        chat.messages.addMessage(reply);
         chat.updated = Date.now();
 
         this.emit(chat.id);
@@ -103,10 +92,51 @@ export class ChatManager extends EventEmitter {
             ? chat.messages.getMessageChainTo(message.parentID)
             : [];
         messages.push(newMessage);
-        
+
+        await this.getReply(messages, message.requestedParameters);
+    }
+
+    public async regenerate(message: Message, requestedParameters: Parameters) {
+        const chat = this.chats.get(message.chatID);
+
+        if (!chat) {
+            throw new Error('Chat not found');
+        }
+
+        const messages: Message[] = message.parentID
+            ? chat.messages.getMessageChainTo(message.parentID)
+            : [];
+
+        await this.getReply(messages, requestedParameters);
+    }
+
+    private async getReply(messages: Message[], requestedParameters: Parameters) {
+        const latestMessage = messages[messages.length - 1];
+        const chat = this.chats.get(latestMessage.chatID);
+
+        if (!chat) {
+            throw new Error('Chat not found');
+        }
+
+        const reply: Message = {
+            id: uuidv4(),
+            parentID: latestMessage.id,
+            chatID: latestMessage.chatID,
+            timestamp: Date.now(),
+            role: 'assistant',
+            content: '',
+            done: false,
+        };
+
+        chat.messages.addMessage(reply);
+        chat.updated = Date.now();
+
+        this.emit(chat.id);
+        channel.postMessage({ type: 'chat-update', data: chat });
+
         const messagesToSend = selectMessagesToSendSafely(messages.map(getOpenAIMessageFromMessage));
 
-        const response = await createStreamingChatCompletion(messagesToSend, message.requestedParameters);
+        const response = await createStreamingChatCompletion(messagesToSend, requestedParameters);
 
         response.on('error', () => {
             if (!reply.content) {
@@ -139,7 +169,7 @@ export class ChatManager extends EventEmitter {
             setTimeout(() => this.search.update(chat), 500);
 
             if (!chat.title) {
-                chat.title = await createTitle(chat, message.requestedParameters.apiKey);
+                chat.title = await createTitle(chat, requestedParameters.apiKey);
                 if (chat.title) {
                     this.emit(chat.id);
                     this.emit('title', chat.id, chat.title);
diff --git a/src/components/message.tsx b/src/components/message.tsx
index b3489bb..f8ac335 100644
--- a/src/components/message.tsx
+++ b/src/components/message.tsx
@@ -1,10 +1,12 @@
 import styled from '@emotion/styled';
-import { Button, CopyButton, Loader } from '@mantine/core';
+import { Button, CopyButton, Loader, Textarea } from '@mantine/core';
+import { useState } from 'react';
 
 import { Message } from "../types";
 import { share } from '../utils';
 import { ElevenLabsReaderButton } from '../elevenlabs';
 import { Markdown } from './markdown';
+import { useAppContext } from '../context';
 
 // hide for everyone but screen readers
 const SROnly = styled.span`
@@ -169,6 +171,17 @@ const EndOfChatMarker = styled.div`
     background: rgba(255, 255, 255, 0.1);
 `;
 
+const Editor = styled.div`
+    max-width: 50rem;
+    margin-left: auto;
+    margin-right: auto;
+    margin-top: 0.5rem;
+
+    .mantine-Button-root {
+        margin-top: 1rem;
+    }
+`;
+
 function getRoleName(role: string, share = false) {
     switch (role) {
         case 'user':
@@ -193,6 +206,10 @@ function InlineLoader() {
 }
 
 export default function MessageComponent(props: { message: Message, last: boolean, share?: boolean }) {
+    const context = useAppContext();
+    const [editing, setEditing] = useState(false);
+    const [content, setContent] = useState('');
+
     if (props.message.role === 'system') {
         return null;
     }
@@ -222,8 +239,30 @@ export default function MessageComponent(props: { message: Message, last: boolea
                         <span>Share</span>
                     </Button>
                 )}
+                {!context.isShare && props.message.role === 'user' && (
+                    <Button variant="subtle" size="sm" compact onClick={() => {
+                        setContent(props.message.content);
+                        setEditing(true);
+                    }}>
+                        <i className="fa fa-edit" />
+                        <span>Edit</span>
+                    </Button>
+                )}
+                {!context.isShare && props.message.role === 'assistant' && (
+                    <Button variant="subtle" size="sm" compact onClick={() => context.regenerateMessage(props.message)}>
+                        <i className="fa fa-refresh" />
+                        <span>Regenerate</span>
+                    </Button>
+                )}
             </div>
-            <Markdown content={props.message.content} className={"content content-" + props.message.id} />
+            {!editing && <Markdown content={props.message.content} className={"content content-" + props.message.id} />}
+            {editing && (<Editor>
+                <Textarea value={content}
+                    onChange={e => setContent(e.currentTarget.value)}
+                    autosize={true} />
+                <Button variant="light" onClick={() => context.editMessage(props.message, content)}>Save changes</Button>
+                <Button variant="subtle" onClick={() => setEditing(false)}>Cancel</Button>
+            </Editor>)}
         </div>
         {props.last && <EndOfChatMarker />}
     </Container>
diff --git a/src/context.tsx b/src/context.tsx
index 970a4ee..10492c3 100644
--- a/src/context.tsx
+++ b/src/context.tsx
@@ -5,7 +5,7 @@ import { backend } from "./backend";
 import ChatManagerInstance, { ChatManager } from "./chat-manager";
 import { defaultElevenLabsVoiceID } from "./elevenlabs";
 import { loadParameters, saveParameters } from "./parameters";
-import { Parameters } from "./types";
+import { Message, Parameters } from "./types";
 import { useChat, UseChatResult } from "./use-chat";
 
 export interface Context {
@@ -33,9 +33,11 @@ export interface Context {
     generating: boolean;
     message: string;
     parameters: Parameters;
-    setMessage: (message: string) => void;
+    setMessage: (message: string, parentID?: string) => void;
     setParameters: (parameters: Parameters) => void;
     onNewMessage: (message?: string) => Promise<boolean>;
+    regenerateMessage: (message: Message) => Promise<boolean>;
+    editMessage: (message: Message, content: string) => Promise<boolean>;
 }
 
 const AppContext = React.createContext<Context>({} as any);
@@ -45,7 +47,7 @@ export function useCreateAppContext(): Context {
     const pathname = useLocation().pathname;
     const isShare = pathname.startsWith('/s/');
     const navigate = useNavigate();
-    
+
     const chatManager = useRef(ChatManagerInstance);
     const currentChat = useChat(chatManager.current, id, isShare);
     const [authenticated, setAuthenticated] = useState(backend?.isAuthenticated || false);
@@ -150,6 +152,75 @@ export function useCreateAppContext(): Context {
         return true;
     }, [chatManager, openaiApiKey, id, parameters, message, currentChat.leaf]);
 
+    const regenerateMessage = useCallback(async (message: Message) => {
+        if (isShare) {
+            return false;
+        }
+
+        if (!openaiApiKey) {
+            setSettingsTab('user');
+            setOption('openai-api-key');
+            return false;
+        }
+
+        setGenerating(true);
+
+        await chatManager.current.regenerate(message, {
+            ...parameters,
+            apiKey: openaiApiKey,
+        });
+
+        setTimeout(() => setGenerating(false), 4000);
+
+        return true;
+    }, [chatManager, openaiApiKey, id, parameters]);
+
+    const editMessage = useCallback(async (message: Message, content: string) => {
+        if (isShare) {
+            return false;
+        }
+
+        if (!content?.trim().length) {
+            return false;
+        }
+
+        if (!openaiApiKey) {
+            setSettingsTab('user');
+            setOption('openai-api-key');
+            return false;
+        }
+
+        setGenerating(true);
+
+        if (id) {
+            await chatManager.current.sendMessage({
+                chatID: id,
+                content: content.trim(),
+                requestedParameters: {
+                    ...parameters,
+                    apiKey: openaiApiKey,
+                },
+                parentID: message.parentID,
+            });
+        } else {
+            const id = await chatManager.current.createChat();
+            await chatManager.current.sendMessage({
+                chatID: id,
+                content: content.trim(),
+                requestedParameters: {
+                    ...parameters,
+                    apiKey: openaiApiKey,
+                },
+                parentID: message.parentID,
+            });
+            navigate('/chat/' + id);
+        }
+
+        setTimeout(() => setGenerating(false), 4000);
+
+        return true;
+    }, [chatManager, openaiApiKey, id, parameters, message, currentChat.leaf]);
+
     const context = useMemo<Context>(() => ({
         authenticated,
         id,
@@ -184,8 +255,10 @@ export function useCreateAppContext(): Context {
         setMessage,
         setParameters,
         onNewMessage,
+        regenerateMessage,
+        editMessage,
     }), [chatManager, authenticated, openaiApiKey, elevenLabsApiKey, settingsTab, option, voiceID,
-        generating, message, parameters, onNewMessage, currentChat]);
+        generating, message, parameters, onNewMessage, regenerateMessage, editMessage, currentChat]);
 
     return context;
 }