feat: add sources for websearch (#1551)

* feat: playwright, spatial parsing, markdown for web search

Co-authored-by: Aaditya Sahay <[email protected]>

* feat: choose multiple clusters if necessary (#2)

* chore: resolve linting failures

* feat: improve paring performance and error messages

* feat: inline citations

* feat: adjust inline citation prompt, less intrusive tokens

* feat: add sources to message when using websearch

* fix: clean up packages

* fix: packages

* fix: packages lol

* fix: make websearch citation work better wiht tools

* fix: use single brackets for sources, only render source element if a matching source is available

* fix: bad import

---------

Co-authored-by: Liam Dyer <[email protected]>
Co-authored-by: Aaditya Sahay <[email protected]>
Co-authored-by: Aaditya Sahay <[email protected]>

package-lock.json

@@ -13434,4 +13434,4 @@
 			}
 		}
 	}
-}
+}

src/lib/components/chat/ChatMessage.svelte

@@ -37,6 +37,28 @@
 	import DOMPurify from "isomorphic-dompurify";
 	import { enhance } from "$app/forms";
 	import { browser } from "$app/environment";
+	import type { WebSearchSource } from "$lib/types/WebSearch";
+
+	function addInlineCitations(md: string, webSearchSources: WebSearchSource[] = []): string {
+		const linkStyle =
+			"color: rgb(59, 130, 246); text-decoration: none; hover:text-decoration: underline;";
+
+		return md.replace(/\[(\d+)\]/g, (match: string) => {
+			const indices: number[] = (match.match(/\d+/g) || []).map(Number);
+			const links: string = indices
+				.map((index: number) => {
+					const source = webSearchSources[index - 1];
+					if (source) {
+						return `<a href="${source.link}" target="_blank" rel="noreferrer" style="${linkStyle}">${index}</a>`;
+					}
+					return "";
+				})
+				.filter(Boolean)
+				.join(", ");
+
+			return links ? ` <sup>${links}</sup>` : match;
+		});
+	}
 
 	function sanitizeMd(md: string) {
 		let ret = md
@@ -114,7 +136,7 @@
 		})
 	);
 
-	$: tokens = marked.lexer(sanitizeMd(message.content ?? ""));
+	$: tokens = marked.lexer(addInlineCitations(sanitizeMd(message.content), webSearchSources));
 
 	$: emptyLoad =
 		!message.content && (webSearchIsDone || (searchUpdates && searchUpdates.length === 0));

src/lib/server/endpoints/preprocessMessages.ts

@@ -17,7 +17,7 @@ export async function preprocessMessages(
 
 function addWebSearchContext(messages: Message[], webSearch: Message["webSearch"]) {
 	const webSearchContext = webSearch?.contextSources
-		.map(({ context }) => context.trim())
+		.map(({ context }, idx) => `Source [${idx + 1}]\n${context.trim()}`)
 		.join("\n\n----------\n\n");
 
 	// No web search context available, skip
@@ -35,7 +35,7 @@ function addWebSearchContext(messages: Message[], webSearch: Message["webSearch"
 	const finalMessage = {
 		...messages[messages.length - 1],
 		content: `I searched the web using the query: ${webSearch.searchQuery}.
-Today is ${currentDate} and here are the results:
+Today is ${currentDate} and here are the results. When answering the question, if you use a source, cite its index inline like this: [1], [2], etc.
 =====================
 ${webSearchContext}
 =====================

src/lib/server/tools/web/search.ts

@@ -25,12 +25,21 @@ const websearch: ConfigTool = {
 	showOutput: false,
 	async *call({ query }, { conv, assistant, messages }) {
 		const webSearchToolResults = yield* runWebSearch(conv, messages, assistant?.rag, String(query));
-		const chunks = webSearchToolResults?.contextSources
-			.map(({ context }) => context)
-			.join("\n------------\n");
+
+		const webSearchContext = webSearchToolResults?.contextSources
+			.map(({ context }, idx) => `Source [${idx + 1}]\n${context.trim()}`)
+			.join("\n\n----------\n\n");
 
 		return {
-			outputs: [{ websearch: chunks }],
+			outputs: [
+				{
+					websearch: webSearchContext,
+				},
+				{
+					instructions:
+						"When answering the question, if you use sources from the websearch results above, cite each index inline individually wrapped like: [1], [2] etc.",
+				},
+			],
 			display: false,
 		};
 	},