Refactor of websearch (#281)

* broke up websearch into multiple endpoints

* Refactored loading to use the load function instead of client side fetching

* lint

* Chat Logo Home Screen Bookmark icons for iOS (#279)

* prettier fix for #279

* fix eslint

---------

Co-authored-by: Carolyn Marie <[email protected]>

src/lib/components/OpenWebSearchResults.svelte

@@ -7,41 +7,17 @@
 
 	import EosIconsLoading from "~icons/eos-icons/loading";
 
-	import { base } from "$app/paths";
-	import { onMount } from "svelte";
-
 	export let loading = false;
 	export let classNames = "";
-	export let webSearchId: string | undefined;
 	export let webSearchMessages: WebSearchMessage[] = [];
 
 	let detailsOpen: boolean;
 	let error: boolean;
-	onMount(() => {
-		if (webSearchMessages.length === 0 && webSearchId) {
-			fetch(`${base}/search/${webSearchId}`)
-				.then((res) => res.json())
-				.then((res) => {
-					webSearchMessages = [...res.messages, { type: "result", id: webSearchId }];
-				})
-				.catch((err) => console.log(err));
-		}
-	});
 	$: error = webSearchMessages.some((message) => message.type === "error");
 </script>
 
 <details
 	class="flex w-fit rounded-xl border border-gray-200 bg-white shadow-sm dark:border-gray-800 dark:bg-gray-900 {classNames} max-w-full"
-	on:toggle={() => {
-		if (webSearchMessages.length === 0 && webSearchId) {
-			fetch(`${base}/search/${webSearchId}`)
-				.then((res) => res.json())
-				.then((res) => {
-					webSearchMessages = [...res.messages, { type: "result", id: webSearchId }];
-				})
-				.catch((err) => console.log(err));
-		}
-	}}
 	bind:open={detailsOpen}
 >
 	<summary

src/lib/components/chat/ChatMessage.svelte

@@ -46,7 +46,6 @@
 	export let isAuthor = true;
 	export let readOnly = false;
 	export let isTapped = false;
-	export let isLast = false;
 
 	export let webSearchMessages: WebSearchMessage[] = [];
 
@@ -99,9 +98,8 @@
 	let webSearchIsDone = true;
 
 	$: webSearchIsDone =
-		!!message.webSearchId ||
-		(webSearchMessages.length > 0 &&
-			webSearchMessages[webSearchMessages.length - 1].type === "result");
+		webSearchMessages.length > 0 &&
+		webSearchMessages[webSearchMessages.length - 1].type === "result";
 </script>
 
 {#if message.from === "assistant"}
@@ -118,17 +116,14 @@
 		<div
 			class="relative min-h-[calc(2rem+theme(spacing[3.5])*2)] min-w-[60px] break-words rounded-2xl border border-gray-100 bg-gradient-to-br from-gray-50 px-5 py-3.5 text-gray-600 prose-pre:my-2 dark:border-gray-800 dark:from-gray-800/40 dark:text-gray-300"
 		>
-			{#if message.webSearchId || (webSearchMessages.length > 0 && isLast)}
-				{#key (message.webSearchId, message.score, loading)}
-					<OpenWebSearchResults
-						classNames={tokens.length ? "mb-3.5" : ""}
-						webSearchId={message.webSearchId}
-						{webSearchMessages}
-						loading={!webSearchIsDone}
-					/>
-				{/key}
+			{#if webSearchMessages && webSearchMessages.length > 0}
+				<OpenWebSearchResults
+					classNames={tokens.length ? "mb-3.5" : ""}
+					{webSearchMessages}
+					loading={!webSearchIsDone}
+				/>
 			{/if}
-			{#if !message.content && (webSearchIsDone || webSearchMessages.length === 0)}
+			{#if !message.content && (webSearchIsDone || (webSearchMessages && webSearchMessages.length === 0))}
 				<IconLoading />
 			{/if}
 

src/lib/components/chat/ChatMessages.svelte

@@ -9,7 +9,6 @@
 	import ChatIntroduction from "./ChatIntroduction.svelte";
 	import ChatMessage from "./ChatMessage.svelte";
 	import type { WebSearchMessage } from "$lib/types/WebSearch";
-	import { page } from "$app/stores";
 
 	export let messages: Message[];
 	export let loading: boolean;
@@ -19,7 +18,9 @@
 	export let settings: LayoutData["settings"];
 	export let models: Model[];
 	export let readOnly: boolean;
+	export let searches: Record<string, WebSearchMessage[]>;
 
+	let webSearchArray: Array<WebSearchMessage[] | undefined> = [];
 	let chatContainer: HTMLElement;
 
 	export let webSearchMessages: WebSearchMessage[] = [];
@@ -33,6 +34,17 @@
 	$: if (messages[messages.length - 1]?.from === "user") {
 		scrollToBottom();
 	}
+
+	$: messages,
+		(webSearchArray = messages.map((message, idx) => {
+			if (message.webSearchId) {
+				return searches[message.webSearchId] ?? [];
+			} else if (idx === messages.length - 1) {
+				return webSearchMessages;
+			} else {
+				return [];
+			}
+		}));
 </script>
 
 <div
@@ -42,19 +54,16 @@
 >
 	<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}
-			{#key (message.id, $page.params.id)}
-				<ChatMessage
-					loading={loading && i === messages.length - 1}
-					{message}
-					{isAuthor}
-					{readOnly}
-					model={currentModel}
-					{webSearchMessages}
-					isLast={i === messages.length - 1}
-					on:retry
-					on:vote
-				/>
-			{/key}
+			<ChatMessage
+				loading={loading && i === messages.length - 1}
+				{message}
+				{isAuthor}
+				{readOnly}
+				model={currentModel}
+				webSearchMessages={webSearchArray[i]}
+				on:retry
+				on:vote
+			/>
 		{:else}
 			<ChatIntroduction {settings} {models} {currentModel} on:message />
 		{/each}
@@ -62,7 +71,6 @@
 			<ChatMessage
 				message={{ from: "assistant", content: "", id: randomUUID() }}
 				model={currentModel}
-				isLast={true}
 				{webSearchMessages}
 			/>
 		{/if}

src/lib/components/chat/ChatWindow.svelte

@@ -24,6 +24,7 @@
 	export let models: Model[];
 	export let settings: LayoutData["settings"];
 	export let webSearchMessages: WebSearchMessage[] = [];
+	export let searches: Record<string, WebSearchMessage[]> = {};
 
 	export let loginRequired = false;
 	$: isReadOnly = !models.some((model) => model.id === currentModel.id);
@@ -59,6 +60,7 @@
 		readOnly={isReadOnly}
 		isAuthor={!shared}
 		{webSearchMessages}
+		{searches}
 		on:message
 		on:vote
 		on:retry={(ev) => {

src/lib/server/websearch/generateQuery.ts

@@ -0,0 +1,25 @@
+import type { Message } from "$lib/types/Message";
+import { generateFromDefaultEndpoint } from "../generateFromDefaultEndpoint";
+import type { BackendModel } from "../models";
+
+export async function generateQuery(messages: Message[], model: BackendModel) {
+	const promptSearchQuery =
+		model.userMessageToken +
+		"The following messages were written by a user, trying to answer a question." +
+		model.messageEndToken +
+		messages
+			.filter((message) => message.from === "user")
+			.map((message) => model.userMessageToken + message.content + model.messageEndToken) +
+		model.userMessageToken +
+		"What plain-text english sentence would you input into Google to answer the last question? Answer with a short (10 words max) simple sentence." +
+		model.messageEndToken +
+		model.assistantMessageToken +
+		"Query: ";
+
+	const searchQuery = await generateFromDefaultEndpoint(promptSearchQuery).then((query) => {
+		const arr = query.split(/\r?\n/);
+		return arr[0].length > 0 ? arr[0] : arr[1];
+	});
+
+	return searchQuery;
+}

src/lib/server/websearch/parseWeb.ts

@@ -0,0 +1,56 @@
+import { JSDOM, VirtualConsole } from "jsdom";
+
+function removeTags(node: Node) {
+	if (node.hasChildNodes()) {
+		node.childNodes.forEach((childNode) => {
+			if (node.nodeName === "SCRIPT" || node.nodeName === "STYLE") {
+				node.removeChild(childNode);
+			} else {
+				removeTags(childNode);
+			}
+		});
+	}
+}
+function naiveInnerText(node: Node): string {
+	const Node = node; // We need Node(DOM's Node) for the constants, but Node doesn't exist in the nodejs global space, and any Node instance references the constants through the prototype chain
+	return [...node.childNodes]
+		.map((childNode) => {
+			switch (childNode.nodeType) {
+				case Node.TEXT_NODE:
+					return node.textContent;
+				case Node.ELEMENT_NODE:
+					return naiveInnerText(childNode);
+				default:
+					return "";
+			}
+		})
+		.join("\n");
+}
+
+export async function parseWeb(url: string) {
+	const abortController = new AbortController();
+	setTimeout(() => abortController.abort(), 10000);
+	const htmlString = await fetch(url, { signal: abortController.signal })
+		.then((response) => response.text())
+		.catch((err) => console.log(err));
+
+	const virtualConsole = new VirtualConsole();
+	virtualConsole.on("error", () => {
+		// No-op to skip console errors.
+	});
+
+	// put the html string into a DOM
+	const dom = new JSDOM(htmlString ?? "", {
+		virtualConsole,
+	});
+
+	const body = dom.window.document.querySelector("body");
+	if (!body) throw new Error("body of the webpage is null");
+
+	removeTags(body);
+
+	// recursively extract text content from the body and then remove newlines and multiple spaces
+	const text = (naiveInnerText(body) ?? "").replace(/ {2}|\r\n|\n|\r/gm, "");
+
+	return text;
+}

src/lib/server/websearch/summarizeWeb.ts

@@ -0,0 +1,23 @@
+import { generateFromDefaultEndpoint } from "../generateFromDefaultEndpoint";
+import type { BackendModel } from "../models";
+
+export async function summarizeWeb(content: string, query: string, model: BackendModel) {
+	const summaryPrompt =
+		model.userMessageToken +
+		content
+			.split(" ")
+			.slice(0, model.parameters?.truncate ?? 0)
+			.join(" ") +
+		model.messageEndToken +
+		model.userMessageToken +
+		`The text above should be summarized to best answer the query: ${query}.` +
+		model.messageEndToken +
+		model.assistantMessageToken +
+		"Summary: ";
+
+	const summary = await generateFromDefaultEndpoint(summaryPrompt).then((txt: string) =>
+		txt.trim()
+	);
+
+	return summary;
+}

src/lib/types/UrlDependency.ts

@@ -1,4 +1,5 @@
 /* eslint-disable no-shadow */
 export enum UrlDependency {
 	ConversationList = "conversation:list",
+	Conversation = "conversation",
 }

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

@@ -2,14 +2,18 @@ import { collections } from "$lib/server/database";
 import { ObjectId } from "mongodb";
 import { error } from "@sveltejs/kit";
 import { authCondition } from "$lib/server/auth";
+import type { WebSearchMessageResult } from "$lib/types/WebSearch";
+import { UrlDependency } from "$lib/types/UrlDependency";
 
-export const load = async ({ params, locals }) => {
+export const load = async ({ params, depends, locals }) => {
 	// todo: add validation on params.id
 	const conversation = await collections.conversations.findOne({
 		_id: new ObjectId(params.id),
 		...authCondition(locals),
 	});
 
+	depends(UrlDependency.Conversation);
+
 	if (!conversation) {
 		const conversationExists =
 			(await collections.conversations.countDocuments({
@@ -26,9 +30,23 @@ export const load = async ({ params, locals }) => {
 		throw error(404, "Conversation not found.");
 	}
 
+	const webSearchesId = conversation.messages
+		.filter((message) => message.webSearchId)
+		.map((message) => new ObjectId(message.webSearchId));
+
+	const results = await collections.webSearches.find({ _id: { $in: webSearchesId } }).toArray();
+
+	const searches = Object.fromEntries(
+		results.map((x) => [
+			x._id.toString(),
+			[...x.messages, { type: "result", id: x._id.toString() } satisfies WebSearchMessageResult],
+		])
+	);
+
 	return {
 		messages: conversation.messages,
 		title: conversation.title,
 		model: conversation.model,
+		searches,
 	};
 };

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

@@ -13,8 +13,9 @@
 	import { randomUUID } from "$lib/utils/randomUuid";
 	import { findCurrentModel } from "$lib/utils/models";
 	import { webSearchParameters } from "$lib/stores/webSearchParameters";
-	import type { WebSearchMessage } from "$lib/types/WebSearch.js";
+	import type { WebSearchMessage } from "$lib/types/WebSearch";
 	import type { Message } from "$lib/types/Message";
+	import { browser } from "$app/environment";
 
 	export let data;
 
@@ -194,6 +195,7 @@
 			await getTextGenerationStream(message, messageId, isRetry, searchResponseId ?? undefined);
 
 			webSearchMessages = [];
+			if (browser) invalidate(UrlDependency.Conversation);
 
 			if (messages.filter((m) => m.from === "user").length === 1) {
 				summarizeTitle($page.params.id)
@@ -266,6 +268,7 @@
 	{pending}
 	{messages}
 	bind:webSearchMessages
+	searches={{ ...data.searches }}
 	on:message={(event) => writeMessage(event.detail)}
 	on:retry={(event) => writeMessage(event.detail.content, event.detail.id)}
 	on:vote={(event) => voteMessage(event.detail.score, event.detail.id)}

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

@@ -1,7 +1,7 @@
 import { buildPrompt } from "$lib/buildPrompt";
 import { authCondition } from "$lib/server/auth";
 import { collections } from "$lib/server/database";
-import { generateFromDefaultEndpoint } from "$lib/server/generateFromDefaultEndpoint.js";
+import { generateFromDefaultEndpoint } from "$lib/server/generateFromDefaultEndpoint";
 import { defaultModel } from "$lib/server/models";
 import { error } from "@sveltejs/kit";
 import { ObjectId } from "mongodb";

src/routes/conversation/[id]/web-search/+server.ts

@@ -1,41 +1,15 @@
 import { authCondition } from "$lib/server/auth";
 import { collections } from "$lib/server/database";
-import { generateFromDefaultEndpoint } from "$lib/server/generateFromDefaultEndpoint.js";
 import { defaultModel } from "$lib/server/models";
-import { searchWeb } from "$lib/server/searchWeb.js";
-import type { Message } from "$lib/types/Message.js";
+import { searchWeb } from "$lib/server/websearch/searchWeb";
+import type { Message } from "$lib/types/Message";
 import { error } from "@sveltejs/kit";
 import { ObjectId } from "mongodb";
 import { z } from "zod";
-import { JSDOM, VirtualConsole } from "jsdom";
-import type { WebSearch } from "$lib/types/WebSearch.js";
-
-function removeTags(node: Node) {
-	if (node.hasChildNodes()) {
-		node.childNodes.forEach((childNode) => {
-			if (node.nodeName === "SCRIPT" || node.nodeName === "STYLE") {
-				node.removeChild(childNode);
-			} else {
-				removeTags(childNode);
-			}
-		});
-	}
-}
-function naiveInnerText(node: Node): string {
-	const Node = node; // We need Node(DOM's Node) for the constants, but Node doesn't exist in the nodejs global space, and any Node instance references the constants through the prototype chain
-	return [...node.childNodes]
-		.map((childNode) => {
-			switch (childNode.nodeType) {
-				case Node.TEXT_NODE:
-					return node.textContent;
-				case Node.ELEMENT_NODE:
-					return naiveInnerText(childNode);
-				default:
-					return "";
-			}
-		})
-		.join("\n");
-}
+import type { WebSearch } from "$lib/types/WebSearch";
+import { generateQuery } from "$lib/server/websearch/generateQuery";
+import { parseWeb } from "$lib/server/websearch/parseWeb";
+import { summarizeWeb } from "$lib/server/websearch/summarizeWeb";
 
 interface GenericObject {
 	[key: string]: GenericObject | unknown;
@@ -82,45 +56,24 @@ export async function GET({ params, locals, url }) {
 				createdAt: new Date(),
 				updatedAt: new Date(),
 			};
-			try {
-				webSearch.messages.push({
-					type: "update",
-					message: "Generating search query",
-				});
-				controller.enqueue(JSON.stringify({ messages: webSearch.messages }));
-
-				const promptSearchQuery =
-					model.userMessageToken +
-					"The following messages were written by a user, trying to answer a question." +
-					model.messageEndToken +
-					messages
-						.filter((message) => message.from === "user")
-						.map((message) => model.userMessageToken + message.content + model.messageEndToken) +
-					model.userMessageToken +
-					"What plain-text english sentence would you input into Google to answer the last question? Answer with a short (10 words max) simple sentence." +
-					model.messageEndToken +
-					model.assistantMessageToken +
-					"Query: ";
-
-				webSearch.searchQuery = await generateFromDefaultEndpoint(promptSearchQuery).then(
-					(query) => {
-						const arr = query.split(/\r?\n/);
-						return arr[0].length > 0 ? arr[0] : arr[1];
-					}
-				);
-				// the model has a tendency to continue answering even when we tell it not to, so the split makes
-				// sure we only get the first line of the response
 
+			function appendUpdate(message: string, args?: string[]) {
 				webSearch.messages.push({
 					type: "update",
-					message: "Searching Google",
-					args: [webSearch.searchQuery],
+					message,
+					args,
 				});
 				controller.enqueue(JSON.stringify({ messages: webSearch.messages }));
+			}
 
+			try {
+				appendUpdate("Generating search query");
+				webSearch.searchQuery = await generateQuery(messages, model);
+
+				appendUpdate("Searching Google", [webSearch.searchQuery]);
 				const results = await searchWeb(webSearch.searchQuery);
-				let text = "";
 
+				let text = "";
 				webSearch.results =
 					(results.organic_results &&
 						results.organic_results.map((el: { link: string }) => el.link)) ??
@@ -129,85 +82,22 @@ export async function GET({ params, locals, url }) {
 				if (results.knowledge_graph) {
 					// if google returns a knowledge graph, we use it
 					webSearch.knowledgeGraph = JSON.stringify(removeLinks(results.knowledge_graph));
-
 					text = webSearch.knowledgeGraph;
-
-					webSearch.messages.push({
-						type: "update",
-						message: "Found a Google knowledge page",
-					});
-					controller.enqueue(JSON.stringify({ messages: webSearch.messages }));
+					appendUpdate("Found a Google knowledge page");
 				} else if (webSearch.results.length > 0) {
 					// otherwise we use the top result from search
 					const topUrl = webSearch.results[0];
+					appendUpdate("Browsing first result", [JSON.stringify(topUrl)]);
 
-					webSearch.messages.push({
-						type: "update",
-						message: "Browsing first result",
-						args: [JSON.stringify(topUrl)],
-					});
-					controller.enqueue(JSON.stringify({ messages: webSearch.messages }));
-
-					// fetch the webpage
-					//10 second timeout:
-					const abortController = new AbortController();
-					setTimeout(() => abortController.abort(), 10000);
-					const htmlString = await fetch(topUrl, { signal: abortController.signal })
-						.then((response) => response.text())
-						.catch((err) => console.log(err));
-
-					const virtualConsole = new VirtualConsole();
-					virtualConsole.on("error", () => {
-						// No-op to skip console errors.
-					});
-
-					// put the html string into a DOM
-					const dom = new JSDOM(htmlString ?? "", {
-						virtualConsole,
-					});
-
-					const body = dom.window.document.querySelector("body");
-					if (!body) throw new Error("body of the webpage is null");
-
-					removeTags(body);
-
-					// recursively extract text content from the body and then remove newlines and multiple spaces
-					text = (naiveInnerText(body) ?? "").replace(/ {2}|\r\n|\n|\r/gm, "");
-
+					text = await parseWeb(topUrl);
 					if (!text) throw new Error("text of the webpage is null");
 				} else {
 					throw new Error("No results found for this search query");
 				}
 
-				webSearch.messages.push({
-					type: "update",
-					message: "Creating summary",
-				});
-				controller.enqueue(JSON.stringify({ messages: webSearch.messages }));
-
-				const summaryPrompt =
-					model.userMessageToken +
-					text
-						.split(" ")
-						.slice(0, model.parameters?.truncate ?? 0)
-						.join(" ") +
-					model.messageEndToken +
-					model.userMessageToken +
-					`The text above should be summarized to best answer the query: ${webSearch.searchQuery}.` +
-					model.messageEndToken +
-					model.assistantMessageToken +
-					"Summary: ";
-
-				webSearch.summary = await generateFromDefaultEndpoint(summaryPrompt).then((txt: string) =>
-					txt.trim()
-				);
-
-				webSearch.messages.push({
-					type: "update",
-					message: "Injecting summary",
-					args: [JSON.stringify(webSearch.summary)],
-				});
-				controller.enqueue(JSON.stringify({ messages: webSearch.messages }));
+				appendUpdate("Creating summary");
+				webSearch.summary = await summarizeWeb(text, webSearch.searchQuery, model);
+				appendUpdate("Injecting summary", [JSON.stringify(webSearch.summary)]);
 			} catch (searchError) {
 				if (searchError instanceof Error) {
 					webSearch.messages.push({
@@ -219,7 +109,6 @@ export async function GET({ params, locals, url }) {
 			}
 
 			const res = await collections.webSearches.insertOne(webSearch);
-
 			webSearch.messages.push({
 				type: "result",
 				id: res.insertedId.toString(),

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

@@ -1,6 +1,8 @@
 import type { PageServerLoad } from "./$types";
 import { collections } from "$lib/server/database";
 import { error } from "@sveltejs/kit";
+import { ObjectId } from "mongodb";
+import type { WebSearchMessageResult } from "$lib/types/WebSearch";
 
 export const load: PageServerLoad = async ({ params }) => {
 	const conversation = await collections.sharedConversations.findOne({
@@ -11,9 +13,23 @@ export const load: PageServerLoad = async ({ params }) => {
 		throw error(404, "Conversation not found");
 	}
 
+	const webSearchesId = conversation.messages
+		.filter((message) => message.webSearchId)
+		.map((message) => new ObjectId(message.webSearchId));
+
+	const results = await collections.webSearches.find({ _id: { $in: webSearchesId } }).toArray();
+
+	const searches = Object.fromEntries(
+		results.map((x) => [
+			x._id.toString(),
+			[...x.messages, { type: "result", id: x._id.toString() } satisfies WebSearchMessageResult],
+		])
+	);
+
 	return {
 		messages: conversation.messages,
 		title: conversation.title,
 		model: conversation.model,
+		searches,
 	};
 };

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

@@ -59,6 +59,7 @@
 	{loading}
 	shared={true}
 	messages={data.messages}
+	searches={data.searches}
 	on:message={(ev) =>
 		createConversation()
 			.then((convId) => {

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

@@ -1,5 +1,5 @@
 import { collections } from "$lib/server/database";
-import { sha256 } from "$lib/utils/sha256.js";
+import { sha256 } from "$lib/utils/sha256";
 import { error } from "@sveltejs/kit";
 import { ObjectId } from "mongodb";