this repo has no description
Add regenerate message button to chat messages
+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
+
}