[Websearch] change context schema (#944)

src/lib/server/preprocessMessages.ts

@@ -11,8 +11,14 @@ export async function preprocessMessages(
 ): Promise<Message[]> {
 	return await Promise.all(
 		structuredClone(messages).map(async (message, idx) => {
+			const webSearchContext = webSearch?.contextSources
+				.map(({ context }) => context)
+				.flat()
+				.sort((a, b) => a.idx - b.idx)
+				.map(({ text }) => text)
+				.join(" ");
 			// start by adding websearch to the last message
-			if (idx === messages.length - 1 && webSearch && webSearch.context) {
+			if (idx === messages.length - 1 && webSearch && webSearchContext?.trim()) {
 				const lastQuestion = messages.findLast((el) => el.from === "user")?.content ?? "";
 				const previousQuestions = messages
 					.filter((el) => el.from === "user")
@@ -23,7 +29,7 @@ export async function preprocessMessages(
 				message.content = `I searched the web using the query: ${webSearch.searchQuery}. 
 Today is ${currentDate} and here are the results:
 =====================
-${webSearch.context}
+${webSearchContext}
 =====================
 ${previousQuestions.length > 0 ? `Previous questions: \n- ${previousQuestions.join("\n- ")}` : ""}
 Answer the question: ${lastQuestion}`;

src/lib/server/websearch/runWebSearch.ts

@@ -36,7 +36,6 @@ export async function runWebSearch(
 		prompt,
 		searchQuery: "",
 		results: [],
-		context: "",
 		contextSources: [],
 		createdAt: new Date(),
 		updatedAt: new Date(),
@@ -153,14 +152,15 @@ export async function runWebSearch(
 		const indices = await findSimilarSentences(embeddingModel, prompt, texts, {
 			topK: topKClosestParagraphs,
 		});
-		webSearch.context = indices.map((idx) => texts[idx]).join("");
 
-		const usedSources = new Set<string>();
 		for (const idx of indices) {
 			const { source } = paragraphChunks[idx];
-			if (!usedSources.has(source.link)) {
-				usedSources.add(source.link);
-				webSearch.contextSources.push(source);
+			const contextWithId = { idx, text: texts[idx] };
+			const usedSource = webSearch.contextSources.find((cSource) => cSource.link === source.link);
+			if (usedSource) {
+				usedSource.context.push(contextWithId);
+			} else {
+				webSearch.contextSources.push({ ...source, context: [contextWithId] });
 			}
 		}
 		updatePad({

src/lib/types/WebSearch.ts

@@ -10,8 +10,7 @@ export interface WebSearch extends Timestamps {
 
 	searchQuery: string;
 	results: WebSearchSource[];
-	context: string;
-	contextSources: WebSearchSource[];
+	contextSources: WebSearchUsedSource[];
 }
 
 export interface WebSearchSource {
@@ -21,6 +20,10 @@ export interface WebSearchSource {
 	text?: string; // You.com provides text of webpage right away
 }
 
+export interface WebSearchUsedSource extends WebSearchSource {
+	context: { idx: number; text: string }[];
+}
+
 export type WebSearchMessageSources = {
 	type: "sources";
 	sources: WebSearchSource[];