[Websearch] change context schema (#944)

huggingface · Mar 25, 2024 · 2e2f16c · 2e2f16c
1 parent 506e596
commit 2e2f16c
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 10 deletions.
diff --git a/src/lib/server/preprocessMessages.ts b/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}`;

diff --git a/src/lib/server/websearch/runWebSearch.ts b/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({

diff --git 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[];