Conversation trees (#223) (#807)

* work on branching

* branching

* work on input field

* wip

* wip

* helper stuff

* pass tests

* add type guards and clean up type requirements

* fixed shared conv type

* more helpers stuff

* clean up code

* addChildren helper

* addSibling

* add test for addSibling & refacto addChildren tests

* backend work pt. 1

* backend done

* add children property to messages for easier frontend rendering

* fix message id type rail

* front-end work on simple linear conversation

* fix title generation

* convert conversations on post

* server side retry works with branching

* clean up buttons

* fix retry feature backend

* Send new messages in any subtree

* make edit previous prompts feature work

* fix padding

* revert unneeded changes

* bring back pending message

* fix front-end streaming

* fix initial message

* Revert "fix initial message"

This reverts commit 6257fe8789c52a20e7506aa453efda55c16e95bb.

* Fix bug subtree state refresh

* fix continue feature on shared conversations

* fix websearch updates

* Fix first message streaming

* fix fornt-end websearch updates

* fix loading icon

* fix bottom padding

* move children nav to below the message

* Show current message in continue & retry

* children nav styling

* you can now edit the first message

* bottom padding on assistant message

* fix test

* lint

* use <form>

* tree navigation

* misc

* mobile: hide download link

* forgot to implem continue feature lol

* fix continue feature on llama 70b

* fix edit mode

* disable submit button & nav when loading

* fix bug when interrupting

* hide arrows in edit mode

* forgot to reset edit mode when submitting retry

* reset editing when switching conversations

---------

Co-authored-by: Victor Mustar <[email protected]>

.env.template

@@ -33,10 +33,6 @@ MODELS=`[
       "name": "meta-llama/Llama-2-70b-chat-hf",
       "description": "The latest and biggest model from Meta, fine-tuned for chat.",
       "websiteUrl": "https://ai.meta.com/llama/",
-      "userMessageToken": "",
-      "userMessageEndToken": " [/INST] ",
-      "assistantMessageToken": "",
-      "assistantMessageEndToken": " </s><s>[INST] ",
       "preprompt": " ",
       "chatPromptTemplate" : "<s>[INST] <<SYS>>\n{{preprompt}}\n<</SYS>>\n\n{{#each messages}}{{#ifUser}}{{content}} [/INST] {{/ifUser}}{{#ifAssistant}}{{content}} </s><s>[INST] {{/ifAssistant}}{{/each}}",
       "promptExamples": [
@@ -58,7 +54,7 @@ MODELS=`[
         "top_k": 50,
         "truncate": 3072,
         "max_new_tokens": 1024,
-        "stop" : ["</s>", " </s><s>[INST] "]
+        "stop" : ["</s>", "</s><s>[INST]"]
       }
     },
     {

src/lib/buildPrompt.ts

@@ -1,94 +1,31 @@
+import type { EndpointParameters } from "./server/endpoints/endpoints";
 import type { BackendModel } from "./server/models";
-import type { Message } from "./types/Message";
-import { format } from "date-fns";
-import type { WebSearch } from "./types/WebSearch";
-import { downloadFile } from "./server/files/downloadFile";
-import type { Conversation } from "./types/Conversation";
 
-interface buildPromptOptions {
-	messages: Pick<Message, "from" | "content" | "files">[];
-	id?: Conversation["_id"];
+type buildPromptOptions = Pick<EndpointParameters, "messages" | "preprompt" | "continueMessage"> & {
 	model: BackendModel;
-	locals?: App.Locals;
-	webSearch?: WebSearch;
-	preprompt?: string;
-	files?: File[];
-	continue?: boolean;
-}
+};
 
 export async function buildPrompt({
 	messages,
 	model,
-	webSearch,
 	preprompt,
-	id,
+	continueMessage,
 }: buildPromptOptions): Promise<string> {
-	let modifiedMessages = [...messages];
-
-	if (webSearch && webSearch.context) {
-		// find index of the last user message
-		const lastUsrMsgIndex = modifiedMessages.map((el) => el.from).lastIndexOf("user");
-
-		// combine all the other previous questions into one string
-		const previousUserMessages = modifiedMessages.filter((el) => el.from === "user").slice(0, -1);
-		const previousQuestions =
-			previousUserMessages.length > 0
-				? `Previous questions: \n${previousUserMessages
-						.map(({ content }) => `- ${content}`)
-						.join("\n")}`
-				: "";
-
-		const currentDate = format(new Date(), "MMMM d, yyyy");
-
-		// update the last user message directly (that way if the last message is an assistant partial answer, we keep the beginning of that answer)
-		modifiedMessages[lastUsrMsgIndex] = {
-			from: "user",
-			content: `I searched the web using the query: ${webSearch.searchQuery}. Today is ${currentDate} and here are the results:
-=====================
-${webSearch.context}
-=====================
-${previousQuestions}
-Answer the question: ${messages[lastUsrMsgIndex].content}`,
-		};
-	}
-	// section to handle potential files input
-	if (model.multimodal) {
-		modifiedMessages = await Promise.all(
-			modifiedMessages.map(async (el) => {
-				let content = el.content;
-
-				if (el.from === "user") {
-					if (el?.files && el.files.length > 0 && id) {
-						const markdowns = await Promise.all(
-							el.files.map(async (hash) => {
-								try {
-									const { content: image, mime } = await downloadFile(hash, id);
-									const b64 = image.toString("base64");
-									return `![](data:${mime};base64,${b64})})`;
-								} catch (e) {
-									console.error(e);
-								}
-							})
-						);
-						content += markdowns.join("\n ");
-					} else {
-						// if no image, append an empty white image
-						content +=
-							"\n![](data:image/png;base64,/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/2wBDAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/wAARCAAQABADAREAAhEBAxEB/8QAHwAAAQUBAQEBAQEAAAAAAAAAAAECAwQFBgcICQoL/8QAtRAAAgEDAwIEAwUFBAQAAAF9AQIDAAQRBRIhMUEGE1FhByJxFDKBkaEII0KxwRVS0fAkM2JyggkKFhcYGRolJicoKSo0NTY3ODk6Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqDhIWGh4iJipKTlJWWl5iZmqKjpKWmp6ipqrKztLW2t7i5usLDxMXGx8jJytLT1NXW19jZ2uHi4+Tl5ufo6erx8vP09fb3+Pn6/8QAHwEAAwEBAQEBAQEBAQAAAAAAAAECAwQFBgcICQoL/8QAtREAAgECBAQDBAcFBAQAAQJ3AAECAxEEBSExBhJBUQdhcRMiMoEIFEKRobHBCSMzUvAVYnLRChYkNOEl8RcYGRomJygpKjU2Nzg5OkNERUZHSElKU1RVVldYWVpjZGVmZ2hpanN0dXZ3eHl6goOEhYaHiImKkpOUlZaXmJmaoqOkpaanqKmqsrO0tba3uLm6wsPExcbHyMnK0tPU1dbX2Nna4uPk5ebn6Onq8vP09fb3+Pn6/9oADAMBAAIRAxEAPwD+/igAoAKACgD/2Q==)";
-					}
-				}
-
-				return { ...el, content };
-			})
-		);
+	let prompt = model
+		.chatPromptRender({ messages, preprompt })
+		// Not super precise, but it's truncated in the model's backend anyway
+		.split(" ")
+		.slice(-(model.parameters?.truncate ?? 0))
+		.join(" ");
+
+	if (continueMessage && model.parameters?.stop) {
+		prompt = model.parameters.stop.reduce((acc: string, curr: string) => {
+			if (acc.endsWith(curr)) {
+				return acc.slice(0, acc.length - curr.length);
+			}
+			return acc;
+		}, prompt.trimEnd());
 	}
 
-	return (
-		model
-			.chatPromptRender({ messages: modifiedMessages, preprompt })
-			// Not super precise, but it's truncated in the model's backend anyway
-			.split(" ")
-			.slice(-(model.parameters?.truncate ?? 0))
-			.join(" ")
-	);
+	return prompt;
 }

src/lib/components/chat/ChatMessage.svelte

@@ -2,7 +2,7 @@
 	import { marked } from "marked";
 	import markedKatex from "marked-katex-extension";
 	import type { Message } from "$lib/types/Message";
-	import { afterUpdate, createEventDispatcher } from "svelte";
+	import { afterUpdate, createEventDispatcher, tick } from "svelte";
 	import { deepestChild } from "$lib/utils/deepestChild";
 	import { page } from "$app/stores";
 
@@ -13,6 +13,9 @@
 	import CarbonDownload from "~icons/carbon/download";
 	import CarbonThumbsUp from "~icons/carbon/thumbs-up";
 	import CarbonThumbsDown from "~icons/carbon/thumbs-down";
+	import CarbonPen from "~icons/carbon/pen";
+	import CarbonChevronLeft from "~icons/carbon/chevron-left";
+	import CarbonChevronRight from "~icons/carbon/chevron-right";
 
 	import { PUBLIC_SEP_TOKEN } from "$lib/constants/publicSepToken";
 	import type { Model } from "$lib/types/Model";
@@ -20,6 +23,7 @@
 	import OpenWebSearchResults from "../OpenWebSearchResults.svelte";
 	import type { WebSearchUpdate } from "$lib/types/MessageUpdate";
 	import { base } from "$app/paths";
+	import { useConvTreeStore } from "$lib/stores/convTree";
 
 	function sanitizeMd(md: string) {
 		let ret = md
@@ -45,16 +49,17 @@
 	}
 
 	export let model: Model;
-	export let message: Message;
+	export let id: Message["id"];
+	export let messages: Message[];
 	export let loading = false;
 	export let isAuthor = true;
 	export let readOnly = false;
 	export let isTapped = false;
 
-	export let webSearchMessages: WebSearchUpdate[];
+	$: message = messages.find((m) => m.id === id) ?? ({} as Message);
 
 	const dispatch = createEventDispatcher<{
-		retry: { content: string; id: Message["id"] };
+		retry: { content?: string; id: Message["id"] };
 		vote: { score: Message["score"]; id: Message["id"] };
 	}>();
 
@@ -63,6 +68,8 @@
 	let pendingTimeout: ReturnType<typeof setTimeout>;
 	let isCopied = false;
 
+	let initialized = false;
+
 	const renderer = new marked.Renderer();
 	// For code blocks with simple backticks
 	renderer.codespan = (code) => {
@@ -91,12 +98,15 @@
 
 	$: tokens = marked.lexer(sanitizeMd(message.content));
 
+	$: emptyLoad =
+		!message.content && (webSearchIsDone || (searchUpdates && searchUpdates.length === 0));
+
 	afterUpdate(() => {
 		loadingEl?.$destroy();
 		clearTimeout(pendingTimeout);
 
 		// Add loading animation to the last message if update takes more than 600ms
-		if (loading) {
+		if ((loading && isLast) || emptyLoad) {
 			pendingTimeout = setTimeout(() => {
 				if (contentEl) {
 					loadingEl = new IconLoading({
@@ -108,11 +118,14 @@
 		}
 	});
 
-	let searchUpdates: WebSearchUpdate[] = [];
+	function handleKeyDown(e: KeyboardEvent) {
+		if (e.key === "Enter" && (e.metaKey || e.ctrlKey)) {
+			editFormEl.requestSubmit();
+		}
+	}
 
-	$: searchUpdates = ((webSearchMessages.length > 0
-		? webSearchMessages
-		: message.updates?.filter(({ type }) => type === "webSearch")) ?? []) as WebSearchUpdate[];
+	$: searchUpdates = (message.updates?.filter(({ type }) => type === "webSearch") ??
+		[]) as WebSearchUpdate[];
 
 	$: downloadLink =
 		message.from === "user" ? `${$page.url.pathname}/message/${message.id}/prompt` : undefined;
@@ -131,11 +144,40 @@
 			isCopied = false;
 		}, 1000);
 	}
+
+	$: editMode = $convTreeStore.editing === message.id;
+	let editContentEl: HTMLTextAreaElement;
+	let editFormEl: HTMLFormElement;
+
+	$: if (editMode) {
+		tick();
+		if (editContentEl) {
+			editContentEl.value = message.content;
+			editContentEl?.focus();
+		}
+	}
+
+	$: isLast = (message && message.children?.length === 0) ?? false;
+
+	$: childrenToRender = 0;
+	$: nChildren = message?.children?.length ?? 0;
+
+	$: {
+		if (initialized) {
+			childrenToRender = Math.max(0, nChildren - 1);
+		} else {
+			childrenToRender = 0;
+			initialized = true;
+		}
+	}
+	const convTreeStore = useConvTreeStore();
+
+	$: if (message.children?.length === 0) $convTreeStore.leaf = message.id;
 </script>
 
 {#if message.from === "assistant"}
 	<div
-		class="group relative -mb-8 flex items-start justify-start gap-4 pb-8 leading-relaxed"
+		class="group relative -mb-6 flex items-start justify-start gap-4 pb-4 leading-relaxed"
 		role="presentation"
 		on:click={() => (isTapped = !isTapped)}
 		on:keydown={() => (isTapped = !isTapped)}
@@ -162,9 +204,6 @@
 					webSearchMessages={searchUpdates}
 				/>
 			{/if}
-			{#if !message.content && (webSearchIsDone || (webSearchMessages && webSearchMessages.length === 0))}
-				<IconLoading />
-			{/if}
 
 			<div
 				class="prose max-w-none max-sm:prose-sm dark:prose-invert prose-headings:font-semibold prose-h1:text-lg prose-h2:text-base prose-h3:text-base prose-pre:bg-gray-800 dark:prose-pre:bg-gray-900"
@@ -179,6 +218,7 @@
 					{/if}
 				{/each}
 			</div>
+
 			<!-- Web Search sources -->
 			{#if webSearchSources?.length}
 				<div class="mt-4 flex flex-wrap items-center gap-x-2 gap-y-1.5 text-sm">
@@ -202,7 +242,7 @@
 		</div>
 		{#if isAuthor && !loading && message.content}
 			<div
-				class="absolute bottom-1 right-0 flex max-md:transition-all md:bottom-0 md:group-hover:visible md:group-hover:opacity-100
+				class="absolute bottom-1 right-0 -mb-4 flex max-md:transition-all md:bottom-0 md:group-hover:visible md:group-hover:opacity-100
 					{message.score ? 'visible opacity-100' : 'invisible max-md:-translate-y-4 max-md:opacity-0'}
 					{isTapped || isCopied ? 'max-md:visible max-md:translate-y-0 max-md:opacity-100' : ''}
 				"
@@ -230,6 +270,14 @@
 				>
 					<CarbonThumbsDown class="h-[1.14em] w-[1.14em]" />
 				</button>
+				<button
+					class="btn rounded-sm p-1 text-sm text-gray-400 focus:ring-0 hover:text-gray-500 dark:text-gray-400 dark:hover:text-gray-300"
+					title="Retry"
+					type="button"
+					on:click={() => dispatch("retry", { id: message.id })}
+				>
+					<CarbonRotate360 />
+				</button>
 				<CopyToClipBoardBtn
 					on:click={() => {
 						isCopied = true;
@@ -240,10 +288,16 @@
 			</div>
 		{/if}
 	</div>
+	<slot name="childrenNav" />
 {/if}
 {#if message.from === "user"}
-	<div class="group relative flex items-start justify-start gap-4 max-sm:text-sm">
-		<div class="flex flex-col">
+	<div
+		class="group relative w-full items-start justify-start gap-4 max-sm:text-sm"
+		role="presentation"
+		on:click={() => (isTapped = !isTapped)}
+		on:keydown={() => (isTapped = !isTapped)}
+	>
+		<div class="flex w-full flex-col">
 			{#if message.files && message.files.length > 0}
 				<div class="mx-auto grid w-fit grid-cols-2 gap-5 px-5">
 					{#each message.files as file}
@@ -266,36 +320,133 @@
 				</div>
 			{/if}
 
-			<div
-				class="max-w-full whitespace-break-spaces break-words rounded-2xl px-5 py-3.5 text-gray-500 dark:text-gray-400"
-			>
-				{message.content.trim()}
+			<div class="flex w-full flex-row flex-nowrap">
+				{#if !editMode}
+					<p
+						class="disabled w-full appearance-none whitespace-break-spaces text-wrap break-words bg-inherit px-5 py-3.5 text-gray-500 dark:text-gray-400"
+					>
+						{message.content.trim()}
+					</p>
+				{:else}
+					<form
+						class="flex w-full flex-col"
+						bind:this={editFormEl}
+						on:submit|preventDefault={() => {
+							dispatch("retry", { content: editContentEl.value, id: message.id });
+							$convTreeStore.editing = null;
+						}}
+					>
+						<textarea
+							class="w-full whitespace-break-spaces break-words rounded-lg bg-gray-100 px-5 py-3.5 text-gray-500 *:h-max dark:bg-gray-800 dark:text-gray-400"
+							bind:this={editContentEl}
+							value={message.content.trim()}
+							on:keydown={handleKeyDown}
+							required
+						/>
+						<div class="flex w-full flex-row flex-nowrap items-center justify-center gap-2 pt-2">
+							<button
+								type="submit"
+								class="btn rounded-lg px-3 py-1.5 text-sm
+								{loading
+									? 'bg-gray-300 text-gray-400 dark:bg-gray-700 dark:text-gray-600'
+									: 'bg-gray-200 text-gray-600 focus:ring-0   hover:text-gray-800 dark:bg-gray-800 dark:text-gray-300 dark:hover:text-gray-200'}
+								"
+								disabled={loading}
+							>
+								Submit
+							</button>
+							<button
+								type="button"
+								class="btn rounded-sm p-2 text-sm text-gray-400 focus:ring-0 hover:text-gray-500 dark:text-gray-400 dark:hover:text-gray-300"
+								on:click={() => {
+									$convTreeStore.editing = null;
+								}}
+							>
+								Cancel
+							</button>
+						</div>
+					</form>
+				{/if}
+				{#if !loading && !editMode}
+					<div
+						class="
+						max-md:opacity-0' invisible absolute
+						right-0 top-3.5 z-10 h-max max-md:-translate-y-4 max-md:transition-all md:bottom-0 md:group-hover:visible md:group-hover:opacity-100 {isTapped ||
+						isCopied
+							? 'max-md:visible max-md:translate-y-0 max-md:opacity-100'
+							: ''}"
+					>
+						<div class="mx-auto flex flex-row flex-nowrap gap-2">
+							{#if downloadLink}
+								<a
+									class="rounded-lg border border-gray-100 bg-gray-100 p-1 text-xs text-gray-400 group-hover:block hover:text-gray-500 max-sm:!hidden md:hidden dark:border-gray-800 dark:bg-gray-800 dark:text-gray-400 dark:hover:text-gray-300"
+									title="Download prompt and parameters"
+									type="button"
+									target="_blank"
+									href={downloadLink}
+								>
+									<CarbonDownload />
+								</a>
+							{/if}
+							{#if !readOnly}
+								<button
+									class="cursor-pointer rounded-lg border border-gray-100 bg-gray-100 p-1 text-xs text-gray-400 group-hover:block hover:text-gray-500 md:hidden lg:-right-2 dark:border-gray-800 dark:bg-gray-800 dark:text-gray-400 dark:hover:text-gray-300"
+									title="Branch"
+									type="button"
+									on:click={() => ($convTreeStore.editing = message.id)}
+								>
+									<CarbonPen />
+								</button>
+							{/if}
+						</div>
+					</div>
+				{/if}
 			</div>
-			{#if !loading}
-				<div class="absolute right-0 top-3.5 flex gap-2 lg:-right-2">
-					{#if downloadLink}
-						<a
-							class="rounded-lg border border-gray-100 p-1 text-xs text-gray-400 group-hover:block hover:text-gray-500 md:hidden dark:border-gray-800 dark:text-gray-400 dark:hover:text-gray-300"
-							title="Download prompt and parameters"
-							type="button"
-							target="_blank"
-							href={downloadLink}
-						>
-							<CarbonDownload />
-						</a>
-					{/if}
-					{#if !readOnly}
-						<button
-							class="cursor-pointer rounded-lg border border-gray-100 p-1 text-xs text-gray-400 group-hover:block hover:text-gray-500 md:hidden lg:-right-2 dark:border-gray-800 dark:text-gray-400 dark:hover:text-gray-300"
-							title="Retry"
-							type="button"
-							on:click={() => dispatch("retry", { content: message.content, id: message.id })}
-						>
-							<CarbonRotate360 />
-						</button>
-					{/if}
-				</div>
-			{/if}
+			<slot name="childrenNav" />
 		</div>
 	</div>
 {/if}
+
+{#if nChildren > 0}
+	<svelte:self
+		{loading}
+		{messages}
+		{isAuthor}
+		{readOnly}
+		{model}
+		id={messages.find((m) => m.id === id)?.children?.[childrenToRender]}
+		on:retry
+		on:vote
+		on:continue
+	>
+		<svelte:fragment slot="childrenNav">
+			{#if nChildren > 1 && $convTreeStore.editing === null}
+				<div
+					class="font-white z-10 -mt-1 ml-3.5 mr-auto flex h-6 w-fit select-none flex-row items-center justify-center gap-1 text-sm"
+				>
+					<button
+						class="inline text-lg font-thin text-gray-400 disabled:pointer-events-none disabled:opacity-25 hover:text-gray-800 dark:text-gray-500 dark:hover:text-gray-200"
+						on:click={() => (childrenToRender = Math.max(0, childrenToRender - 1))}
+						disabled={childrenToRender === 0 || loading}
+					>
+						<CarbonChevronLeft class="text-sm" />
+					</button>
+					<span class=" text-gray-400 dark:text-gray-500">
+						{childrenToRender + 1} / {nChildren}
+					</span>
+					<button
+						class="inline text-lg font-thin text-gray-400 disabled:pointer-events-none disabled:opacity-25 hover:text-gray-800 dark:text-gray-500 dark:hover:text-gray-200"
+						on:click={() =>
+							(childrenToRender = Math.min(
+								message?.children?.length ?? 1 - 1,
+								childrenToRender + 1
+							))}
+						disabled={childrenToRender === nChildren - 1 || loading}
+					>
+						<CarbonChevronRight class="text-sm" />
+					</button>
+				</div>
+			{/if}
+		</svelte:fragment>
+	</svelte:self>
+{/if}

src/lib/components/chat/ChatMessages.svelte

@@ -1,106 +0,0 @@
-<script lang="ts">
-	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 { randomUUID } from "$lib/utils/randomUuid";
-	import type { Model } from "$lib/types/Model";
-	import ChatIntroduction from "./ChatIntroduction.svelte";
-	import ChatMessage from "./ChatMessage.svelte";
-	import type { WebSearchUpdate } from "$lib/types/MessageUpdate";
-	import { browser } from "$app/environment";
-	import SystemPromptModal from "../SystemPromptModal.svelte";
-	import type { Assistant } from "$lib/types/Assistant";
-	import AssistantIntroduction from "./AssistantIntroduction.svelte";
-	import { page } from "$app/stores";
-	import { base } from "$app/paths";
-
-	export let messages: Message[];
-	export let loading: boolean;
-	export let pending: boolean;
-	export let isAuthor: boolean;
-	export let currentModel: Model;
-	export let assistant: Assistant | undefined;
-	export let models: Model[];
-	export let preprompt: string | undefined;
-	export let readOnly: boolean;
-
-	let chatContainer: HTMLElement;
-
-	export let webSearchMessages: WebSearchUpdate[] = [];
-
-	async function scrollToBottom() {
-		await tick();
-		chatContainer.scrollTop = chatContainer.scrollHeight;
-	}
-
-	// If last message is from user, scroll to bottom
-	$: if (browser && messages[messages.length - 1]?.from === "user") {
-		scrollToBottom();
-	}
-</script>
-
-<div
-	class="scrollbar-custom mr-1 h-full overflow-y-auto"
-	use:snapScrollToBottom={messages.length ? [...messages, ...webSearchMessages] : false}
-	bind:this={chatContainer}
->
-	<div class="mx-auto flex h-full max-w-3xl flex-col gap-6 px-5 pt-6 sm:gap-8 xl:max-w-4xl">
-		{#each messages as message, i}
-			{#if i === 0 && $page.data?.assistant}
-				<a
-					class="mx-auto flex items-center gap-1.5 rounded-full border border-gray-100 bg-gray-50 py-1 pl-1 pr-3 text-sm text-gray-800 hover:bg-gray-100 dark:border-gray-800 dark:bg-gray-800 dark:text-gray-200 dark:hover:bg-gray-700"
-					href="{base}/settings/assistants/{$page.data.assistant._id}"
-				>
-					{#if $page.data?.assistant.avatar}
-						<img
-							src="{base}/settings/assistants/{$page.data?.assistant._id.toString()}/avatar.jpg?hash=${$page
-								.data.assistant.avatar}"
-							alt="Avatar"
-							class="size-5 rounded-full object-cover"
-						/>
-					{:else}
-						<div
-							class="flex size-6 items-center justify-center rounded-full bg-gray-300 font-bold uppercase text-gray-500"
-						>
-							{$page.data?.assistant.name[0]}
-						</div>
-					{/if}
-
-					{$page.data.assistant.name}
-				</a>
-			{:else if i === 0 && preprompt && preprompt != currentModel.preprompt}
-				<SystemPromptModal preprompt={preprompt ?? ""} />
-			{/if}
-			<ChatMessage
-				loading={loading && i === messages.length - 1}
-				{message}
-				{isAuthor}
-				{readOnly}
-				model={currentModel}
-				webSearchMessages={i === messages.length - 1 ? webSearchMessages : []}
-				on:retry
-				on:vote
-				on:continue
-			/>
-		{:else}
-			{#if !assistant}
-				<ChatIntroduction {models} {currentModel} on:message />
-			{:else}
-				<AssistantIntroduction {assistant} on:message />
-			{/if}
-		{/each}
-		{#if pending && messages[messages.length - 1]?.from === "user"}
-			<ChatMessage
-				message={{ from: "assistant", content: "", id: randomUUID() }}
-				model={currentModel}
-				{webSearchMessages}
-			/>
-		{/if}
-		<div class="h-44 flex-none" />
-	</div>
-	<ScrollToBottomBtn
-		class="bottom-36 right-4 max-md:hidden lg:right-10"
-		scrollNode={chatContainer}
-	/>
-</div>

src/lib/components/chat/ChatWindow.svelte

@@ -1,6 +1,6 @@
 <script lang="ts">
 	import type { Message } from "$lib/types/Message";
-	import { createEventDispatcher, onDestroy } from "svelte";
+	import { createEventDispatcher, onDestroy, tick } from "svelte";
 
 	import CarbonSendAltFilled from "~icons/carbon/send-alt-filled";
 	import CarbonExport from "~icons/carbon/export";
@@ -11,13 +11,11 @@
 
 	import EosIconsLoading from "~icons/eos-icons/loading";
 
-	import ChatMessages from "./ChatMessages.svelte";
 	import ChatInput from "./ChatInput.svelte";
 	import StopGeneratingBtn from "../StopGeneratingBtn.svelte";
 	import type { Model } from "$lib/types/Model";
 	import WebSearchToggle from "../WebSearchToggle.svelte";
 	import LoginModal from "../LoginModal.svelte";
-	import type { WebSearchUpdate } from "$lib/types/MessageUpdate";
 	import { page } from "$app/stores";
 	import FileDropzone from "./FileDropzone.svelte";
 	import RetryBtn from "../RetryBtn.svelte";
@@ -26,15 +24,23 @@
 	import type { Assistant } from "$lib/types/Assistant";
 	import { base } from "$app/paths";
 	import ContinueBtn from "../ContinueBtn.svelte";
+	import AssistantIntroduction from "./AssistantIntroduction.svelte";
+	import ChatMessage from "./ChatMessage.svelte";
+	import ScrollToBottomBtn from "../ScrollToBottomBtn.svelte";
+	import { browser } from "$app/environment";
+	import { snapScrollToBottom } from "$lib/actions/snapScrollToBottom";
+	import SystemPromptModal from "../SystemPromptModal.svelte";
+	import ChatIntroduction from "./ChatIntroduction.svelte";
+	import { useConvTreeStore } from "$lib/stores/convTree";
 
 	export let messages: Message[] = [];
 	export let loading = false;
 	export let pending = false;
+
 	export let shared = false;
 	export let currentModel: Model;
 	export let models: Model[];
 	export let assistant: Assistant | undefined = undefined;
-	export let webSearchMessages: WebSearchUpdate[] = [];
 	export let preprompt: string | undefined = undefined;
 	export let files: File[] = [];
 
@@ -50,7 +56,7 @@
 		message: string;
 		share: void;
 		stop: void;
-		retry: { id: Message["id"]; content: string };
+		retry: { id: Message["id"]; content?: string };
 		continue: { id: Message["id"] };
 	}>();
 
@@ -76,7 +82,11 @@
 	const onDragOver = (e: DragEvent) => {
 		e.preventDefault();
 	};
-	$: lastIsError = messages[messages.length - 1]?.from === "user" && !loading;
+
+	const convTreeStore = useConvTreeStore();
+
+	$: lastMessage = browser && (messages.find((m) => m.id == $convTreeStore.leaf) as Message);
+	$: lastIsError = lastMessage && lastMessage.from === "user" && !loading;
 
 	$: sources = files.map((file) => file2base64(file));
 
@@ -96,6 +106,18 @@
 			clearTimeout(timeout);
 		}
 	});
+
+	let chatContainer: HTMLElement;
+
+	async function scrollToBottom() {
+		await tick();
+		chatContainer.scrollTop = chatContainer.scrollHeight;
+	}
+
+	// If last message is from user, scroll to bottom
+	$: if (lastMessage && lastMessage.from === "user") {
+		scrollToBottom();
+	}
 </script>
 
 <div class="relative min-h-0 min-w-0">
@@ -106,31 +128,79 @@
 			}}
 		/>
 	{/if}
-	<ChatMessages
-		{loading}
-		{pending}
-		{currentModel}
-		{models}
-		{assistant}
-		{messages}
-		readOnly={isReadOnly}
-		isAuthor={!shared}
-		{webSearchMessages}
-		{preprompt}
-		on:message={(ev) => {
-			if ($page.data.loginRequired) {
-				loginModalOpen = true;
-			} else {
-				dispatch("message", ev.detail);
-			}
-		}}
-		on:vote
-		on:continue
-		on:retry={(ev) => {
-			if (!loading) dispatch("retry", ev.detail);
-		}}
-	/>
+	<div
+		class="scrollbar-custom mr-1 h-full overflow-y-auto"
+		use:snapScrollToBottom={messages.length ? [...messages] : false}
+		bind:this={chatContainer}
+	>
+		<div class="mx-auto flex h-full max-w-3xl flex-col gap-6 px-5 pt-6 sm:gap-8 xl:max-w-4xl">
+			{#if $page.data?.assistant}
+				<a
+					class="mx-auto flex items-center gap-1.5 rounded-full border border-gray-100 bg-gray-50 py-1 pl-1 pr-3 text-sm text-gray-800 hover:bg-gray-100 dark:border-gray-800 dark:bg-gray-800 dark:text-gray-200 dark:hover:bg-gray-700"
+					href="{base}/settings/assistants/{$page.data.assistant._id}"
+				>
+					{#if $page.data?.assistant.avatar}
+						<img
+							src="{base}/settings/assistants/{$page.data?.assistant._id.toString()}/avatar.jpg?hash=${$page
+								.data.assistant.avatar}"
+							alt="Avatar"
+							class="size-5 rounded-full object-cover"
+						/>
+					{:else}
+						<div
+							class="flex size-6 items-center justify-center rounded-full bg-gray-300 font-bold uppercase text-gray-500"
+						>
+							{$page.data?.assistant.name[0]}
+						</div>
+					{/if}
+
+					{$page.data.assistant.name}
+				</a>
+			{:else if preprompt && preprompt != currentModel.preprompt}
+				<SystemPromptModal preprompt={preprompt ?? ""} />
+			{/if}
 
+			{#if messages.length > 0}
+				<div class="flex h-max flex-col gap-6 pb-52">
+					<ChatMessage
+						{loading}
+						{messages}
+						id={messages[0].id}
+						isAuthor={!shared}
+						readOnly={isReadOnly}
+						model={currentModel}
+						on:retry
+						on:vote
+						on:continue
+					/>
+				</div>
+			{:else if pending}
+				<ChatMessage
+					loading={true}
+					messages={[
+						{
+							id: "0-0-0-0-0",
+							content: "",
+							from: "assistant",
+							children: [],
+						},
+					]}
+					id={"0-0-0-0-0"}
+					isAuthor={!shared}
+					readOnly={isReadOnly}
+					model={currentModel}
+				/>
+			{:else if !assistant}
+				<ChatIntroduction {models} {currentModel} on:message />
+			{:else}
+				<AssistantIntroduction {assistant} on:message />
+			{/if}
+		</div>
+		<ScrollToBottomBtn
+			class="bottom-36 right-4 max-md:hidden lg:right-10"
+			scrollNode={chatContainer}
+		/>
+	</div>
 	<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 max-md:border-t max-md:bg-white sm:px-5 md:py-8 xl:max-w-4xl dark:border-gray-800 dark:from-gray-900 dark:to-gray-900/0 max-md:dark:bg-gray-900 [&>*]:pointer-events-auto"
 	>
@@ -169,23 +239,28 @@
 				{:else if lastIsError}
 					<RetryBtn
 						classNames="ml-auto"
-						on:click={() =>
-							dispatch("retry", {
-								id: messages[messages.length - 1].id,
-								content: messages[messages.length - 1].content,
-							})}
+						on:click={() => {
+							if (lastMessage && lastMessage.ancestors) {
+								dispatch("retry", {
+									id: lastMessage.ancestors[lastMessage.ancestors.length - 1],
+								});
+							}
+						}}
 					/>
 				{:else}
 					<div class="ml-auto gap-2">
 						{#if currentModel.multimodal}
 							<UploadBtn bind:files classNames="ml-auto" />
 						{/if}
-						{#if messages && messages[messages.length - 1]?.interrupted && !isReadOnly}
+						{#if messages && lastMessage && lastMessage.interrupted && !isReadOnly}
 							<ContinueBtn
-								on:click={() =>
-									dispatch("continue", {
-										id: messages[messages.length - 1].id,
-									})}
+								on:click={() => {
+									if (lastMessage && lastMessage.ancestors) {
+										dispatch("continue", {
+											id: lastMessage?.id,
+										});
+									}
+								}}
 							/>
 						{/if}
 					</div>

src/lib/server/database.ts

@@ -62,6 +62,12 @@ client.on("open", () => {
 			{ partialFilterExpression: { userId: { $exists: true } } }
 		)
 		.catch(console.error);
+	conversations
+		.createIndex(
+			{ "message.id": 1, "message.ancestors": 1 },
+			{ partialFilterExpression: { userId: { $exists: true } } }
+		)
+		.catch(console.error);
 	abortedGenerations.createIndex({ updatedAt: 1 }, { expireAfterSeconds: 30 }).catch(console.error);
 	abortedGenerations.createIndex({ conversationId: 1 }, { unique: true }).catch(console.error);
 	sharedConversations.createIndex({ hash: 1 }, { unique: true }).catch(console.error);

src/lib/server/endpoints/aws/endpointAws.ts

@@ -36,11 +36,11 @@ export async function endpointAws(
 		region,
 	});
 
-	return async ({ conversation }) => {
+	return async ({ messages, preprompt, continueMessage }) => {
 		const prompt = await buildPrompt({
-			messages: conversation.messages,
-			webSearch: conversation.messages[conversation.messages.length - 1].webSearch,
-			preprompt: conversation.preprompt,
+			messages,
+			continueMessage,
+			preprompt,
 			model,
 		});
 

src/lib/server/endpoints/endpoints.ts

@@ -8,13 +8,10 @@ import endpointLlamacpp, { endpointLlamacppParametersSchema } from "./llamacpp/e
 import endpointOllama, { endpointOllamaParametersSchema } from "./ollama/endpointOllama";
 
 // parameters passed when generating text
-interface EndpointParameters {
-	conversation: {
-		messages: Omit<Conversation["messages"][0], "id">[];
-		preprompt?: Conversation["preprompt"];
-		_id?: Conversation["_id"];
-	};
-	continue?: boolean;
+export interface EndpointParameters {
+	messages: Omit<Conversation["messages"][0], "id">[];
+	preprompt?: Conversation["preprompt"];
+	continueMessage?: boolean; // used to signal that the last message will be extended
 }
 
 interface CommonEndpoint {

src/lib/server/endpoints/llamacpp/endpointLlamacpp.ts

@@ -19,11 +19,11 @@ export function endpointLlamacpp(
 	input: z.input<typeof endpointLlamacppParametersSchema>
 ): Endpoint {
 	const { url, model } = endpointLlamacppParametersSchema.parse(input);
-	return async ({ conversation }) => {
+	return async ({ messages, preprompt, continueMessage }) => {
 		const prompt = await buildPrompt({
-			messages: conversation.messages,
-			webSearch: conversation.messages[conversation.messages.length - 1].webSearch,
-			preprompt: conversation.preprompt,
+			messages,
+			continueMessage,
+			preprompt,
 			model,
 		});
 

src/lib/server/endpoints/ollama/endpointOllama.ts

@@ -14,11 +14,11 @@ export const endpointOllamaParametersSchema = z.object({
 export function endpointOllama(input: z.input<typeof endpointOllamaParametersSchema>): Endpoint {
 	const { url, model, ollamaName } = endpointOllamaParametersSchema.parse(input);
 
-	return async ({ conversation }) => {
+	return async ({ messages, preprompt, continueMessage }) => {
 		const prompt = await buildPrompt({
-			messages: conversation.messages,
-			webSearch: conversation.messages[conversation.messages.length - 1].webSearch,
-			preprompt: conversation.preprompt,
+			messages,
+			continueMessage,
+			preprompt,
 			model,
 		});
 

src/lib/server/endpoints/openai/endpointOai.ts

@@ -4,7 +4,6 @@ import { openAIChatToTextGenerationStream } from "./openAIChatToTextGenerationSt
 import { buildPrompt } from "$lib/buildPrompt";
 import { OPENAI_API_KEY } from "$env/static/private";
 import type { Endpoint } from "../endpoints";
-import { format } from "date-fns";
 
 export const endpointOAIParametersSchema = z.object({
 	weight: z.number().int().positive().default(1),
@@ -37,16 +36,18 @@ export async function endpointOai(
 	});
 
 	if (completion === "completions") {
-		return async ({ conversation }) => {
+		return async ({ messages, preprompt, continueMessage }) => {
+			const prompt = await buildPrompt({
+				messages,
+				continueMessage,
+				preprompt,
+				model,
+			});
+
 			return openAICompletionToTextGenerationStream(
 				await openai.completions.create({
 					model: model.id ?? model.name,
-					prompt: await buildPrompt({
-						messages: conversation.messages,
-						webSearch: conversation.messages[conversation.messages.length - 1].webSearch,
-						preprompt: conversation.preprompt,
-						model,
-					}),
+					prompt,
 					stream: true,
 					max_tokens: model.parameters?.max_new_tokens,
 					stop: model.parameters?.stop,
@@ -57,48 +58,20 @@ export async function endpointOai(
 			);
 		};
 	} else if (completion === "chat_completions") {
-		return async ({ conversation }) => {
-			let messages = conversation.messages;
-			const webSearch = conversation.messages[conversation.messages.length - 1].webSearch;
-
-			if (webSearch && webSearch.context) {
-				const lastMsg = messages.slice(-1)[0];
-				const messagesWithoutLastUsrMsg = messages.slice(0, -1);
-				const previousUserMessages = messages.filter((el) => el.from === "user").slice(0, -1);
-
-				const previousQuestions =
-					previousUserMessages.length > 0
-						? `Previous questions: \n${previousUserMessages
-								.map(({ content }) => `- ${content}`)
-								.join("\n")}`
-						: "";
-				const currentDate = format(new Date(), "MMMM d, yyyy");
-				messages = [
-					...messagesWithoutLastUsrMsg,
-					{
-						from: "user",
-						content: `I searched the web using the query: ${webSearch.searchQuery}. Today is ${currentDate} and here are the results:
-						=====================
-						${webSearch.context}
-						=====================
-						${previousQuestions}
-						Answer the question: ${lastMsg.content} 
-						`,
-					},
-				];
-			}
-
-			const messagesOpenAI = messages.map((message) => ({
+		return async ({ messages, preprompt }) => {
+			let messagesOpenAI = messages.map((message) => ({
 				role: message.from,
 				content: message.content,
 			}));
 
+			if (messagesOpenAI?.[0]?.role !== "system") {
+				messagesOpenAI = [{ role: "system", content: preprompt ?? "" }, ...messagesOpenAI];
+			}
+
 			return openAIChatToTextGenerationStream(
 				await openai.chat.completions.create({
 					model: model.id ?? model.name,
-					messages: conversation.preprompt
-						? [{ role: "system", content: conversation.preprompt }, ...messagesOpenAI]
-						: messagesOpenAI,
+					messages: messagesOpenAI,
 					stream: true,
 					max_tokens: model.parameters?.max_new_tokens,
 					stop: model.parameters?.stop,

src/lib/server/endpoints/tgi/endpointTgi.ts

@@ -16,25 +16,14 @@ export const endpointTgiParametersSchema = z.object({
 export function endpointTgi(input: z.input<typeof endpointTgiParametersSchema>): Endpoint {
 	const { url, accessToken, model, authorization } = endpointTgiParametersSchema.parse(input);
 
-	return async ({ conversation, continue: messageContinue }) => {
-		let prompt = await buildPrompt({
-			messages: conversation.messages,
-			webSearch: conversation.messages[conversation.messages.length - 1].webSearch,
-			preprompt: conversation.preprompt,
+	return async ({ messages, preprompt, continueMessage }) => {
+		const prompt = await buildPrompt({
+			messages,
+			preprompt,
 			model,
-			id: conversation._id,
+			continueMessage,
 		});
 
-		if (messageContinue) {
-			// start with the full prompt, and for each stop token, try to remove it from the end of the prompt
-			prompt = model.parameters.stop.reduce((acc: string, curr: string) => {
-				if (acc.endsWith(curr)) {
-					return acc.slice(0, acc.length - curr.length);
-				}
-				return acc;
-			}, prompt.trimEnd());
-		}
-
 		return textGenerationStream(
 			{
 				parameters: { ...model.parameters, return_full_text: false },

src/lib/server/generateFromDefaultEndpoint.ts

@@ -10,7 +10,7 @@ export async function generateFromDefaultEndpoint({
 }): Promise<string> {
 	const endpoint = await smallModel.getEndpoint();
 
-	const tokenStream = await endpoint({ conversation: { messages, preprompt } });
+	const tokenStream = await endpoint({ messages, preprompt });
 
 	for await (const output of tokenStream) {
 		// if not generated_text is here it means the generation is not done

src/lib/server/preprocessMessages.ts

@@ -0,0 +1,57 @@
+import type { Conversation } from "$lib/types/Conversation";
+import type { Message } from "$lib/types/Message";
+import { format } from "date-fns";
+import { downloadFile } from "./files/downloadFile";
+
+export async function preprocessMessages(
+	messages: Message[],
+	multimodal: boolean,
+	id: Conversation["_id"]
+): Promise<Message[]> {
+	return await Promise.all(
+		messages.map(async (message, idx) => {
+			// start by adding websearch to the last message
+			if (idx === messages.length - 1 && message.webSearch && message.webSearch.context) {
+				const lastUsrMsgIndex = messages.map((el) => el.from).lastIndexOf("user");
+				const previousUserMessages = messages.filter((el) => el.from === "user").slice(0, -1);
+				const previousQuestions =
+					previousUserMessages.length > 0
+						? `Previous questions: \n${previousUserMessages
+								.map(({ content }) => `- ${content}`)
+								.join("\n")}`
+						: "";
+				const currentDate = format(new Date(), "MMMM d, yyyy");
+
+				message.content = `I searched the web using the query: ${message.webSearch.searchQuery}. Today is ${currentDate} and here are the results:
+=====================
+${message.webSearch.context}
+=====================
+${previousQuestions}
+Answer the question: ${messages[lastUsrMsgIndex].content}`;
+			}
+			// handle files if model is multimodal
+			if (multimodal) {
+				if (message.files && message.files.length > 0) {
+					const markdowns = await Promise.all(
+						message.files.map(async (hash) => {
+							try {
+								const { content: image, mime } = await downloadFile(hash, id);
+								const b64 = image.toString("base64");
+								return `![](data:${mime};base64,${b64})})`;
+							} catch (e) {
+								console.error(e);
+							}
+						})
+					);
+					message.content += markdowns.join("\n ");
+				} else {
+					// if no image, append an empty white image
+					message.content +=
+						"\n![](data:image/png;base64,/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/2wBDAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/wAARCAAQABADAREAAhEBAxEB/8QAHwAAAQUBAQEBAQEAAAAAAAAAAAECAwQFBgcICQoL/8QAtRAAAgEDAwIEAwUFBAQAAAF9AQIDAAQRBRIhMUEGE1FhByJxFDKBkaEII0KxwRVS0fAkM2JyggkKFhcYGRolJicoKSo0NTY3ODk6Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqDhIWGh4iJipKTlJWWl5iZmqKjpKWmp6ipqrKztLW2t7i5usLDxMXGx8jJytLT1NXW19jZ2uHi4+Tl5ufo6erx8vP09fb3+Pn6/8QAHwEAAwEBAQEBAQEBAQAAAAAAAAECAwQFBgcICQoL/8QAtREAAgECBAQDBAcFBAQAAQJ3AAECAxEEBSExBhJBUQdhcRMiMoEIFEKRobHBCSMzUvAVYnLRChYkNOEl8RcYGRomJygpKjU2Nzg5OkNERUZHSElKU1RVVldYWVpjZGVmZ2hpanN0dXZ3eHl6goOEhYaHiImKkpOUlZaXmJmaoqOkpaanqKmqsrO0tba3uLm6wsPExcbHyMnK0tPU1dbX2Nna4uPk5ebn6Onq8vP09fb3+Pn6/9oADAMBAAIRAxEAPwD+/igAoAKACgD/2Q==)";
+				}
+			}
+
+			return message;
+		})
+	);
+}

src/lib/server/websearch/runWebSearch.ts

@@ -17,13 +17,10 @@ const DOMAIN_BLOCKLIST = ["youtube.com", "twitter.com"];
 
 export async function runWebSearch(
 	conv: Conversation,
-	prompt: string,
+	messages: Message[],
 	updatePad: (upd: MessageUpdate) => void
 ) {
-	const messages = (() => {
-		return [...conv.messages, { content: prompt, from: "user", id: crypto.randomUUID() }];
-	})() satisfies Message[];
-
+	const prompt = messages[messages.length - 1].content;
 	const webSearch: WebSearch = {
 		prompt,
 		searchQuery: "",

src/lib/stores/convTree.ts

@@ -0,0 +1,25 @@
+import type { Message } from "$lib/types/Message";
+import { getContext, setContext } from "svelte";
+import { writable, type Writable } from "svelte/store";
+
+// used to store the id of the message that is the currently displayed leaf of the conversation tree
+// (that is the last message in the current branch of the conversation tree)
+
+interface ConvTreeStore {
+	leaf: Message["id"] | null;
+	editing: Message["id"] | null;
+}
+
+export function useConvTreeStore() {
+	return getContext<Writable<ConvTreeStore>>("convTreeStore");
+}
+
+export function createConvTreeStore() {
+	const convTreeStore = writable<ConvTreeStore>({
+		leaf: null,
+		editing: null,
+	});
+	setContext("convTreeStore", convTreeStore);
+
+	return convTreeStore;
+}

src/lib/types/Conversation.ts

@@ -14,6 +14,7 @@ export interface Conversation extends Timestamps {
 	embeddingModel: string;
 
 	title: string;
+	rootMessageId?: Message["id"];
 	messages: Message[];
 
 	meta?: {

src/lib/types/Message.ts

@@ -3,7 +3,7 @@ import type { Timestamps } from "./Timestamps";
 import type { WebSearch } from "./WebSearch";
 
 export type Message = Partial<Timestamps> & {
-	from: "user" | "assistant";
+	from: "user" | "assistant" | "system";
 	id: ReturnType<typeof crypto.randomUUID>;
 	content: string;
 	updates?: MessageUpdate[];
@@ -12,4 +12,10 @@ export type Message = Partial<Timestamps> & {
 	score?: -1 | 0 | 1;
 	files?: string[]; // can contain either the hash of the file or the b64 encoded image data on the client side when uploading
 	interrupted?: boolean;
+
+	// needed for conversation trees
+	ancestors?: Message["id"][];
+
+	// goes one level deep
+	children?: Message["id"][];
 };

src/lib/types/SharedConversation.ts

@@ -1,17 +1,17 @@
-import type { Assistant } from "./Assistant";
-import type { Message } from "./Message";
-import type { Timestamps } from "./Timestamps";
+import type { Conversation } from "./Conversation";
 
-export interface SharedConversation extends Timestamps {
+export type SharedConversation = Pick<
+	Conversation,
+	| "model"
+	| "embeddingModel"
+	| "title"
+	| "rootMessageId"
+	| "messages"
+	| "preprompt"
+	| "assistantId"
+	| "createdAt"
+	| "updatedAt"
+> & {
 	_id: string;
-
 	hash: string;
-
-	model: string;
-	embeddingModel: string;
-
-	title: string;
-	messages: Message[];
-	preprompt?: string;
-	assistantId?: Assistant["_id"];
-}
+};

src/lib/utils/tree/addChildren.spec.ts

@@ -0,0 +1,102 @@
+import { collections } from "$lib/server/database";
+import { ObjectId } from "mongodb";
+import { describe, expect, it } from "vitest";
+
+import { insertLegacyConversation, insertSideBranchesConversation } from "./treeHelpers.spec";
+import { addChildren } from "./addChildren";
+import type { Message } from "$lib/types/Message";
+
+const newMessage: Omit<Message, "id"> = {
+	content: "new message",
+	from: "user",
+};
+
+Object.freeze(newMessage);
+
+describe("addChildren", async () => {
+	it("should let you append on legacy conversations", async () => {
+		const convId = await insertLegacyConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		const convLength = conv.messages.length;
+
+		addChildren(conv, newMessage, conv.messages[conv.messages.length - 1].id);
+		expect(conv.messages.length).toEqual(convLength + 1);
+	});
+	it("should not let you create branches on legacy conversations", async () => {
+		const convId = await insertLegacyConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		expect(() => addChildren(conv, newMessage, conv.messages[0].id)).toThrow();
+	});
+	it("should not let you create a message that already exists", async () => {
+		const convId = await insertLegacyConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		const messageThatAlreadyExists: Message = {
+			id: conv.messages[0].id,
+			content: "new message",
+			from: "user",
+		};
+
+		expect(() => addChildren(conv, messageThatAlreadyExists, conv.messages[0].id)).toThrow();
+	});
+	it("should let you create branches on conversations with subtrees", async () => {
+		const convId = await insertSideBranchesConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		const nChildren = conv.messages[0].children?.length;
+		if (!nChildren) throw new Error("No children found");
+		addChildren(conv, newMessage, conv.messages[0].id);
+		expect(conv.messages[0].children?.length).toEqual(nChildren + 1);
+	});
+
+	it("should let you create a new leaf", async () => {
+		const convId = await insertSideBranchesConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		const parentId = conv.messages[conv.messages.length - 1].id;
+		const nChildren = conv.messages[conv.messages.length - 1].children?.length;
+
+		if (nChildren === undefined) throw new Error("No children found");
+		expect(nChildren).toEqual(0);
+
+		addChildren(conv, newMessage, parentId);
+		expect(conv.messages[conv.messages.length - 2].children?.length).toEqual(nChildren + 1);
+	});
+
+	it("should let you append to an empty conversation without specifying a parentId", async () => {
+		const conv = {
+			_id: new ObjectId(),
+			rootMessageId: undefined,
+			messages: [] as Message[],
+		};
+
+		addChildren(conv, newMessage);
+		expect(conv.messages.length).toEqual(1);
+		expect(conv.rootMessageId).toEqual(conv.messages[0].id);
+	});
+
+	it("should throw if you don't specify a parentId in a conversation with messages", async () => {
+		const convId = await insertLegacyConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		expect(() => addChildren(conv, newMessage)).toThrow();
+	});
+
+	it("should return the id of the new message", async () => {
+		const convId = await insertLegacyConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		expect(addChildren(conv, newMessage, conv.messages[conv.messages.length - 1].id)).toEqual(
+			conv.messages[conv.messages.length - 1].id
+		);
+	});
+});

src/lib/utils/tree/addChildren.ts

@@ -0,0 +1,52 @@
+import type { Conversation } from "$lib/types/Conversation";
+import type { Message } from "$lib/types/Message";
+
+export function addChildren(
+	conv: Pick<Conversation, "messages" | "rootMessageId">,
+	message: Omit<Message, "id">,
+	parentId?: Message["id"]
+): Message["id"] {
+	// if this is the first message we just push it
+	if (conv.messages.length === 0) {
+		const messageId = crypto.randomUUID();
+		conv.rootMessageId = messageId;
+		conv.messages.push({
+			...message,
+			ancestors: [],
+			id: messageId,
+		});
+		return messageId;
+	}
+
+	if (!parentId) {
+		throw new Error("You need to specify a parentId if this is not the first message");
+	}
+
+	const messageId = crypto.randomUUID();
+	if (!conv.rootMessageId) {
+		// if there is no parentId we just push the message
+		if (!!parentId && parentId !== conv.messages[conv.messages.length - 1].id) {
+			throw new Error("This is a legacy conversation, you can only append to the last message");
+		}
+		conv.messages.push({ ...message, id: messageId });
+		return messageId;
+	}
+
+	const ancestors = [...(conv.messages.find((m) => m.id === parentId)?.ancestors ?? []), parentId];
+	conv.messages.push({
+		...message,
+		ancestors,
+		id: messageId,
+		children: [],
+	});
+
+	const parent = conv.messages.find((m) => m.id === parentId);
+
+	if (parent) {
+		if (parent.children) {
+			parent.children.push(messageId);
+		} else parent.children = [messageId];
+	}
+
+	return messageId;
+}

src/lib/utils/tree/addSibling.spec.ts

@@ -0,0 +1,80 @@
+import { collections } from "$lib/server/database";
+import { ObjectId } from "mongodb";
+import { describe, expect, it } from "vitest";
+
+import { insertLegacyConversation, insertSideBranchesConversation } from "./treeHelpers.spec";
+import type { Message } from "$lib/types/Message";
+import { addSibling } from "./addSibling";
+
+const newMessage: Omit<Message, "id"> = {
+	content: "new message",
+	from: "user",
+};
+
+Object.freeze(newMessage);
+
+describe("addSibling", async () => {
+	it("should fail on empty conversations", () => {
+		const conv = {
+			_id: new ObjectId(),
+			rootMessageId: undefined,
+			messages: [],
+		};
+
+		expect(() => addSibling(conv, newMessage, "not-a-real-id-test")).toThrow(
+			"Cannot add a sibling to an empty conversation"
+		);
+	});
+
+	it("should fail on legacy conversations", async () => {
+		const convId = await insertLegacyConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		expect(() => addSibling(conv, newMessage, conv.messages[0].id)).toThrow(
+			"Cannot add a sibling to a legacy conversation"
+		);
+	});
+
+	it("should fail if the sibling message doesn't exist", async () => {
+		const convId = await insertSideBranchesConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		expect(() => addSibling(conv, newMessage, "not-a-real-id-test")).toThrow(
+			"The sibling message doesn't exist"
+		);
+	});
+
+	// TODO: This behaviour should be fixed, we do not need to fail on the root message.
+	it("should fail if the sibling message is the root message", async () => {
+		const convId = await insertSideBranchesConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+		if (!conv.rootMessageId) throw new Error("Root message not found");
+
+		expect(() => addSibling(conv, newMessage, conv.rootMessageId as Message["id"])).toThrow(
+			"The sibling message is the root message, therefore we can't add a sibling"
+		);
+	});
+
+	it("should add a sibling to a message", async () => {
+		const convId = await insertSideBranchesConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		// add sibling and check children count for parnets
+
+		const nChildren = conv.messages[1].children?.length;
+		const siblingId = addSibling(conv, newMessage, conv.messages[2].id);
+		const nChildrenNew = conv.messages[1].children?.length;
+
+		if (!nChildren) throw new Error("No children found");
+
+		expect(nChildrenNew).toBe(nChildren + 1);
+
+		// make sure siblings have the same ancestors
+		const sibling = conv.messages.find((m) => m.id === siblingId);
+		expect(sibling?.ancestors).toEqual(conv.messages[2].ancestors);
+	});
+});

src/lib/utils/tree/addSibling.ts

@@ -0,0 +1,45 @@
+import type { Conversation } from "$lib/types/Conversation";
+import type { Message } from "$lib/types/Message";
+
+export function addSibling(
+	conv: Pick<Conversation, "messages" | "rootMessageId">,
+	message: Omit<Message, "id">,
+	siblingId: Message["id"]
+): Message["id"] {
+	if (conv.messages.length === 0) {
+		throw new Error("Cannot add a sibling to an empty conversation");
+	}
+	if (!conv.rootMessageId) {
+		throw new Error("Cannot add a sibling to a legacy conversation");
+	}
+
+	const sibling = conv.messages.find((m) => m.id === siblingId);
+
+	if (!sibling) {
+		throw new Error("The sibling message doesn't exist");
+	}
+
+	if (!sibling.ancestors || sibling.ancestors?.length === 0) {
+		throw new Error("The sibling message is the root message, therefore we can't add a sibling");
+	}
+
+	const messageId = crypto.randomUUID();
+
+	conv.messages.push({
+		...message,
+		id: messageId,
+		ancestors: sibling.ancestors,
+		children: [],
+	});
+
+	const nearestAncestorId = sibling.ancestors[sibling.ancestors.length - 1];
+	const nearestAncestor = conv.messages.find((m) => m.id === nearestAncestorId);
+
+	if (nearestAncestor) {
+		if (nearestAncestor.children) {
+			nearestAncestor.children.push(messageId);
+		} else nearestAncestor.children = [messageId];
+	}
+
+	return messageId;
+}

src/lib/utils/tree/buildSubtree.spec.ts

@@ -0,0 +1,110 @@
+import { collections } from "$lib/server/database";
+import { ObjectId } from "mongodb";
+import { describe, expect, it } from "vitest";
+
+import {
+	insertLegacyConversation,
+	insertLinearBranchConversation,
+	insertSideBranchesConversation,
+} from "./treeHelpers.spec";
+import { buildSubtree } from "./buildSubtree";
+
+describe("buildSubtree", () => {
+	it("a subtree in a legacy conversation should be just a slice", async () => {
+		const convId = await insertLegacyConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		// check middle
+		const id = conv.messages[2].id;
+		const subtree = buildSubtree(conv, id);
+		expect(subtree).toEqual(conv.messages.slice(0, 3));
+
+		// check zero
+		const id2 = conv.messages[0].id;
+		const subtree2 = buildSubtree(conv, id2);
+		expect(subtree2).toEqual(conv.messages.slice(0, 1));
+
+		//check full length
+		const id3 = conv.messages[conv.messages.length - 1].id;
+		const subtree3 = buildSubtree(conv, id3);
+		expect(subtree3).toEqual(conv.messages);
+	});
+
+	it("a subtree in a linear branch conversation should be the ancestors and the message", async () => {
+		const convId = await insertLinearBranchConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		// check middle
+		const id = conv.messages[1].id;
+		const subtree = buildSubtree(conv, id);
+		expect(subtree).toEqual([conv.messages[0], conv.messages[1]]);
+
+		// check zero
+		const id2 = conv.messages[0].id;
+		const subtree2 = buildSubtree(conv, id2);
+		expect(subtree2).toEqual([conv.messages[0]]);
+
+		//check full length
+		const id3 = conv.messages[conv.messages.length - 1].id;
+		const subtree3 = buildSubtree(conv, id3);
+		expect(subtree3).toEqual(conv.messages);
+	});
+
+	it("should throw an error if the message is not found", async () => {
+		const convId = await insertLinearBranchConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		const id = "not-a-real-id-test";
+
+		expect(() => buildSubtree(conv, id)).toThrow("Message not found");
+	});
+
+	it("should throw an error if the ancestor is not found", async () => {
+		const convId = await insertLinearBranchConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		const id = "1-1-1-1-2";
+
+		conv.messages[1].ancestors = ["not-a-real-id-test"];
+
+		expect(() => buildSubtree(conv, id)).toThrow("Ancestor not found");
+	});
+
+	it("should work on empty conversations", () => {
+		const conv = {
+			_id: new ObjectId(),
+			rootMessageId: undefined,
+			messages: [],
+		};
+
+		const subtree = buildSubtree(conv, "not-a-real-id-test");
+		expect(subtree).toEqual([]);
+	});
+
+	it("should work for conversation with subtrees", async () => {
+		const convId = await insertSideBranchesConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		const subtree = buildSubtree(conv, "1-1-1-1-2");
+		expect(subtree).toEqual([conv.messages[0], conv.messages[1]]);
+
+		const subtree2 = buildSubtree(conv, "1-1-1-1-4");
+		expect(subtree2).toEqual([
+			conv.messages[0],
+			conv.messages[1],
+			conv.messages[2],
+			conv.messages[3],
+		]);
+
+		const subtree3 = buildSubtree(conv, "1-1-1-1-6");
+		expect(subtree3).toEqual([conv.messages[0], conv.messages[4], conv.messages[5]]);
+
+		const subtree4 = buildSubtree(conv, "1-1-1-1-7");
+		expect(subtree4).toEqual([conv.messages[0], conv.messages[4], conv.messages[6]]);
+	});
+});

src/lib/utils/tree/buildSubtree.ts

@@ -0,0 +1,28 @@
+import type { Conversation } from "$lib/types/Conversation";
+import type { Message } from "$lib/types/Message";
+
+export function buildSubtree(
+	conv: Pick<Conversation, "messages" | "rootMessageId">,
+	id: Message["id"]
+): Message[] {
+	if (!conv.rootMessageId) {
+		if (conv.messages.length === 0) return [];
+		// legacy conversation slice up to id
+		const index = conv.messages.findIndex((m) => m.id === id);
+		if (index === -1) throw new Error("Message not found");
+		return conv.messages.slice(0, index + 1);
+	} else {
+		// find the message with the right id then create the ancestor tree
+		const message = conv.messages.find((m) => m.id === id);
+		if (!message) throw new Error("Message not found");
+
+		return [
+			...(message.ancestors?.map((ancestorId) => {
+				const ancestor = conv.messages.find((m) => m.id === ancestorId);
+				if (!ancestor) throw new Error("Ancestor not found");
+				return ancestor;
+			}) ?? []),
+			message,
+		];
+	}
+}

src/lib/utils/tree/convertLegacyConversation.spec.ts

@@ -0,0 +1,31 @@
+import { collections } from "$lib/server/database";
+import { ObjectId } from "mongodb";
+import { describe, expect, it } from "vitest";
+
+import { convertLegacyConversation } from "./convertLegacyConversation";
+import { insertLegacyConversation } from "./treeHelpers.spec";
+
+describe("convertLegacyConversation", () => {
+	it("should convert a legacy conversation", async () => {
+		const convId = await insertLegacyConversation();
+		const conv = await collections.conversations.findOne({ _id: new ObjectId(convId) });
+		if (!conv) throw new Error("Conversation not found");
+
+		const newConv = convertLegacyConversation(conv);
+
+		expect(newConv.rootMessageId).toBe(newConv.messages[0].id);
+		expect(newConv.messages[0].ancestors).toEqual([]);
+		expect(newConv.messages[1].ancestors).toEqual([newConv.messages[0].id]);
+		expect(newConv.messages[0].children).toEqual([newConv.messages[1].id]);
+	});
+	it("should work on empty conversations", async () => {
+		const conv = {
+			_id: new ObjectId(),
+			rootMessageId: undefined,
+			messages: [],
+		};
+		const newConv = convertLegacyConversation(conv);
+		expect(newConv.rootMessageId).toBe(undefined);
+		expect(newConv.messages).toEqual([]);
+	});
+});

src/lib/utils/tree/convertLegacyConversation.ts

@@ -0,0 +1,35 @@
+import type { Conversation } from "$lib/types/Conversation";
+import type { Message } from "$lib/types/Message";
+
+export function convertLegacyConversation(
+	conv: Pick<Conversation, "messages" | "rootMessageId" | "preprompt">
+): Pick<Conversation, "messages" | "rootMessageId" | "preprompt"> {
+	if (conv.rootMessageId) return conv; // not a legacy conversation
+	if (conv.messages.length === 0) return conv; // empty conversation
+	const messages = [
+		{
+			from: "system",
+			content: conv.preprompt ?? "",
+			createdAt: new Date(),
+			updatedAt: new Date(),
+			id: crypto.randomUUID(),
+		} satisfies Message,
+		...conv.messages,
+	];
+
+	const rootMessageId = messages[0].id;
+
+	const newMessages = messages.map((message, index) => {
+		return {
+			...message,
+			ancestors: messages.slice(0, index).map((m) => m.id),
+			children: index < messages.length - 1 ? [messages[index + 1].id] : [],
+		};
+	});
+
+	return {
+		...conv,
+		rootMessageId,
+		messages: newMessages,
+	};
+}

src/lib/utils/tree/isMessageId.spec.ts

@@ -0,0 +1,14 @@
+import { describe, expect, it } from "vitest";
+import { isMessageId } from "./isMessageId";
+
+describe("isMessageId", () => {
+	it("should return true for a valid message id", () => {
+		expect(isMessageId(crypto.randomUUID())).toBe(true);
+	});
+	it("should return false for an invalid message id", () => {
+		expect(isMessageId("1-2-3-4")).toBe(false);
+	});
+	it("should return false for an empty string", () => {
+		expect(isMessageId("")).toBe(false);
+	});
+});

src/lib/utils/tree/isMessageId.ts

@@ -0,0 +1,5 @@
+import type { Message } from "$lib/types/Message";
+
+export function isMessageId(id: string): id is Message["id"] {
+	return id.split("-").length === 5;
+}

src/lib/utils/tree/treeHelpers.spec.ts

@@ -0,0 +1,164 @@
+import { collections } from "$lib/server/database";
+import { ObjectId } from "mongodb";
+import { describe, expect, it } from "vitest";
+
+// function used to insert conversations used for testing
+
+export const insertLegacyConversation = async () => {
+	const res = await collections.conversations.insertOne({
+		_id: new ObjectId(),
+		createdAt: new Date(),
+		updatedAt: new Date(),
+		title: "legacy conversation",
+		model: "",
+		embeddingModel: "",
+		messages: [
+			{
+				id: "1-1-1-1-1",
+				from: "user",
+				content: "Hello, world! I am a user",
+			},
+			{
+				id: "1-1-1-1-2",
+				from: "assistant",
+				content: "Hello, world! I am an assistant.",
+			},
+			{
+				id: "1-1-1-1-3",
+				from: "user",
+				content: "Hello, world! I am a user.",
+			},
+			{
+				id: "1-1-1-1-4",
+				from: "assistant",
+				content: "Hello, world! I am an assistant.",
+			},
+		],
+	});
+	return res.insertedId;
+};
+
+export const insertLinearBranchConversation = async () => {
+	const res = await collections.conversations.insertOne({
+		_id: new ObjectId(),
+		createdAt: new Date(),
+		updatedAt: new Date(),
+		title: "linear branch conversation",
+		model: "",
+		embeddingModel: "",
+
+		rootMessageId: "1-1-1-1-1",
+		messages: [
+			{
+				id: "1-1-1-1-1",
+				from: "user",
+				content: "Hello, world! I am a user",
+				ancestors: [],
+				children: ["1-1-1-1-2"],
+			},
+			{
+				id: "1-1-1-1-2",
+				from: "assistant",
+				content: "Hello, world! I am an assistant.",
+				ancestors: ["1-1-1-1-1"],
+				children: ["1-1-1-1-3"],
+			},
+			{
+				id: "1-1-1-1-3",
+				from: "user",
+				content: "Hello, world! I am a user.",
+				ancestors: ["1-1-1-1-1", "1-1-1-1-2"],
+				children: ["1-1-1-1-4"],
+			},
+			{
+				id: "1-1-1-1-4",
+				from: "assistant",
+				content: "Hello, world! I am an assistant.",
+				ancestors: ["1-1-1-1-1", "1-1-1-1-2", "1-1-1-1-3"],
+				children: [],
+			},
+		],
+	});
+	return res.insertedId;
+};
+
+export const insertSideBranchesConversation = async () => {
+	const res = await collections.conversations.insertOne({
+		_id: new ObjectId(),
+		createdAt: new Date(),
+		updatedAt: new Date(),
+		title: "side branches conversation",
+		model: "",
+		embeddingModel: "",
+		rootMessageId: "1-1-1-1-1",
+		messages: [
+			{
+				id: "1-1-1-1-1",
+				from: "user",
+				content: "Hello, world, root message!",
+				ancestors: [],
+				children: ["1-1-1-1-2", "1-1-1-1-5"],
+			},
+			{
+				id: "1-1-1-1-2",
+				from: "assistant",
+				content: "Hello, response to root message!",
+				ancestors: ["1-1-1-1-1"],
+				children: ["1-1-1-1-3"],
+			},
+			{
+				id: "1-1-1-1-3",
+				from: "user",
+				content: "Hello, follow up question!",
+				ancestors: ["1-1-1-1-1", "1-1-1-1-2"],
+				children: ["1-1-1-1-4"],
+			},
+			{
+				id: "1-1-1-1-4",
+				from: "assistant",
+				content: "Hello, response from follow up question!",
+				ancestors: ["1-1-1-1-1", "1-1-1-1-2", "1-1-1-1-3"],
+				children: [],
+			},
+			{
+				id: "1-1-1-1-5",
+				from: "assistant",
+				content: "Hello, alternative assistant answer!",
+				ancestors: ["1-1-1-1-1"],
+				children: ["1-1-1-1-6", "1-1-1-1-7"],
+			},
+			{
+				id: "1-1-1-1-6",
+				from: "user",
+				content: "Hello, follow up question to alternative answer!",
+				ancestors: ["1-1-1-1-1", "1-1-1-1-5"],
+				children: [],
+			},
+			{
+				id: "1-1-1-1-7",
+				from: "user",
+				content: "Hello, alternative follow up question to alternative answer!",
+				ancestors: ["1-1-1-1-1", "1-1-1-1-5"],
+				children: [],
+			},
+		],
+	});
+	return res.insertedId;
+};
+
+describe("inserting conversations", () => {
+	it("should insert a legacy conversation", async () => {
+		const id = await insertLegacyConversation();
+		expect(id).toBeDefined();
+	});
+
+	it("should insert a linear branch conversation", async () => {
+		const id = await insertLinearBranchConversation();
+		expect(id).toBeDefined();
+	});
+
+	it("should insert a side branches conversation", async () => {
+		const id = await insertSideBranchesConversation();
+		expect(id).toBeDefined();
+	});
+});

src/routes/conversation/+server.ts

@@ -12,7 +12,6 @@ export const POST: RequestHandler = async ({ locals, request }) => {
 	const body = await request.text();
 
 	let title = "";
-	let messages: Message[] = [];
 
 	const values = z
 		.object({
@@ -23,6 +22,19 @@ export const POST: RequestHandler = async ({ locals, request }) => {
 		})
 		.parse(JSON.parse(body));
 
+	let messages: Message[] = [
+		{
+			id: crypto.randomUUID(),
+			from: "system",
+			content: values.preprompt ?? "",
+			createdAt: new Date(),
+			updatedAt: new Date(),
+			children: [],
+			ancestors: [],
+		},
+	];
+
+	let rootMessageId: Message["id"] = messages[0].id;
 	let embeddingModel: string;
 
 	if (values.fromShare) {
@@ -36,6 +48,7 @@ export const POST: RequestHandler = async ({ locals, request }) => {
 
 		title = conversation.title;
 		messages = conversation.messages;
+		rootMessageId = conversation.rootMessageId ?? rootMessageId;
 		values.model = conversation.model;
 		values.preprompt = conversation.preprompt;
 		values.assistantId = conversation.assistantId?.toString();
@@ -69,6 +82,7 @@ export const POST: RequestHandler = async ({ locals, request }) => {
 	const res = await collections.conversations.insertOne({
 		_id: new ObjectId(),
 		title: title || "New Chat",
+		rootMessageId,
 		messages,
 		model: values.model,
 		preprompt: preprompt === model?.preprompt ? model?.preprompt : preprompt,

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

@@ -3,6 +3,7 @@ import { ObjectId } from "mongodb";
 import { error } from "@sveltejs/kit";
 import { authCondition } from "$lib/server/auth";
 import { UrlDependency } from "$lib/types/UrlDependency";
+import { convertLegacyConversation } from "$lib/utils/tree/convertLegacyConversation.js";
 
 export const load = async ({ params, depends, locals }) => {
 	let conversation;
@@ -45,16 +46,19 @@ export const load = async ({ params, depends, locals }) => {
 		}
 	}
 
+	const convertedConv = { ...conversation, ...convertLegacyConversation(conversation) };
+
 	return {
-		messages: conversation.messages,
-		title: conversation.title,
-		model: conversation.model,
-		preprompt: conversation.preprompt,
-		assistant: conversation.assistantId
+		messages: convertedConv.messages,
+		title: convertedConv.title,
+		model: convertedConv.model,
+		preprompt: convertedConv.preprompt,
+		rootMessageId: convertedConv.rootMessageId,
+		assistant: convertedConv.assistantId
 			? JSON.parse(
 					JSON.stringify(
 						await collections.assistants.findOne({
-							_id: new ObjectId(conversation.assistantId),
+							_id: new ObjectId(convertedConv.assistantId),
 						})
 					)
 			  )

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

@@ -9,20 +9,21 @@
 	import { shareConversation } from "$lib/shareConversation";
 	import { UrlDependency } from "$lib/types/UrlDependency";
 	import { ERROR_MESSAGES, error } from "$lib/stores/errors";
-	import { randomUUID } from "$lib/utils/randomUuid";
 	import { findCurrentModel } from "$lib/utils/models";
 	import { webSearchParameters } from "$lib/stores/webSearchParameters";
 	import type { Message } from "$lib/types/Message";
-	import type { MessageUpdate, WebSearchUpdate } from "$lib/types/MessageUpdate";
+	import type { MessageUpdate } from "$lib/types/MessageUpdate";
 	import titleUpdate from "$lib/stores/titleUpdate";
 	import file2base64 from "$lib/utils/file2base64";
+	import { addChildren } from "$lib/utils/tree/addChildren";
+	import { addSibling } from "$lib/utils/tree/addSibling";
+	import { createConvTreeStore } from "$lib/stores/convTree";
+
 	export let data;
 
 	let messages = data.messages;
 	let lastLoadedMessages = data.messages;
 
-	let webSearchMessages: WebSearchUpdate[] = [];
-
 	// Since we modify the messages array locally, we don't want to reset it if an old version is passed
 	$: if (data.messages !== lastLoadedMessages) {
 		messages = data.messages;
@@ -66,12 +67,12 @@
 	// this function is used to send new message to the backends
 	async function writeMessage({
 		prompt,
-		messageId = randomUUID(),
+		messageId = $convTreeStore.leaf ?? undefined,
 		isRetry = false,
 		isContinue = false,
 	}: {
 		prompt?: string;
-		messageId?: ReturnType<typeof randomUUID>;
+		messageId?: ReturnType<typeof crypto.randomUUID>;
 		isRetry?: boolean;
 		isContinue?: boolean;
 	}): Promise<void> {
@@ -80,25 +81,7 @@
 			loading = true;
 			pending = true;
 
-			// first we check if the messageId already exists, indicating a retry
-
-			let msgIndex = messages.findIndex((msg) => msg.id === messageId);
-
-			if (msgIndex === -1) {
-				msgIndex = messages.length - 1;
-			}
-			if (isRetry && messages[msgIndex].from === "assistant") {
-				throw new Error("Trying to retry a message that is not from user");
-			}
-
-			if (isContinue && messages[msgIndex].from === "user") {
-				throw new Error("Trying to continue a message that is not from assistant");
-			}
-
-			// const isNewMessage = !isRetry && !isContinue;
-
 			const module = await import("browser-image-resizer");
-
 			// currently, only IDEFICS is supported by TGI
 			// the size of images is hardcoded to 224x224 in TGI
 			// this will need to be configurable when support for more models is added
@@ -114,35 +97,101 @@
 				})
 			);
 
-			// slice up to the point of the retry
-			if (isRetry) {
-				messages = [
-					...messages.slice(0, msgIndex),
+			let messageToWriteToId: Message["id"] | undefined = undefined;
+			// used for building the prompt, subtree of the conversation that goes from the latest message to the root
+
+			if (isContinue && messageId) {
+				if ((messages.find((msg) => msg.id === messageId)?.children?.length ?? 0) > 0) {
+					$error = "Can only continue the last message";
+				} else {
+					messageToWriteToId = messageId;
+				}
+			} else if (isRetry && messageId) {
+				// two cases, if we're retrying a user message with a newPrompt set,
+				// it means we're editing a user message
+				// if we're retrying on an assistant message, newPrompt cannot be set
+				// it means we're retrying the last assistant message for a new answer
+
+				const messageToRetry = messages.find((message) => message.id === messageId);
+
+				if (!messageToRetry) {
+					$error = "Message not found";
+				}
+
+				if (messageToRetry?.from === "user" && prompt) {
+					// add a sibling to this message from the user, with the alternative prompt
+					// add a children to that sibling, where we can write to
+					const newUserMessageId = addSibling(
+						{
+							messages,
+							rootMessageId: data.rootMessageId,
+						},
+						{ from: "user", content: prompt },
+						messageId
+					);
+					messageToWriteToId = addChildren(
+						{
+							messages,
+							rootMessageId: data.rootMessageId,
+						},
+						{ from: "assistant", content: "", files: resizedImages },
+						newUserMessageId
+					);
+				} else if (messageToRetry?.from === "assistant") {
+					// we're retrying an assistant message, to generate a new answer
+					// just add a sibling to the assistant answer where we can write to
+					messageToWriteToId = addSibling(
+						{
+							messages,
+							rootMessageId: data.rootMessageId,
+						},
+						{ from: "assistant", content: "" },
+						messageId
+					);
+				}
+			} else {
+				// just a normal linear conversation, so we add the user message
+				// and the blank assistant message back to back
+				const newUserMessageId = addChildren(
 					{
-						from: "user",
-						content: messages[msgIndex].content,
-						id: messageId,
-						files: messages[msgIndex].files,
+						messages,
+						rootMessageId: data.rootMessageId,
 					},
-				];
-			} else if (!isContinue) {
-				// or add a new message if its not a continue request
-				if (!prompt) {
-					throw new Error("Prompt is undefined");
-				}
-				messages = [
-					...messages,
 					{
 						from: "user",
 						content: prompt ?? "",
-						id: messageId,
 						files: resizedImages,
+						createdAt: new Date(),
+						updatedAt: new Date(),
+					},
+					messageId
+				);
+
+				if (!data.rootMessageId) {
+					data.rootMessageId = newUserMessageId;
+				}
+
+				messageToWriteToId = addChildren(
+					{
+						messages,
+						rootMessageId: data.rootMessageId,
+					},
+					{
+						from: "assistant",
+						content: "",
+						createdAt: new Date(),
+						updatedAt: new Date(),
 					},
-				];
+					newUserMessageId
+				);
 			}
 
-			files = [];
+			messages = [...messages];
+			const messageToWriteTo = messages.find((message) => message.id === messageToWriteToId);
 
+			if (!messageToWriteTo) {
+				throw new Error("Message to write to not found");
+			}
 			// disable websearch if assistant is present
 			const hasAssistant = !!$page.data.assistant;
 
@@ -221,25 +270,16 @@
 								invalidate(UrlDependency.Conversation);
 							} else if (update.type === "stream") {
 								pending = false;
-
-								let lastMessage = messages[messages.length - 1];
-
-								if (lastMessage.from !== "assistant") {
-									messages = [
-										...messages,
-										{ from: "assistant", id: randomUUID(), content: update.token },
-									];
-								} else {
-									lastMessage.content += update.token;
-									messages = [...messages];
-								}
+								messageToWriteTo.content += update.token;
+								messages = [...messages];
 							} else if (update.type === "webSearch") {
-								webSearchMessages = [...webSearchMessages, update];
+								messageToWriteTo.updates = [...(messageToWriteTo.updates ?? []), update];
+								messages = [...messages];
 							} else if (update.type === "status") {
 								if (update.status === "title" && update.message) {
-									const conv = data.conversations.find(({ id }) => id === $page.params.id);
-									if (conv) {
-										conv.title = update.message;
+									const convInData = data.conversations.find(({ id }) => id === $page.params.id);
+									if (convInData) {
+										convInData.title = update.message;
 
 										$titleUpdate = {
 											title: update.message,
@@ -265,11 +305,7 @@
 				});
 			}
 
-			webSearchMessages = [];
-
-			const lastMessage = messages[messages.length - 1];
-			lastMessage.updates = messageUpdates;
-
+			messageToWriteTo.updates = messageUpdates;
 			await invalidate(UrlDependency.ConversationList);
 		} catch (err) {
 			if (err instanceof Error && err.message.includes("overloaded")) {
@@ -285,6 +321,7 @@
 		} finally {
 			loading = false;
 			pending = false;
+			await invalidate(UrlDependency.Conversation);
 		}
 	}
 
@@ -336,7 +373,7 @@
 		}
 	}
 
-	async function onRetry(event: CustomEvent<{ id: Message["id"]; content: string }>) {
+	async function onRetry(event: CustomEvent<{ id: Message["id"]; content?: string }>) {
 		if (!data.shared) {
 			await writeMessage({
 				prompt: event.detail.content,
@@ -363,11 +400,26 @@
 	async function onContinue(event: CustomEvent<{ id: Message["id"] }>) {
 		if (!data.shared) {
 			writeMessage({ messageId: event.detail.id, isContinue: true });
+		} else {
+			await convFromShared()
+				.then(async (convId) => {
+					await goto(`${base}/conversation/${convId}`, { invalidateAll: true });
+				})
+				.then(
+					async () =>
+						await writeMessage({
+							messageId: event.detail.id,
+							isContinue: true,
+						})
+				)
+				.finally(() => (loading = false));
 		}
 	}
 
-	$: $page.params.id, (($isAborted = true), (loading = false));
+	$: $page.params.id, (($isAborted = true), (loading = false), ($convTreeStore.editing = null));
 	$: title = data.conversations.find((conv) => conv.id === $page.params.id)?.title ?? data.title;
+
+	const convTreeStore = createConvTreeStore();
 </script>
 
 <svelte:head>
@@ -386,7 +438,6 @@
 	{messages}
 	shared={data.shared}
 	preprompt={data.preprompt}
-	bind:webSearchMessages
 	bind:files
 	on:message={onMessage}
 	on:retry={onRetry}

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

@@ -9,11 +9,16 @@ import { ObjectId } from "mongodb";
 import { z } from "zod";
 import type { MessageUpdate } from "$lib/types/MessageUpdate";
 import { runWebSearch } from "$lib/server/websearch/runWebSearch";
-import type { WebSearch } from "$lib/types/WebSearch";
 import { abortedGenerations } from "$lib/server/abortedGenerations";
 import { summarize } from "$lib/server/summarize";
 import { uploadFile } from "$lib/server/files/uploadFile";
 import sizeof from "image-size";
+import { convertLegacyConversation } from "$lib/utils/tree/convertLegacyConversation";
+import { isMessageId } from "$lib/utils/tree/isMessageId";
+import { buildSubtree } from "$lib/utils/tree/buildSubtree.js";
+import { addChildren } from "$lib/utils/tree/addChildren.js";
+import { addSibling } from "$lib/utils/tree/addSibling.js";
+import { preprocessMessages } from "$lib/server/preprocessMessages.js";
 
 export async function POST({ request, locals, params, getClientAddress }) {
 	const id = z.string().parse(params.id);
@@ -28,6 +33,29 @@ export async function POST({ request, locals, params, getClientAddress }) {
 	}
 
 	// check if the user has access to the conversation
+	const convBeforeCheck = await collections.conversations.findOne({
+		_id: convId,
+		...authCondition(locals),
+	});
+
+	if (convBeforeCheck && !convBeforeCheck.rootMessageId) {
+		const res = await collections.conversations.updateOne(
+			{
+				_id: convId,
+			},
+			{
+				$set: {
+					...convBeforeCheck,
+					...convertLegacyConversation(convBeforeCheck),
+				},
+			}
+		);
+
+		if (!res.acknowledged) {
+			throw error(500, "Failed to convert conversation");
+		}
+	}
+
 	const conv = await collections.conversations.findOne({
 		_id: convId,
 		...authCondition(locals),
@@ -97,8 +125,8 @@ export async function POST({ request, locals, params, getClientAddress }) {
 		files: b64files,
 	} = z
 		.object({
+			id: z.string().uuid().refine(isMessageId).optional(), // parent message id to append to for a normal message, or the message id for a retry/continue
 			inputs: z.optional(z.string().trim().min(1)),
-			id: z.optional(z.string().uuid()),
 			is_retry: z.optional(z.boolean()),
 			is_continue: z.optional(z.boolean()),
 			web_search: z.optional(z.boolean()),
@@ -138,59 +166,93 @@ export async function POST({ request, locals, params, getClientAddress }) {
 		hashes = await Promise.all(files.map(async (file) => await uploadFile(file, conv)));
 	}
 
-	// can only call isContinue on the last message id
-	if (isContinue && conv.messages[conv.messages.length - 1].id !== messageId) {
-		throw error(400, "Can only continue the last message");
-	}
+	// we will append tokens to the content of this message
+	let messageToWriteToId: Message["id"] | undefined = undefined;
+	// used for building the prompt, subtree of the conversation that goes from the latest message to the root
+	let messagesForPrompt: Message[] = [];
 
-	// get the list of messages
-	// while checking for retries
-	let messages = (() => {
-		// for retries we slice and rewrite the last user message
-		if (isRetry && messageId) {
-			// if the message is a retry, replace the message and remove the messages after it
-			let retryMessageIdx = conv.messages.findIndex((message) => message.id === messageId);
+	if (isContinue && messageId) {
+		// if it's the last message and we continue then we build the prompt up to the last message
+		// we will strip the end tokens afterwards when the prompt is built
+		if ((conv.messages.find((msg) => msg.id === messageId)?.children?.length ?? 0) > 0) {
+			throw error(400, "Can only continue the last message");
+		}
+		messageToWriteToId = messageId;
+		messagesForPrompt = buildSubtree(conv, messageId);
+	} else if (isRetry && messageId) {
+		// two cases, if we're retrying a user message with a newPrompt set,
+		// it means we're editing a user message
+		// if we're retrying on an assistant message, newPrompt cannot be set
+		// it means we're retrying the last assistant message for a new answer
+
+		const messageToRetry = conv.messages.find((message) => message.id === messageId);
+
+		if (!messageToRetry) {
+			throw error(404, "Message not found");
+		}
 
-			if (retryMessageIdx === -1) {
-				retryMessageIdx = conv.messages.length;
-			}
+		if (messageToRetry.from === "user" && newPrompt) {
+			// add a sibling to this message from the user, with the alternative prompt
+			// add a children to that sibling, where we can write to
+			const newUserMessageId = addSibling(conv, { from: "user", content: newPrompt }, messageId);
+			messageToWriteToId = addChildren(
+				conv,
+				{ from: "assistant", content: "", files: hashes },
+				newUserMessageId
+			);
+			messagesForPrompt = buildSubtree(conv, newUserMessageId);
+		} else if (messageToRetry.from === "assistant") {
+			// we're retrying an assistant message, to generate a new answer
+			// just add a sibling to the assistant answer where we can write to
+			messageToWriteToId = addSibling(conv, { from: "assistant", content: "" }, messageId);
+			messagesForPrompt = buildSubtree(conv, messageId);
+			messagesForPrompt.pop(); // don't need the latest assistant message in the prompt since we're retrying it
+		}
+	} else {
+		// just a normal linear conversation, so we add the user message
+		// and the blank assistant message back to back
+		const newUserMessageId = addChildren(
+			conv,
+			{
+				from: "user",
+				content: newPrompt ?? "",
+				files: hashes,
+				createdAt: new Date(),
+				updatedAt: new Date(),
+			},
+			messageId
+		);
 
-			return [
-				...conv.messages.slice(0, retryMessageIdx),
-				{
-					content: conv.messages[retryMessageIdx]?.content,
-					from: "user",
-					id: messageId as Message["id"],
-					updatedAt: new Date(),
-					files: conv.messages[retryMessageIdx]?.files,
-				},
-			];
-		} else if (isContinue && messageId) {
-			// for continue we do nothing and expand the last assistant message
-			return conv.messages;
-		} else {
-			// in normal conversation we add an extra user message
-			return [
-				...conv.messages,
-				{
-					content: newPrompt ?? "",
-					from: "user",
-					id: (messageId as Message["id"]) || crypto.randomUUID(),
-					createdAt: new Date(),
-					updatedAt: new Date(),
-					files: hashes,
-				},
-			];
-		} // else append the message at the bottom
-	})() satisfies Message[];
+		messageToWriteToId = addChildren(
+			conv,
+			{
+				from: "assistant",
+				content: "",
+				createdAt: new Date(),
+				updatedAt: new Date(),
+			},
+			newUserMessageId
+		);
+		// build the prompt from the user message
+		messagesForPrompt = buildSubtree(conv, newUserMessageId);
+	}
 
+	const messageToWriteTo = conv.messages.find((message) => message.id === messageToWriteToId);
+	if (!messageToWriteTo) {
+		throw error(500, "Failed to create message");
+	}
+	if (messagesForPrompt.length === 0) {
+		throw error(500, "Failed to create prompt");
+	}
+
+	// update the conversation with the new messages
 	await collections.conversations.updateOne(
 		{
 			_id: convId,
 		},
 		{
 			$set: {
-				messages,
+				messages: conv.messages,
 				title: conv.title,
 				updatedAt: new Date(),
 			},
@@ -202,13 +264,10 @@ export async function POST({ request, locals, params, getClientAddress }) {
 	// we now build the stream
 	const stream = new ReadableStream({
 		async start(controller) {
-			const updates: MessageUpdate[] = isContinue
-				? conv.messages[conv.messages.length - 1].updates ?? []
-				: [];
-
+			messageToWriteTo.updates ??= [];
 			function update(newUpdate: MessageUpdate) {
 				if (newUpdate.type !== "stream") {
-					updates.push(newUpdate);
+					messageToWriteTo?.updates?.push(newUpdate);
 				}
 
 				if (newUpdate.type === "stream" && newUpdate.token === "") {
@@ -225,10 +284,21 @@ export async function POST({ request, locals, params, getClientAddress }) {
 			update({ type: "status", status: "started" });
 
 			const summarizeIfNeeded = (async () => {
-				if (conv.title === "New Chat" && messages.length === 1) {
+				if (conv.title === "New Chat" && conv.messages.length === 3) {
 					try {
-						conv.title = (await summarize(messages[0].content)) ?? conv.title;
+						conv.title = (await summarize(conv.messages[1].content)) ?? conv.title;
 						update({ type: "status", status: "title", message: conv.title });
+						await collections.conversations.updateOne(
+							{
+								_id: convId,
+							},
+							{
+								$set: {
+									title: conv?.title,
+									updatedAt: new Date(),
+								},
+							}
+						);
 					} catch (e) {
 						console.error(e);
 					}
@@ -241,31 +311,33 @@ export async function POST({ request, locals, params, getClientAddress }) {
 				},
 				{
 					$set: {
-						messages,
 						title: conv.title,
 						updatedAt: new Date(),
 					},
 				}
 			);
 
-			let webSearchResults: WebSearch | undefined;
-
+			// perform websearch if needed
 			if (webSearch && !isContinue && !conv.assistantId) {
-				webSearchResults = await runWebSearch(conv, messages[messages.length - 1].content, update);
-				messages[messages.length - 1].webSearch = webSearchResults;
-			} else if (isContinue) {
-				webSearchResults = messages[messages.length - 1].webSearch;
+				messageToWriteTo.webSearch = await runWebSearch(conv, messagesForPrompt, update);
 			}
 
-			conv.messages = messages;
+			// inject websearch result & optionally images into the messages
+			const processedMessages = await preprocessMessages(
+				messagesForPrompt,
+				model.multimodal,
+				convId
+			);
 
-			const previousContent = isContinue
-				? conv.messages.find((message) => message.id === messageId)?.content ?? ""
-				: "";
+			const previousText = messageToWriteTo.content;
 
 			try {
 				const endpoint = await model.getEndpoint();
-				for await (const output of await endpoint({ conversation: conv, continue: isContinue })) {
+				for await (const output of await endpoint({
+					messages: processedMessages,
+					preprompt: conv.preprompt,
+					continueMessage: isContinue,
+				})) {
 					// if not generated_text is here it means the generation is not done
 					if (!output.generated_text) {
 						// else we get the next token
@@ -274,63 +346,33 @@ export async function POST({ request, locals, params, getClientAddress }) {
 								type: "stream",
 								token: output.token.text,
 							});
-
-							// if the last message is not from assistant, it means this is the first token
-							const lastMessage = messages[messages.length - 1];
-
-							if (lastMessage?.from !== "assistant") {
-								// so we create a new message
-								messages = [
-									...messages,
-									// id doesn't match the backend id but it's not important for assistant messages
-									// First token has a space at the beginning, trim it
-									{
-										from: "assistant",
-										content: output.token.text.trimStart(),
-										webSearch: webSearchResults,
-										updates,
-										id: crypto.randomUUID(),
-										createdAt: new Date(),
-										updatedAt: new Date(),
-									},
-								];
-							} else {
-								// abort check
-								const date = abortedGenerations.get(convId.toString());
-								if (date && date > promptedAt) {
-									break;
-								}
-
-								if (!output) {
-									break;
-								}
-
-								// otherwise we just concatenate tokens
-								lastMessage.content += output.token.text;
+							// abort check
+							const date = abortedGenerations.get(convId.toString());
+							if (date && date > promptedAt) {
+								break;
+							}
+							// no output check
+							if (!output) {
+								break;
 							}
+
+							// otherwise we just concatenate tokens
+							messageToWriteTo.content += output.token.text;
 						}
 					} else {
-						let interrupted = !output.token.special;
+						messageToWriteTo.interrupted = !output.token.special;
 						// add output.generated text to the last message
 						// strip end tokens from the output.generated_text
 						const text = (model.parameters.stop ?? []).reduce((acc: string, curr: string) => {
 							if (acc.endsWith(curr)) {
-								interrupted = false;
+								messageToWriteTo.interrupted = false;
 								return acc.slice(0, acc.length - curr.length);
 							}
 							return acc;
 						}, output.generated_text.trimEnd());
 
-						messages = [
-							...messages.slice(0, -1),
-							{
-								...messages[messages.length - 1],
-								content: previousContent + text,
-								interrupted, // if its a special token it finished on its own, else it was interrupted
-								updates,
-								updatedAt: new Date(),
-							},
-						];
+						messageToWriteTo.content = previousText + text;
+						messageToWriteTo.updatedAt = new Date();
 					}
 				}
 			} catch (e) {
@@ -343,7 +385,7 @@ export async function POST({ request, locals, params, getClientAddress }) {
 				},
 				{
 					$set: {
-						messages,
+						messages: conv.messages,
 						title: conv?.title,
 						updatedAt: new Date(),
 					},
@@ -355,7 +397,7 @@ export async function POST({ request, locals, params, getClientAddress }) {
 
 			update({
 				type: "finalAnswer",
-				text: messages[messages.length - 1].content,
+				text: messageToWriteTo.content,
 			});
 
 			await summarizeIfNeeded;
@@ -370,7 +412,7 @@ export async function POST({ request, locals, params, getClientAddress }) {
 					},
 					{
 						$set: {
-							messages,
+							messages: conv.messages,
 							title: conv.title,
 							updatedAt: new Date(),
 						},

src/routes/conversation/[id]/message/[messageId]/prompt/+server.ts

@@ -2,6 +2,8 @@ import { buildPrompt } from "$lib/buildPrompt";
 import { authCondition } from "$lib/server/auth";
 import { collections } from "$lib/server/database";
 import { models } from "$lib/server/models";
+import { buildSubtree } from "$lib/utils/tree/buildSubtree";
+import { isMessageId } from "$lib/utils/tree/isMessageId";
 import { error } from "@sveltejs/kit";
 import { ObjectId } from "mongodb";
 
@@ -24,7 +26,7 @@ export async function GET({ params, locals }) {
 
 	const messageIndex = conv.messages.findIndex((msg) => msg.id === messageId);
 
-	if (messageIndex === -1) {
+	if (!isMessageId(messageId) || messageIndex === -1) {
 		throw error(404, "Message not found");
 	}
 
@@ -34,11 +36,10 @@ export async function GET({ params, locals }) {
 		throw error(404, "Conversation model not found");
 	}
 
-	const messagesUpTo = conv.messages.slice(0, messageIndex + 1);
+	const messagesUpTo = buildSubtree(conv, messageId);
 
 	const prompt = await buildPrompt({
 		preprompt: conv.preprompt,
-		webSearch: messagesUpTo[messagesUpTo.length - 1].webSearch,
 		messages: messagesUpTo,
 		model,
 	});

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

@@ -32,10 +32,11 @@ export async function POST({ params, url, locals }) {
 
 	const shared: SharedConversation = {
 		_id: nanoid(7),
-		createdAt: new Date(),
-		messages: conversation.messages,
 		hash,
+		createdAt: new Date(),
 		updatedAt: new Date(),
+		rootMessageId: conversation.rootMessageId,
+		messages: conversation.messages,
 		title: conversation.title,
 		model: conversation.model,
 		embeddingModel: conversation.embeddingModel,