diff --git a/src/lib/server/preprocessMessages.ts b/src/lib/server/preprocessMessages.ts index 74346e2fca..53768fa6f6 100644 --- a/src/lib/server/preprocessMessages.ts +++ b/src/lib/server/preprocessMessages.ts @@ -11,8 +11,14 @@ export async function preprocessMessages( ): Promise { 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}`; diff --git a/src/lib/server/websearch/runWebSearch.ts b/src/lib/server/websearch/runWebSearch.ts index 27f361a080..501c2f28ae 100644 --- a/src/lib/server/websearch/runWebSearch.ts +++ b/src/lib/server/websearch/runWebSearch.ts @@ -153,14 +153,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(); 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({ diff --git a/src/lib/types/WebSearch.ts b/src/lib/types/WebSearch.ts index e4ea514fcd..b7e3c40eb9 100644 --- a/src/lib/types/WebSearch.ts +++ b/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[];