blair.nz / cai
this repo has no description
0
fork

Configure Feed

Select the types of activity you want to include in your feed.

Add edit last message functionality to chat messages

James Blair 26bdd32d 7245b1e6

+511 -37
+102
apps/api/src/routes/chats.ts
··· 218 218 }, 219 219 }); 220 220 }) 221 + .post( 222 + "/:id/messages/edit-last", 223 + zValidator("json", messageSchema), 224 + async (c) => { 225 + const { id } = c.req.param(); 226 + const body = c.req.valid("json"); 227 + const userId = c.get("userId"); 228 + 229 + const chat = await db.query.chats.findFirst({ 230 + where: and(eq(schema.chats.id, id), eq(schema.chats.userId, userId)), 231 + }); 232 + 233 + if (!chat) return c.json({ message: "Not found" }, 404); 234 + 235 + const character = characters.find( 236 + (character) => character.id === chat.characterId, 237 + ); 238 + 239 + if (!character) return c.json({ message: "Invalid character" }, 400); 240 + 241 + const messages = await db.query.messages.findMany({ 242 + where: and( 243 + eq(schema.messages.chatId, id), 244 + eq(schema.messages.userId, userId), 245 + ), 246 + orderBy: asc(schema.messages.createdAt), 247 + }); 248 + const lastAssistantMessage = messages.at(-1); 249 + const lastUserMessage = messages.at(-2); 250 + 251 + if ( 252 + !lastUserMessage || 253 + lastUserMessage.role !== "user" || 254 + !lastAssistantMessage || 255 + lastAssistantMessage.role !== "assistant" 256 + ) { 257 + return c.json({ message: "No editable message found" }, 400); 258 + } 259 + 260 + db.update(schema.messages) 261 + .set({ text: body.text }) 262 + .where( 263 + and( 264 + eq(schema.messages.id, lastUserMessage.id), 265 + eq(schema.messages.chatId, id), 266 + eq(schema.messages.userId, userId), 267 + ), 268 + ) 269 + .run(); 270 + 271 + const system = "prompt" in character ? character.prompt : undefined; 272 + const res = streamChat({ 273 + system, 274 + modelId: chat.modelId, 275 + messages: [ 276 + ...messages.slice(0, -2).map((message) => ({ 277 + role: message.role, 278 + content: message.text, 279 + })), 280 + { 281 + role: "user", 282 + content: body.text, 283 + }, 284 + ], 285 + }); 286 + 287 + const stream = new ReadableStream({ 288 + async start(controller) { 289 + const encoder = new TextEncoder(); 290 + let assistantText = ""; 291 + 292 + try { 293 + for await (const text of res.textStream) { 294 + assistantText += text; 295 + controller.enqueue(encoder.encode(text)); 296 + } 297 + 298 + db.update(schema.messages) 299 + .set({ text: assistantText }) 300 + .where( 301 + and( 302 + eq(schema.messages.id, lastAssistantMessage.id), 303 + eq(schema.messages.chatId, id), 304 + eq(schema.messages.userId, userId), 305 + ), 306 + ) 307 + .run(); 308 + 309 + controller.close(); 310 + } catch (error) { 311 + controller.error(error); 312 + } 313 + }, 314 + }); 315 + 316 + return new Response(stream, { 317 + headers: { 318 + "Content-Type": "text/plain; charset=utf-8", 319 + }, 320 + }); 321 + }, 322 + ) 221 323 .patch("/:id", zValidator("json", chatSchema.partial()), async (c) => { 222 324 const { id } = c.req.param(); 223 325 const body = c.req.valid("json");
+14 -2
apps/web/src/components/chat-message-actions.tsx
··· 1 1 import { CopyButton } from "./copy-button"; 2 + import { EditMessageButton } from "./edit-message-button"; 2 3 import { Row } from "./layout"; 3 4 import { RegenerateMessageButton } from "./regenerate-message-button"; 4 5 5 6 interface ChatMessageActionsProps { 7 + canEdit?: boolean; 6 8 canRegenerate?: boolean; 9 + isEditing?: boolean; 10 + isPending?: boolean; 7 11 isRegenerating?: boolean; 12 + onEdit?: () => void; 8 13 onRegenerate?: () => void; 9 14 text: string; 10 15 } 11 16 12 17 export function ChatMessageActions({ 18 + canEdit, 13 19 canRegenerate, 20 + isEditing, 21 + isPending, 14 22 isRegenerating, 23 + onEdit, 15 24 onRegenerate, 16 25 text, 17 26 }: ChatMessageActionsProps) { ··· 20 29 gap="xs" 21 30 className="opacity-0 transition-opacity duration-150 group-focus-within:opacity-100 group-hover:opacity-100" 22 31 > 32 + <CopyButton text={text} /> 33 + {canEdit && onEdit && ( 34 + <EditMessageButton disabled={isEditing || isPending} onClick={onEdit} /> 35 + )} 23 36 {canRegenerate && onRegenerate && ( 24 37 <RegenerateMessageButton 25 - disabled={isRegenerating} 38 + disabled={isPending || isRegenerating} 26 39 onClick={onRegenerate} 27 40 /> 28 41 )} 29 - <CopyButton text={text} /> 30 42 </Row> 31 43 ); 32 44 }
+44
apps/web/src/components/chat-message-bubble.tsx
··· 1 + import { cva, type VariantProps } from "class-variance-authority"; 2 + import Markdown from "react-markdown"; 3 + import { cn } from "#/lib/utils"; 4 + import { TypingIndicator } from "./typing-indicator"; 5 + 6 + const chatMessageBubble = cva( 7 + "rounded-2xl px-4 py-3 text-sm leading-relaxed wrap-anywhere max-w-[80%]", 8 + { 9 + variants: { 10 + author: { 11 + assistant: "bg-primary text-primary-foreground", 12 + user: "bg-secondary text-secondary-foreground", 13 + }, 14 + }, 15 + defaultVariants: { 16 + author: "assistant", 17 + }, 18 + }, 19 + ); 20 + 21 + type ChatMessageBubbleProps = VariantProps<typeof chatMessageBubble> & { 22 + text: string; 23 + className?: string; 24 + isTyping?: boolean; 25 + }; 26 + 27 + export function ChatMessageBubble({ 28 + className, 29 + author, 30 + text, 31 + isTyping, 32 + }: ChatMessageBubbleProps) { 33 + return ( 34 + <div 35 + className={cn( 36 + chatMessageBubble({ author }), 37 + !isTyping && "markdown", 38 + className, 39 + )} 40 + > 41 + {isTyping ? <TypingIndicator /> : <Markdown>{text}</Markdown>} 42 + </div> 43 + ); 44 + }
+54
apps/web/src/components/chat-message-edit-form.tsx
··· 1 + import { Row, Stack } from "./layout"; 2 + import { Button } from "./ui/button"; 3 + import { Textarea } from "./ui/textarea"; 4 + 5 + interface ChatMessageEditFormProps { 6 + draftText: string; 7 + disabled?: boolean; 8 + isSaveDisabled?: boolean; 9 + onCancel?: () => void; 10 + onChange: (text: string) => void; 11 + onSubmit: (event: React.FormEvent<HTMLFormElement>) => void; 12 + } 13 + 14 + export function ChatMessageEditForm({ 15 + draftText, 16 + disabled, 17 + isSaveDisabled, 18 + onCancel, 19 + onChange, 20 + onSubmit, 21 + }: ChatMessageEditFormProps) { 22 + return ( 23 + <form onSubmit={onSubmit} className="w-full max-w-[80%]"> 24 + <Stack gap="xs" className="w-full"> 25 + <Textarea 26 + autoFocus 27 + value={draftText} 28 + disabled={disabled} 29 + className="wrap-anywhere min-h-25 focus-visible:border-transparent focus-visible:ring-0" 30 + onChange={(event) => onChange(event.currentTarget.value)} 31 + onKeyDown={(event) => { 32 + if (event.key === "Enter" && !event.shiftKey) { 33 + event.preventDefault(); 34 + event.currentTarget.form?.requestSubmit(); 35 + } 36 + }} 37 + /> 38 + <Row gap="xs" justify="end"> 39 + <Button 40 + type="button" 41 + variant="ghost" 42 + disabled={disabled} 43 + onClick={onCancel} 44 + > 45 + Cancel 46 + </Button> 47 + <Button type="submit" disabled={isSaveDisabled}> 48 + Save 49 + </Button> 50 + </Row> 51 + </Stack> 52 + </form> 53 + ); 54 + }
+57 -30
apps/web/src/components/chat-message.tsx
··· 1 - import { cva, type VariantProps } from "class-variance-authority"; 2 - import Markdown from "react-markdown"; 1 + import { useEffect, useState } from "react"; 3 2 import { cn } from "#/lib/utils"; 4 3 import { ChatMessageActions } from "./chat-message-actions"; 4 + import { ChatMessageBubble } from "./chat-message-bubble"; 5 + import { ChatMessageEditForm } from "./chat-message-edit-form"; 5 6 import { Stack } from "./layout"; 6 - import { TypingIndicator } from "./typing-indicator"; 7 7 8 - const chatMessage = cva( 9 - "rounded-2xl px-4 py-3 text-sm leading-relaxed wrap-anywhere max-w-[80%]", 10 - { 11 - variants: { 12 - author: { 13 - assistant: "bg-primary text-primary-foreground", 14 - user: "bg-secondary text-secondary-foreground", 15 - }, 16 - }, 17 - defaultVariants: { 18 - author: "assistant", 19 - }, 20 - }, 21 - ); 22 - 23 - type ChatMessageProps = VariantProps<typeof chatMessage> & { 8 + interface ChatMessageProps { 9 + author?: "assistant" | "user" | null; 24 10 text: string; 25 11 className?: string; 12 + canEdit?: boolean; 26 13 canRegenerate?: boolean; 14 + isEditing?: boolean; 15 + isPending?: boolean; 27 16 isRegenerating?: boolean; 17 + onCancelEdit?: () => void; 18 + onEdit?: () => void; 19 + onSaveEdit?: (text: string) => Promise<void> | void; 28 20 onRegenerate?: () => void; 29 - }; 21 + } 30 22 31 23 export function ChatMessage({ 24 + canEdit, 32 25 canRegenerate, 33 26 className, 34 27 author, 35 28 text, 29 + isEditing, 30 + isPending, 36 31 isRegenerating, 32 + onCancelEdit, 33 + onEdit, 34 + onSaveEdit, 37 35 onRegenerate, 38 36 }: ChatMessageProps) { 39 37 const isTyping = author === "assistant" && text.length === 0; 38 + const [draftText, setDraftText] = useState(text); 39 + 40 + useEffect(() => { 41 + if (!isEditing) { 42 + setDraftText(text); 43 + } 44 + }, [isEditing, text]); 45 + 46 + const isSaveDisabled = 47 + isPending || !draftText.trim() || draftText.trim() === text.trim(); 48 + 49 + async function handleSubmit(event: React.FormEvent<HTMLFormElement>) { 50 + event.preventDefault(); 51 + if (!onSaveEdit || isSaveDisabled) return; 52 + 53 + await onSaveEdit(draftText); 54 + } 40 55 41 56 return ( 42 57 <Stack ··· 48 63 : "items-end self-end", 49 64 )} 50 65 > 51 - <div 52 - className={cn( 53 - chatMessage({ author }), 54 - !isTyping && "markdown", 55 - className, 56 - )} 57 - > 58 - {isTyping ? <TypingIndicator /> : <Markdown>{text}</Markdown>} 59 - </div> 60 - {!isTyping && text && ( 66 + {isEditing ? ( 67 + <ChatMessageEditForm 68 + draftText={draftText} 69 + disabled={isPending} 70 + isSaveDisabled={isSaveDisabled} 71 + onCancel={onCancelEdit} 72 + onChange={setDraftText} 73 + onSubmit={handleSubmit} 74 + /> 75 + ) : ( 76 + <ChatMessageBubble 77 + author={author} 78 + className={className} 79 + isTyping={isTyping} 80 + text={text} 81 + /> 82 + )} 83 + {!isTyping && text && !isEditing && ( 61 84 <ChatMessageActions 85 + canEdit={canEdit} 62 86 canRegenerate={canRegenerate} 87 + isEditing={isEditing} 88 + isPending={isPending} 63 89 isRegenerating={isRegenerating} 90 + onEdit={onEdit} 64 91 onRegenerate={onRegenerate} 65 92 text={text} 66 93 />
+40 -1
apps/web/src/components/chat-messages.tsx
··· 1 - import { useEffect, useLayoutEffect, useRef } from "react"; 1 + import { useEffect, useLayoutEffect, useRef, useState } from "react"; 2 + import { useEditLastMessage } from "#/hooks/use-edit-last-message"; 2 3 import { useRegenerateLastMessage } from "#/hooks/use-regenerate-last-message"; 3 4 import type { ChatMessage as ChatMessageType } from "#/lib/types"; 4 5 import { ChatMessage } from "./chat-message"; ··· 27 28 const hasInitializedScrollRef = useRef(false); 28 29 const shouldStickToBottomRef = useRef(true); 29 30 const previousMessagesRef = useRef(messages); 31 + const [editingMessageId, setEditingMessageId] = useState<string | null>(null); 32 + const editLastMessage = useEditLastMessage(); 30 33 const regenerateLastMessage = useRegenerateLastMessage(); 34 + const lastAssistantMessage = messages.at(-1); 35 + const lastEditableUserMessage = 36 + lastAssistantMessage?.role === "assistant" && 37 + lastAssistantMessage.text.trim().length > 0 && 38 + messages.at(-2)?.role === "user" 39 + ? messages.at(-2) 40 + : undefined; 31 41 32 42 useEffect(() => { 33 43 const container = containerRef.current; ··· 47 57 }; 48 58 }, []); 49 59 60 + useEffect(() => { 61 + if ( 62 + editingMessageId && 63 + lastEditableUserMessage && 64 + editingMessageId !== lastEditableUserMessage.id 65 + ) { 66 + setEditingMessageId(null); 67 + } 68 + 69 + if (!lastEditableUserMessage && editingMessageId) { 70 + setEditingMessageId(null); 71 + } 72 + }, [editingMessageId, lastEditableUserMessage]); 73 + 50 74 useLayoutEffect(() => { 51 75 const container = containerRef.current; 52 76 if (!(container instanceof HTMLDivElement) || messages.length === 0) return; ··· 85 109 <ChatMessage 86 110 key={message.id} 87 111 author={message.role} 112 + canEdit={message.id === lastEditableUserMessage?.id} 88 113 canRegenerate={ 89 114 message.id === messages.at(-1)?.id && 90 115 message.role === "assistant" && 91 116 message.text.trim().length > 0 92 117 } 118 + isEditing={message.id === editingMessageId} 119 + isPending={ 120 + editLastMessage.isPending || regenerateLastMessage.isPending 121 + } 93 122 isRegenerating={regenerateLastMessage.isPending} 123 + onCancelEdit={() => { 124 + setEditingMessageId(null); 125 + }} 126 + onEdit={() => { 127 + setEditingMessageId(message.id); 128 + }} 129 + onSaveEdit={async (text) => { 130 + await editLastMessage.mutateAsync({ chatId, text }); 131 + setEditingMessageId(null); 132 + }} 94 133 onRegenerate={() => { 95 134 regenerateLastMessage.mutate({ chatId }); 96 135 }}
+1 -2
apps/web/src/components/copy-button.tsx
··· 29 29 return ( 30 30 <Button 31 31 type="button" 32 - size="icon-xs" 32 + size="icon-sm" 33 33 variant="ghost" 34 34 onClick={handleCopy} 35 - className="text-muted-foreground hover:text-foreground" 36 35 aria-label={isCopied ? "Copied message" : "Copy message"} 37 36 title={isCopied ? "Copied" : "Copy message"} 38 37 >
+26
apps/web/src/components/edit-message-button.tsx
··· 1 + import { PencilIcon } from "lucide-react"; 2 + import { Button } from "./ui/button"; 3 + 4 + interface EditMessageButtonProps { 5 + disabled?: boolean; 6 + onClick: () => void; 7 + } 8 + 9 + export function EditMessageButton({ 10 + disabled, 11 + onClick, 12 + }: EditMessageButtonProps) { 13 + return ( 14 + <Button 15 + type="button" 16 + size="icon-sm" 17 + variant="ghost" 18 + onClick={onClick} 19 + disabled={disabled} 20 + aria-label="Edit message" 21 + title="Edit message" 22 + > 23 + <PencilIcon /> 24 + </Button> 25 + ); 26 + }
+1 -2
apps/web/src/components/regenerate-message-button.tsx
··· 13 13 return ( 14 14 <Button 15 15 type="button" 16 - size="icon-xs" 16 + size="icon-sm" 17 17 variant="ghost" 18 18 onClick={onClick} 19 19 disabled={disabled} 20 - className="text-muted-foreground hover:text-foreground" 21 20 aria-label="Regenerate message" 22 21 title="Regenerate message" 23 22 >
+172
apps/web/src/hooks/use-edit-last-message.ts
··· 1 + import { useAuth } from "@clerk/clerk-react"; 2 + import { 3 + type QueryClient, 4 + useMutation, 5 + useQueryClient, 6 + } from "@tanstack/react-query"; 7 + import { client } from "#/lib/api"; 8 + import { chatQueryKey, chatsQueryKey } from "#/lib/chat-query"; 9 + import type { Chat, ChatMessage } from "#/lib/types"; 10 + 11 + function setMessages( 12 + queryClient: QueryClient, 13 + userId: string | null | undefined, 14 + chatId: string, 15 + update: (messages: ChatMessage[]) => ChatMessage[], 16 + ) { 17 + queryClient.setQueryData<Chat | undefined>( 18 + chatQueryKey(userId, chatId), 19 + (chat) => { 20 + if (!chat) return chat; 21 + 22 + return { 23 + ...chat, 24 + messages: update(chat.messages), 25 + }; 26 + }, 27 + ); 28 + } 29 + 30 + function patchMessage( 31 + queryClient: QueryClient, 32 + userId: string | null | undefined, 33 + chatId: string, 34 + messageId: string, 35 + update: (message: ChatMessage) => ChatMessage, 36 + ) { 37 + setMessages(queryClient, userId, chatId, (messages) => 38 + messages.map((message) => 39 + message.id === messageId ? update(message) : message, 40 + ), 41 + ); 42 + } 43 + 44 + async function editLastMessage(chatId: string, text: string) { 45 + const response = await client.api.chats[":id"].messages["edit-last"].$post({ 46 + param: { id: chatId }, 47 + json: { text }, 48 + }); 49 + 50 + if (!response.ok) { 51 + throw new Error("Failed to edit message"); 52 + } 53 + 54 + return response; 55 + } 56 + 57 + async function readTextStream( 58 + response: Response, 59 + onChunk: (chunk: string) => void, 60 + ) { 61 + if (!response.body) { 62 + onChunk(await response.text()); 63 + return; 64 + } 65 + 66 + const reader = response.body.getReader(); 67 + const decoder = new TextDecoder(); 68 + 69 + while (true) { 70 + const { done, value } = await reader.read(); 71 + 72 + if (done) break; 73 + 74 + onChunk(decoder.decode(value, { stream: true })); 75 + } 76 + 77 + const remainingText = decoder.decode(); 78 + 79 + if (remainingText) { 80 + onChunk(remainingText); 81 + } 82 + } 83 + 84 + export function useEditLastMessage() { 85 + const { userId } = useAuth(); 86 + const queryClient = useQueryClient(); 87 + 88 + return useMutation({ 89 + mutationFn: async ({ chatId, text }: { chatId: string; text: string }) => { 90 + const trimmedText = text.trim(); 91 + 92 + if (!trimmedText) { 93 + throw new Error("Message is required"); 94 + } 95 + 96 + const chat = queryClient.getQueryData<Chat | undefined>( 97 + chatQueryKey(userId, chatId), 98 + ); 99 + const lastAssistantMessage = chat?.messages.at(-1); 100 + const lastUserMessage = chat?.messages.at(-2); 101 + 102 + if ( 103 + !lastUserMessage || 104 + lastUserMessage.role !== "user" || 105 + !lastAssistantMessage || 106 + lastAssistantMessage.role !== "assistant" 107 + ) { 108 + throw new Error("No editable message found"); 109 + } 110 + 111 + await queryClient.cancelQueries({ 112 + queryKey: chatQueryKey(userId, chatId), 113 + }); 114 + 115 + patchMessage( 116 + queryClient, 117 + userId, 118 + chatId, 119 + lastUserMessage.id, 120 + (message) => ({ 121 + ...message, 122 + text: trimmedText, 123 + }), 124 + ); 125 + patchMessage( 126 + queryClient, 127 + userId, 128 + chatId, 129 + lastAssistantMessage.id, 130 + (message) => ({ 131 + ...message, 132 + text: "", 133 + }), 134 + ); 135 + 136 + try { 137 + const response = await editLastMessage(chatId, trimmedText); 138 + await readTextStream(response, (chunk) => { 139 + patchMessage( 140 + queryClient, 141 + userId, 142 + chatId, 143 + lastAssistantMessage.id, 144 + (message) => ({ 145 + ...message, 146 + text: `${message.text}${chunk}`, 147 + }), 148 + ); 149 + }); 150 + } catch (error) { 151 + patchMessage( 152 + queryClient, 153 + userId, 154 + chatId, 155 + lastAssistantMessage.id, 156 + (message) => ({ 157 + ...message, 158 + text: "Failed to stream response.", 159 + }), 160 + ); 161 + throw error; 162 + } finally { 163 + await queryClient.invalidateQueries({ 164 + queryKey: chatQueryKey(userId, chatId), 165 + }); 166 + await queryClient.invalidateQueries({ 167 + queryKey: chatsQueryKey(userId), 168 + }); 169 + } 170 + }, 171 + }); 172 + }