this repo has no description
Add regenerate message button to chat messages (#8)
### TL;DR
Adds the ability to regenerate the last assistant message in a chat.
### What changed?
- Added a `POST /:id/messages/regenerate` API endpoint that streams a regenerated response for the last assistant message in a chat, replacing its content in the database once streaming completes.
- Added a `RegenerateMessageButton` component (using a `RotateCcwIcon`) that appears in the message actions row for the last assistant message.
- Added a `useRegenerateLastMessage` hook that handles optimistically clearing the last assistant message text, streaming the new response chunk-by-chunk into the query cache, and invalidating relevant queries on completion or error.
- The regenerate button is only shown on the most recent assistant message that has non-empty text, and is disabled while a regeneration is in progress.
### How to test?
1. Open an existing chat that has at least one assistant response.
2. Hover over the last assistant message to reveal the action buttons.
3. Click the regenerate (↺) button and confirm the message clears and streams a new response.
4. Verify the regenerated content persists after the stream completes.
5. Confirm the button is disabled while regeneration is in progress and re-enables afterward.
### Why make this change?
Users may want to get a different response from the assistant without having to resend their message. This feature provides a straightforward way to retry the last assistant message in a chat.
authored by
James Blair
and committed by
GitHub
7245b1e6
45f8b8d9
+304
-5
+74
apps/api/src/routes/chats.ts
+74
apps/api/src/routes/chats.ts
···
144
144
},
145
145
});
146
146
})
147
+
.post("/:id/messages/regenerate", async (c) => {
148
+
const { id } = c.req.param();
149
+
const userId = c.get("userId");
150
+
151
+
const chat = await db.query.chats.findFirst({
152
+
where: and(eq(schema.chats.id, id), eq(schema.chats.userId, userId)),
153
+
});
154
+
155
+
if (!chat) return c.json({ message: "Not found" }, 404);
156
+
157
+
const character = characters.find(
158
+
(character) => character.id === chat.characterId,
159
+
);
160
+
161
+
if (!character) return c.json({ message: "Invalid character" }, 400);
162
+
163
+
const messages = await db.query.messages.findMany({
164
+
where: and(
165
+
eq(schema.messages.chatId, id),
166
+
eq(schema.messages.userId, userId),
167
+
),
168
+
orderBy: asc(schema.messages.createdAt),
169
+
});
170
+
const lastAssistantMessage = messages.at(-1);
171
+
172
+
if (!lastAssistantMessage || lastAssistantMessage.role !== "assistant") {
173
+
return c.json({ message: "No assistant message to regenerate" }, 400);
174
+
}
175
+
176
+
const system = "prompt" in character ? character.prompt : undefined;
177
+
const res = streamChat({
178
+
system,
179
+
modelId: chat.modelId,
180
+
messages: messages.slice(0, -1).map((message) => ({
181
+
role: message.role,
182
+
content: message.text,
183
+
})),
184
+
});
185
+
186
+
const stream = new ReadableStream({
187
+
async start(controller) {
188
+
const encoder = new TextEncoder();
189
+
let assistantText = "";
190
+
191
+
try {
192
+
for await (const text of res.textStream) {
193
+
assistantText += text;
194
+
controller.enqueue(encoder.encode(text));
195
+
}
196
+
197
+
db.update(schema.messages)
198
+
.set({ text: assistantText })
199
+
.where(
200
+
and(
201
+
eq(schema.messages.id, lastAssistantMessage.id),
202
+
eq(schema.messages.chatId, id),
203
+
eq(schema.messages.userId, userId),
204
+
),
205
+
)
206
+
.run();
207
+
208
+
controller.close();
209
+
} catch (error) {
210
+
controller.error(error);
211
+
}
212
+
},
213
+
});
214
+
215
+
return new Response(stream, {
216
+
headers: {
217
+
"Content-Type": "text/plain; charset=utf-8",
218
+
},
219
+
});
220
+
})
147
221
.patch("/:id", zValidator("json", chatSchema.partial()), async (c) => {
148
222
const { id } = c.req.param();
149
223
const body = c.req.valid("json");
+16
-1
apps/web/src/components/chat-message-actions.tsx
+16
-1
apps/web/src/components/chat-message-actions.tsx
···
1
1
import { CopyButton } from "./copy-button";
2
2
import { Row } from "./layout";
3
+
import { RegenerateMessageButton } from "./regenerate-message-button";
3
4
4
5
interface ChatMessageActionsProps {
6
+
canRegenerate?: boolean;
7
+
isRegenerating?: boolean;
8
+
onRegenerate?: () => void;
5
9
text: string;
6
10
}
7
11
8
-
export function ChatMessageActions({ text }: ChatMessageActionsProps) {
12
+
export function ChatMessageActions({
13
+
canRegenerate,
14
+
isRegenerating,
15
+
onRegenerate,
16
+
text,
17
+
}: ChatMessageActionsProps) {
9
18
return (
10
19
<Row
11
20
gap="xs"
12
21
className="opacity-0 transition-opacity duration-150 group-focus-within:opacity-100 group-hover:opacity-100"
13
22
>
23
+
{canRegenerate && onRegenerate && (
24
+
<RegenerateMessageButton
25
+
disabled={isRegenerating}
26
+
onClick={onRegenerate}
27
+
/>
28
+
)}
14
29
<CopyButton text={text} />
15
30
</Row>
16
31
);
+19
-2
apps/web/src/components/chat-message.tsx
+19
-2
apps/web/src/components/chat-message.tsx
···
23
23
type ChatMessageProps = VariantProps<typeof chatMessage> & {
24
24
text: string;
25
25
className?: string;
26
+
canRegenerate?: boolean;
27
+
isRegenerating?: boolean;
28
+
onRegenerate?: () => void;
26
29
};
27
30
28
-
export function ChatMessage({ className, author, text }: ChatMessageProps) {
31
+
export function ChatMessage({
32
+
canRegenerate,
33
+
className,
34
+
author,
35
+
text,
36
+
isRegenerating,
37
+
onRegenerate,
38
+
}: ChatMessageProps) {
29
39
const isTyping = author === "assistant" && text.length === 0;
30
40
31
41
return (
···
47
57
>
48
58
{isTyping ? <TypingIndicator /> : <Markdown>{text}</Markdown>}
49
59
</div>
50
-
{!isTyping && text && <ChatMessageActions text={text} />}
60
+
{!isTyping && text && (
61
+
<ChatMessageActions
62
+
canRegenerate={canRegenerate}
63
+
isRegenerating={isRegenerating}
64
+
onRegenerate={onRegenerate}
65
+
text={text}
66
+
/>
67
+
)}
51
68
</Stack>
52
69
);
53
70
}
+13
-1
apps/web/src/components/chat-messages.tsx
+13
-1
apps/web/src/components/chat-messages.tsx
···
1
1
import { useEffect, useLayoutEffect, useRef } from "react";
2
+
import { useRegenerateLastMessage } from "#/hooks/use-regenerate-last-message";
2
3
import type { ChatMessage as ChatMessageType } from "#/lib/types";
3
4
import { ChatMessage } from "./chat-message";
4
5
import { Stack } from "./layout";
5
6
6
7
interface ChatMessageProps {
8
+
chatId: string;
7
9
messages: ChatMessageType[];
8
10
}
9
11
···
20
22
});
21
23
}
22
24
23
-
export function ChatMessages({ messages }: ChatMessageProps) {
25
+
export function ChatMessages({ chatId, messages }: ChatMessageProps) {
24
26
const containerRef = useRef<HTMLDivElement>(null);
25
27
const hasInitializedScrollRef = useRef(false);
26
28
const shouldStickToBottomRef = useRef(true);
27
29
const previousMessagesRef = useRef(messages);
30
+
const regenerateLastMessage = useRegenerateLastMessage();
28
31
29
32
useEffect(() => {
30
33
const container = containerRef.current;
···
82
85
<ChatMessage
83
86
key={message.id}
84
87
author={message.role}
88
+
canRegenerate={
89
+
message.id === messages.at(-1)?.id &&
90
+
message.role === "assistant" &&
91
+
message.text.trim().length > 0
92
+
}
93
+
isRegenerating={regenerateLastMessage.isPending}
94
+
onRegenerate={() => {
95
+
regenerateLastMessage.mutate({ chatId });
96
+
}}
85
97
text={message.text}
86
98
/>
87
99
))}
+1
-1
apps/web/src/components/chat.tsx
+1
-1
apps/web/src/components/chat.tsx
+154
apps/web/src/hooks/use-regenerate-last-message.ts
+154
apps/web/src/hooks/use-regenerate-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 regenerateLastMessage(chatId: string) {
45
+
const response = await client.api.chats[":id"].messages.regenerate.$post({
46
+
param: { id: chatId },
47
+
});
48
+
49
+
if (!response.ok) {
50
+
throw new Error("Failed to regenerate message");
51
+
}
52
+
53
+
return response;
54
+
}
55
+
56
+
async function readTextStream(
57
+
response: Response,
58
+
onChunk: (chunk: string) => void,
59
+
) {
60
+
if (!response.body) {
61
+
onChunk(await response.text());
62
+
return;
63
+
}
64
+
65
+
const reader = response.body.getReader();
66
+
const decoder = new TextDecoder();
67
+
68
+
while (true) {
69
+
const { done, value } = await reader.read();
70
+
71
+
if (done) break;
72
+
73
+
onChunk(decoder.decode(value, { stream: true }));
74
+
}
75
+
76
+
const remainingText = decoder.decode();
77
+
78
+
if (remainingText) {
79
+
onChunk(remainingText);
80
+
}
81
+
}
82
+
83
+
export function useRegenerateLastMessage() {
84
+
const { userId } = useAuth();
85
+
const queryClient = useQueryClient();
86
+
87
+
return useMutation({
88
+
mutationFn: async ({ chatId }: { chatId: string }) => {
89
+
const chat = queryClient.getQueryData<Chat | undefined>(
90
+
chatQueryKey(userId, chatId),
91
+
);
92
+
const lastAssistantMessage = [...(chat?.messages ?? [])]
93
+
.reverse()
94
+
.find(
95
+
(message) =>
96
+
message.role === "assistant" && message.text.trim().length > 0,
97
+
);
98
+
99
+
if (!lastAssistantMessage) {
100
+
throw new Error("No assistant message to regenerate");
101
+
}
102
+
103
+
await queryClient.cancelQueries({
104
+
queryKey: chatQueryKey(userId, chatId),
105
+
});
106
+
107
+
patchMessage(
108
+
queryClient,
109
+
userId,
110
+
chatId,
111
+
lastAssistantMessage.id,
112
+
(message) => ({
113
+
...message,
114
+
text: "",
115
+
}),
116
+
);
117
+
118
+
try {
119
+
const response = await regenerateLastMessage(chatId);
120
+
await readTextStream(response, (chunk) => {
121
+
patchMessage(
122
+
queryClient,
123
+
userId,
124
+
chatId,
125
+
lastAssistantMessage.id,
126
+
(message) => ({
127
+
...message,
128
+
text: `${message.text}${chunk}`,
129
+
}),
130
+
);
131
+
});
132
+
} catch (error) {
133
+
patchMessage(
134
+
queryClient,
135
+
userId,
136
+
chatId,
137
+
lastAssistantMessage.id,
138
+
(message) => ({
139
+
...message,
140
+
text: "Failed to stream response.",
141
+
}),
142
+
);
143
+
throw error;
144
+
} finally {
145
+
await queryClient.invalidateQueries({
146
+
queryKey: chatQueryKey(userId, chatId),
147
+
});
148
+
await queryClient.invalidateQueries({
149
+
queryKey: chatsQueryKey(userId),
150
+
});
151
+
}
152
+
},
153
+
});
154
+
}