edit and regenerate messages

2023-03-09 01:38:59 -08:00 · 2023-03-09 01:38:59 -08:00 · aca53b79a0
parent 63bee57226
commit aca53b79a0
3 changed files with 163 additions and 21 deletions
--- 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);
@ -104,9 +93,50 @@ export class ChatManager extends EventEmitter {
            : [];
        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);
--- 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>
--- 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);
@ -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;
 }