✨ Retry messages (#121)

Co-authored-by: Victor Mustar <victor.mustar@gmail.com>

src/lib/components/chat/ChatMessage.svelte

@@ -1,11 +1,12 @@
 <script lang="ts">
 	import { marked } from "marked";
 	import type { Message } from "$lib/types/Message";
-	import { afterUpdate } from "svelte";
+	import { afterUpdate, createEventDispatcher } from "svelte";
 	import { deepestChild } from "$lib/utils/deepestChild";
 
 	import CodeBlock from "../CodeBlock.svelte";
 	import IconLoading from "../icons/IconLoading.svelte";
+	import CarbonRotate360 from "~icons/carbon/rotate-360";
 
 	function sanitizeMd(md: string) {
 		return md
@@ -25,6 +26,8 @@
 	export let message: Message;
 	export let loading: boolean = false;
 
+	const dispatch = createEventDispatcher<{ retry: void }>();
+
 	let contentEl: HTMLElement;
 	let loadingEl: any;
 	let pendingTimeout: NodeJS.Timeout;
@@ -69,7 +72,7 @@
 		<img
 			alt=""
 			src="https://huggingface.co/avatars/2edb18bd0206c16b433841a47f53fa8e.svg"
-			class="mt-5 h-3 w-3 flex-none rounded-full shadow-lg"
+			class="mt-5 h-3 w-3 flex-none select-none rounded-full shadow-lg"
 		/>
 		<div
 			class="relative min-h-[calc(2rem+theme(spacing[3.5])*2)] min-w-[100px] rounded-2xl border border-gray-100 bg-gradient-to-br from-gray-50 px-5 py-3.5 text-gray-600 prose-pre:my-2 dark:border-gray-800 dark:from-gray-800/40 dark:text-gray-300"
@@ -93,10 +96,20 @@
 	</div>
 {/if}
 {#if message.from === "user"}
-	<div class="flex items-start justify-start gap-4 max-sm:text-sm">
+	<div class="group relative flex items-start justify-start gap-4 max-sm:text-sm">
 		<div class="mt-5 h-3 w-3 flex-none rounded-full" />
 		<div class="whitespace-break-spaces rounded-2xl px-5 py-3.5 text-gray-500 dark:text-gray-400">
 			{message.content.trim()}
 		</div>
+		{#if !loading && message.id}
+			<button
+				class="absolute right-0 top-3.5 cursor-pointer rounded-lg border border-gray-100 p-1 text-xs text-gray-400 group-hover:block hover:text-gray-500 dark:border-gray-800 dark:text-gray-400 dark:hover:text-gray-300 md:hidden lg:-right-2"
+				title="Retry"
+				type="button"
+				on:click={() => dispatch("retry")}
+			>
+				<CarbonRotate360 />
+			</button>
+		{/if}
 	</div>
 {/if}

src/lib/components/chat/ChatMessages.svelte

@@ -2,10 +2,13 @@
 	import type { Message } from "$lib/types/Message";
 	import { snapScrollToBottom } from "$lib/actions/snapScrollToBottom";
 	import ScrollToBottomBtn from "$lib/components/ScrollToBottomBtn.svelte";
-	import { tick } from "svelte";
+	import { createEventDispatcher, tick } from "svelte";
 
 	import ChatIntroduction from "./ChatIntroduction.svelte";
 	import ChatMessage from "./ChatMessage.svelte";
+	import { randomUUID } from "$lib/utils/randomUuid";
+
+	const dispatch = createEventDispatcher<{ retry: { id: Message["id"]; content: string } }>();
 
 	export let messages: Message[];
 	export let loading: boolean;
@@ -31,12 +34,16 @@
 >
 	<div class="mx-auto flex h-full max-w-3xl flex-col gap-5 px-5 pt-6 sm:gap-8 xl:max-w-4xl">
 		{#each messages as message, i}
-			<ChatMessage loading={loading && i === messages.length - 1} {message} />
+			<ChatMessage
+				loading={loading && i === messages.length - 1}
+				{message}
+				on:retry={() => dispatch("retry", { id: message.id, content: message.content })}
+			/>
 		{:else}
 			<ChatIntroduction on:message />
 		{/each}
 		{#if pending}
-			<ChatMessage message={{ from: "assistant", content: "" }} />
+			<ChatMessage message={{ from: "assistant", content: "", id: randomUUID() }} />
 		{/if}
 		<div class="h-32 flex-none" />
 	</div>

src/lib/components/chat/ChatWindow.svelte

@@ -17,7 +17,12 @@
 
 	let message: string;
 
-	const dispatch = createEventDispatcher<{ message: string; share: void; stop: void }>();
+	const dispatch = createEventDispatcher<{
+		message: string;
+		share: void;
+		stop: void;
+		retry: { id: Message["id"]; content: string };
+	}>();
 
 	const handleSubmit = () => {
 		if (loading) return;
@@ -27,7 +32,15 @@
 </script>
 
 <div class="relative min-h-0 min-w-0">
-	<ChatMessages {loading} {pending} {messages} on:message />
+	<ChatMessages
+		{loading}
+		{pending}
+		{messages}
+		on:message
+		on:retry={(ev) => {
+			if (!loading) dispatch("retry", ev.detail);
+		}}
+	/>
 	<div
 		class="dark:via-gray-80 pointer-events-none absolute inset-x-0 bottom-0 z-0 mx-auto flex w-full max-w-3xl flex-col items-center justify-center bg-gradient-to-t from-white via-white/80 to-white/0 px-3.5 py-4 dark:border-gray-800 dark:from-gray-900 dark:to-gray-900/0 max-md:border-t max-md:bg-white max-md:dark:bg-gray-900 sm:px-5 md:py-8 xl:max-w-4xl [&>*]:pointer-events-auto"
 	>

src/lib/stores/pendingMessageIdToRetry.ts

@@ -0,0 +1,4 @@
+import type { Message } from "$lib/types/Message";
+import { writable } from "svelte/store";
+
+export const pendingMessageIdToRetry = writable<Message["id"] | null>(null);

src/lib/types/Message.ts

@@ -1,4 +1,5 @@
 export interface Message {
 	from: "user" | "assistant";
+	id: ReturnType<typeof crypto.randomUUID>;
 	content: string;
 }

src/lib/utils/randomUuid.ts

@@ -0,0 +1,14 @@
+type UUID = ReturnType<typeof crypto.randomUUID>;
+
+export function randomUUID(): UUID {
+	// Only on old safari / ios
+	if (!("randomUUID" in crypto)) {
+		return "10000000-1000-4000-8000-100000000000".replace(/[018]/g, (c) =>
+			(
+				Number(c) ^
+				(crypto.getRandomValues(new Uint8Array(1))[0] & (15 >> (Number(c) / 4)))
+			).toString(16)
+		) as UUID;
+	}
+	return crypto.randomUUID();
+}

src/routes/conversation/[id]/+page.svelte

@@ -1,6 +1,7 @@
 <script lang="ts">
 	import ChatWindow from "$lib/components/chat/ChatWindow.svelte";
 	import { pendingMessage } from "$lib/stores/pendingMessage";
+	import { pendingMessageIdToRetry } from "$lib/stores/pendingMessageIdToRetry";
 	import { onMount } from "svelte";
 	import { page } from "$app/stores";
 	import { textGenerationStream } from "@huggingface/inference";
@@ -10,6 +11,7 @@
 	import { shareConversation } from "$lib/shareConversation";
 	import { UrlDependency } from "$lib/types/UrlDependency";
 	import { error } from "$lib/stores/errors";
+	import { randomUUID } from "$lib/utils/randomUuid";
 
 	export let data;
 
@@ -26,7 +28,7 @@
 	let loading = false;
 	let pending = false;
 
-	async function getTextGenerationStream(inputs: string) {
+	async function getTextGenerationStream(inputs: string, messageId: string, isRetry = false) {
 		let conversationId = $page.params.id;
 
 		const response = textGenerationStream(
@@ -48,6 +50,8 @@
 				},
 			},
 			{
+				id: messageId,
+				is_retry: isRetry,
 				use_cache: false,
 			}
 		);
@@ -89,7 +93,11 @@
 
 				if (lastMessage?.from !== "assistant") {
 					// First token has a space at the beginning, trim it
-					messages = [...messages, { from: "assistant", content: data.token.text.trimStart() }];
+					messages = [
+						...messages,
+						// id doesn't match the backend id but it's not important for assistant messages
+						{ from: "assistant", content: data.token.text.trimStart(), id: randomUUID() },
+					];
 				} else {
 					lastMessage.content += data.token.text;
 					messages = [...messages];
@@ -104,7 +112,7 @@
 		});
 	}
 
-	async function writeMessage(message: string) {
+	async function writeMessage(message: string, messageId = crypto.randomUUID()) {
 		if (!message.trim()) return;
 
 		try {
@@ -112,9 +120,18 @@
 			loading = true;
 			pending = true;
 
-			messages = [...messages, { from: "user", content: message }];
+			let retryMessageIndex = messages.findIndex((msg) => msg.id === messageId);
+			const isRetry = retryMessageIndex !== -1;
+			if (!isRetry) {
+				retryMessageIndex = messages.length;
+			}
+
+			messages = [
+				...messages.slice(0, retryMessageIndex),
+				{ from: "user", content: message, id: messageId },
+			];
 
-			await getTextGenerationStream(message);
+			await getTextGenerationStream(message, messageId, isRetry);
 
 			if (messages.filter((m) => m.from === "user").length === 1) {
 				summarizeTitle($page.params.id)
@@ -135,9 +152,11 @@
 	onMount(async () => {
 		if ($pendingMessage) {
 			const val = $pendingMessage;
+			const messageId = $pendingMessageIdToRetry || undefined;
 			$pendingMessage = "";
+			$pendingMessageIdToRetry = null;
 
-			writeMessage(val);
+			writeMessage(val, messageId);
 		}
 	});
 
@@ -153,6 +172,7 @@
 	{pending}
 	{messages}
 	on:message={(message) => writeMessage(message.detail)}
+	on:retry={(message) => writeMessage(message.detail.content, message.detail.id)}
 	on:share={() => shareConversation($page.params.id, data.title)}
 	on:stop={() => (isAborted = true)}
 />

src/routes/conversation/[id]/+server.ts

@@ -27,10 +27,43 @@ export async function POST({ request, fetch, locals, params }) {
 		throw error(404, "Conversation not found");
 	}
 
-	// Todo: validate prompt with zod? or aktype
 	const json = await request.json();
-
-	const messages = [...conv.messages, { from: "user", content: json.inputs }] satisfies Message[];
+	const {
+		inputs: newPrompt,
+		options: { id: messageId, is_retry },
+	} = z
+		.object({
+			inputs: z.string().trim().min(1),
+			options: z.object({
+				id: z.optional(z.string().uuid()),
+				is_retry: z.optional(z.boolean()),
+			}),
+		})
+		.parse(json);
+
+	const messages = (() => {
+		if (is_retry && messageId) {
+			let retryMessageIdx = conv.messages.findIndex((message) => message.id === messageId);
+			if (retryMessageIdx === -1) {
+				retryMessageIdx = conv.messages.length;
+			}
+			return [
+				...conv.messages.slice(0, retryMessageIdx),
+				{ content: newPrompt, from: "user", id: messageId as Message["id"] },
+			];
+		}
+		return [
+			...conv.messages,
+			{ content: newPrompt, from: "user", id: (messageId as Message["id"]) || crypto.randomUUID() },
+		];
+	})() satisfies Message[];
+
+	// Todo: on-the-fly migration, remove later
+	for (const message of messages) {
+		if (!message.id) {
+			message.id = crypto.randomUUID();
+		}
+	}
 	const prompt = buildPrompt(messages);
 
 	const randomEndpoint = modelEndpoint();
@@ -62,7 +95,7 @@ export async function POST({ request, fetch, locals, params }) {
 
 		generated_text = trimSuffix(trimPrefix(generated_text, "<|startoftext|>"), PUBLIC_SEP_TOKEN);
 
-		messages.push({ from: "assistant", content: generated_text });
+		messages.push({ from: "assistant", content: generated_text, id: crypto.randomUUID() });
 
 		await collections.conversations.updateOne(
 			{

src/routes/r/[id]/+page.svelte

@@ -5,13 +5,14 @@
 	import ChatWindow from "$lib/components/chat/ChatWindow.svelte";
 	import { ERROR_MESSAGES, error } from "$lib/stores/errors";
 	import { pendingMessage } from "$lib/stores/pendingMessage";
+	import { pendingMessageIdToRetry } from "$lib/stores/pendingMessageIdToRetry";
 	import type { PageData } from "./$types";
 
 	export let data: PageData;
 
 	let loading = false;
 
-	async function createConversation(message: string) {
+	async function createConversation() {
 		try {
 			loading = true;
 			const res = await fetch(`${base}/conversation`, {
@@ -32,16 +33,11 @@
 
 			const { conversationId } = await res.json();
 
-			// Ugly hack to use a store as temp storage, feel free to improve ^^
-			pendingMessage.set(message);
-
-			// invalidateAll to update list of conversations
-			await goto(`${base}/conversation/${conversationId}`, { invalidateAll: true });
+			return conversationId;
 		} catch (err) {
 			error.set(ERROR_MESSAGES.default);
 			console.error(String(err));
-		} finally {
-			loading = false;
+			throw err;
 		}
 	}
 
@@ -65,8 +61,22 @@
 </svelte:head>
 
 <ChatWindow
-	on:message={(ev) => createConversation(ev.detail)}
+	on:message={(ev) =>
+		createConversation()
+			.then((convId) => {
+				$pendingMessage = ev.detail;
+				return goto(`${base}/conversation/${convId}`, { invalidateAll: true });
+			})
+			.finally(() => (loading = false))}
 	on:share={shareConversation}
+	on:retry={(ev) =>
+		createConversation()
+			.then((convId) => {
+				$pendingMessageIdToRetry = ev.detail.id;
+				$pendingMessage = ev.detail.content;
+				return goto(`${base}/conversation/${convId}`, { invalidateAll: true });
+			})
+			.finally(() => (loading = false))}
 	messages={data.messages}
 	{loading}
 />