Add websearch controls for assistants

.env.template

@@ -248,3 +248,5 @@ EXPOSE_API=true
 ALTERNATIVE_REDIRECT_URLS=`[
  huggingchat://login/callback
 ]`
+
+WEBSEARCH_BLOCKLIST=`["youtube.com", "twitter.com"]`

src/lib/components/AssistantSettings.svelte

@@ -71,6 +71,14 @@
  let deleteExistingAvatar = false;
 
  let loading = false;
+
+ let ragMode: false | "links" | "domains" | "all" = assistant?.rag?.allowAll
+ ? "all"
+ : assistant?.rag?.links?.length ?? 0 > 0
+ ? "links"
+ : (assistant?.rag?.allowList?.length ?? 0) > 0
+ ? "domains"
+ : false;
 </script>
 
 <form
@@ -103,6 +111,24 @@
  }
  }
 
+ formData.delete("ragMode");
+
+ if (ragMode === false) {
+ formData.set("ragAllowAll", "false");
+ formData.set("ragLinkList", "");
+ formData.set("ragDomainList", "");
+ } else if (ragMode === "all") {
+ formData.set("ragAllowAll", "true");
+ formData.set("ragLinkList", "");
+ formData.set("ragDomainList", "");
+ } else if (ragMode === "links") {
+ formData.set("ragAllowAll", "false");
+ formData.set("ragDomainList", "");
+ } else if (ragMode === "domains") {
+ formData.set("ragAllowAll", "false");
+ formData.set("ragLinkList", "");
+ }
+
  return async ({ result }) => {
  loading = false;
  await applyAction(result);
@@ -255,7 +281,7 @@
  </label>
  </div>
 
- <label class="flex flex-col">
+ <div class="flex flex-col">
  <span class="mb-1 text-sm font-semibold"> Instructions (system prompt) </span>
  <textarea
  name="preprompt"
@@ -264,7 +290,89 @@
  value={assistant?.preprompt ?? ""}
  />
  <p class="text-xs text-red-500">{getError("preprompt", form)}</p>
- </label>
+ <div class="flex min-h-44 flex-col flex-nowrap">
+ <span class="my-2 text-smd font-semibold"> RAG Settings</span>
+
+ <label>
+ <input
+ checked={!ragMode}
+ on:change={() => (ragMode = false)}
+ type="radio"
+ name="ragMode"
+ value={false}
+ />
+ <span class="my-2 text-sm" class:font-semibold={!ragMode}> Disabled </span>
+ {#if !ragMode}
+ <span class="block text-xs text-gray-500">
+ Assistant won't look for information on the web.
+ </span>
+ {/if}
+ </label>
+
+ <label>
+ <input
+ checked={ragMode === "all"}
+ on:change={() => (ragMode = "all")}
+ type="radio"
+ name="ragMode"
+ value={"all"}
+ />
+ <span class="my-2 text-sm" class:font-semibold={ragMode === "all"}> Enabled </span>
+ {#if ragMode === "all"}
+ <span class="block text-xs text-gray-500">
+ Assistant can access any content on the web.
+ </span>
+ {/if}
+ </label>
+ <label>
+ <input
+ checked={ragMode === "links"}
+ on:change={() => (ragMode = "links")}
+ type="radio"
+ name="ragMode"
+ value={false}
+ />
+ <span class="my-2 text-sm" class:font-semibold={ragMode === "links"}> Links </span>
+ </label>
+ {#if ragMode === "links"}
+ <span class="mb-2 text-xs text-gray-500">
+ Specify max 3 direct URLs the assistant will access. HTML & plaintext only. Separate the
+ list elements with a semicolon.
+ </span>
+ <input
+ name="ragLinkList"
+ class="w-full rounded-lg border-2 border-gray-200 bg-gray-100 p-2"
+ placeholder="https://raw.githubusercontent.com/huggingface/chat-ui/main/README.md"
+ value={assistant?.rag?.links.join(";") ?? ""}
+ />
+ <p class="text-xs text-red-500">{getError("ragLinkList", form)}</p>
+ {/if}
+
+ <label>
+ <input
+ checked={ragMode === "domains"}
+ on:change={() => (ragMode = "domains")}
+ type="radio"
+ name="ragMode"
+ value={false}
+ />
+ <span class="my-2 text-sm" class:font-semibold={ragMode === "domains"}> Domains </span>
+ </label>
+ {#if ragMode === "domains"}
+ <span class="mb-2 text-xs text-gray-500">
+ Specify allowed domains for web search, separe the list elements with a semicolon.
+ </span>
+
+ <input
+ name="ragDomainList"
+ class="w-full rounded-lg border-2 border-gray-200 bg-gray-100 p-2"
+ placeholder="wikipedia.org;bbc.com"
+ value={assistant?.rag?.allowList?.join(";") ?? ""}
+ />
+ <p class="text-xs text-red-500">{getError("ragDomainList", form)}</p>
+ {/if}
+ </div>
+ </div>
  </div>
 
  <div class="mt-5 flex justify-end gap-2">

src/lib/server/websearch/generateQuery.ts

@@ -1,19 +1,6 @@
 import type { Message } from "$lib/types/Message";
 import { format } from "date-fns";
 import { generateFromDefaultEndpoint } from "../generateFromDefaultEndpoint";
-import { WEBSEARCH_ALLOWLIST, WEBSEARCH_BLOCKLIST } from "$env/static/private";
-import { z } from "zod";
-import JSON5 from "json5";
-
-const listSchema = z.array(z.string()).default([]);
-
-const allowList = listSchema.parse(JSON5.parse(WEBSEARCH_ALLOWLIST));
-const blockList = listSchema.parse(JSON5.parse(WEBSEARCH_BLOCKLIST));
-
-const queryModifier = [
- ...allowList.map((item) => "site:" + item),
- ...blockList.map((item) => "-site:" + item),
-].join(" ");
 
 export async function generateQuery(messages: Message[]) {
  const currentDate = format(new Date(), "MMMM d, yyyy");
@@ -79,5 +66,5 @@ Current Question: Where is it being hosted?`,
  preprompt: `You are tasked with generating web search queries. Give me an appropriate query to answer my question for google search. Answer with only the query. Today is ${currentDate}`,
  });
 
- return (queryModifier + " " + webQuery).trim();
+ return webQuery.trim();
 }

src/lib/server/websearch/parseWeb.ts

@@ -3,30 +3,37 @@ import { JSDOM, VirtualConsole } from "jsdom";
 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();
+ const r = await fetch(url, { signal: abortController.signal }).catch();
 
- const virtualConsole = new VirtualConsole();
- virtualConsole.on("error", () => {
- // No-op to skip console errors.
- });
+ if (r.headers.get("content-type")?.includes("text/html")) {
+ 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,
- });
+  // put the html string into a DOM
+  const dom = new JSDOM((await r.text()) ?? "", {
+  virtualConsole,
+  });
 
- const { document } = dom.window;
- const textElTags = "p";
- const paragraphs = document.querySelectorAll(textElTags);
- if (!paragraphs.length) {
- throw new Error(`webpage doesn't have any "${textElTags}" element`);
- }
- const paragraphTexts = Array.from(paragraphs).map((p) => p.textContent);
+  const { document } = dom.window;
+  const textElTags = "p";
+  const paragraphs = document.querySelectorAll(textElTags);
+  if (!paragraphs.length) {
+  throw new Error(`webpage doesn't have any "${textElTags}" element`);
+  }
+  const paragraphTexts = Array.from(paragraphs).map((p) => p.textContent);
 
- // combine text contents from paragraphs and then remove newlines and multiple spaces
- const text = paragraphTexts.join(" ").replace(/ {2}|\r\n|\n|\r/gm, "");
+  // combine text contents from paragraphs and then remove newlines and multiple spaces
+  const text = paragraphTexts.join(" ").replace(/ {2}|\r\n|\n|\r/gm, "");
 
- return text;
+ return text;
+ } else if (
+ r.headers.get("content-type")?.includes("text/plain") ||
+ r.headers.get("content-type")?.includes("text/markdown")
+ ) {
+ return r.text();
+ } else {
+ throw new Error("Unsupported content type");
+ }
 }

src/lib/server/websearch/runWebSearch.ts

@@ -13,12 +13,21 @@ import { defaultEmbeddingModel, embeddingModels } from "$lib/server/embeddingMod
 const MAX_N_PAGES_SCRAPE = 10 as const;
 const MAX_N_PAGES_EMBED = 5 as const;
 
-const DOMAIN_BLOCKLIST = ["youtube.com", "twitter.com"];
+import { WEBSEARCH_ALLOWLIST, WEBSEARCH_BLOCKLIST } from "$env/static/private";
+import { z } from "zod";
+import JSON5 from "json5";
+import type { Assistant } from "$lib/types/Assistant";
+
+const listSchema = z.array(z.string()).default([]);
+
+const allowList = listSchema.parse(JSON5.parse(WEBSEARCH_ALLOWLIST));
+const blockList = listSchema.parse(JSON5.parse(WEBSEARCH_BLOCKLIST));
 
 export async function runWebSearch(
  conv: Conversation,
  prompt: string,
- updatePad: (upd: MessageUpdate) => void
+ updatePad: (upd: MessageUpdate) => void,
+ ragSettings?: Assistant["rag"]
 ) {
  const messages = (() => {
  return [...conv.messages, { content: prompt, from: "user", id: crypto.randomUUID() }];
@@ -39,26 +48,48 @@ export async function runWebSearch(
  }
 
  try {
- webSearch.searchQuery = await generateQuery(messages);
- const searchProvider = getWebSearchProvider();
- appendUpdate(`Searching ${searchProvider}`, [webSearch.searchQuery]);
- const results = await searchWeb(webSearch.searchQuery);
- webSearch.results =
- (results.organic_results &&
- results.organic_results.map((el: { title?: string; link: string; text?: string }) => {
- try {
- const { title, link, text } = el;
- const { hostname } = new URL(link);
- return { title, link, hostname, text };
- } catch (e) {
- // Ignore Errors
- return null;
- }
- })) ??
- [];
+ // if the assistant specified direct links, skip the websearch
+ if (ragSettings && ragSettings?.links.length > 0) {
+ appendUpdate("Using links specified in assistant directly. Skipping websearch");
+ webSearch.results = ragSettings.links.map((link) => {
+ return { link, hostname: new URL(link).hostname, title: "", text: "" };
+ });
+ } else {
+ webSearch.searchQuery = await generateQuery(messages);
+ const searchProvider = getWebSearchProvider();
+ appendUpdate(`Searching ${searchProvider}`, [webSearch.searchQuery]);
+
+ if (ragSettings && ragSettings?.allowList.length > 0) {
+ appendUpdate("Filtering results to only domains specified in assistant");
+ webSearch.searchQuery +=
+ " " + ragSettings.allowList.map((item) => "site:" + item).join(" ");
+ }
+
+ // handle the global lists
+ webSearch.searchQuery +=
+ allowList.map((item) => "site:" + item).join(" ") +
+ " " +
+ blockList.map((item) => "-site:" + item).join(" ");
+
+ const results = await searchWeb(webSearch.searchQuery);
+ webSearch.results =
+ (results.organic_results &&
+ results.organic_results.map((el: { title?: string; link: string; text?: string }) => {
+ try {
+ const { title, link, text } = el;
+ const { hostname } = new URL(link);
+ return { title, link, hostname, text };
+ } catch (e) {
+ // Ignore Errors
+ return null;
+ }
+ })) ??
+ [];
+ }
+
  webSearch.results = webSearch.results.filter((value) => value !== null);
  webSearch.results = webSearch.results
- .filter(({ link }) => !DOMAIN_BLOCKLIST.some((el) => link.includes(el))) // filter out blocklist links
+ .filter(({ link }) => !blockList.some((el) => link.includes(el))) // filter out blocklist links
  .slice(0, MAX_N_PAGES_SCRAPE); // limit to first 10 links only
 
  // fetch the model

src/lib/types/Assistant.ts

@@ -14,4 +14,9 @@ export interface Assistant extends Timestamps {
  preprompt: string;
  userCount?: number;
  featured?: boolean;
+ rag?: {
+ allowAll: boolean;
+ allowList: string[];
+ links: string[];
+ };
 }

src/lib/utils/parseStringToList.ts

@@ -0,0 +1,10 @@
+export function parseStringToList(links: unknown): string[] {
+ if (typeof links !== "string") {
+ return [];
+ }
+
+ return links
+ .split(";")
+ .map((link) => link.trim())
+ .filter((link) => link.length > 0);
+}

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

@@ -250,13 +250,31 @@ export async function POST({ request, locals, params, getClientAddress }) {
 
  let webSearchResults: WebSearch | undefined;
 
- if (webSearch && !isContinue && !conv.assistantId) {
- webSearchResults = await runWebSearch(conv, messages[messages.length - 1].content, update);
+ // check if assistant has a rag
+ const assistantRAG = await collections.assistants
+ .findOne({ _id: conv.assistantId })
+ .then((a) => {
+ return a?.rag;
+ });
+
+ const assistantHasRAG =
+ assistantRAG &&
+ (assistantRAG.links.length > 0 ||
+ assistantRAG.allowList.length > 0 ||
+ assistantRAG.allowAll);
+
+ // if websearch is enabled and the assistant is not specified or it is and has a rag
+ if (!isContinue && ((webSearch && !conv.assistantId) || assistantHasRAG)) {
+ webSearchResults = await runWebSearch(
+ conv,
+ messages[messages.length - 1].content,
+ update,
+ assistantRAG
+ );
  messages[messages.length - 1].webSearch = webSearchResults;
  } else if (isContinue) {
  webSearchResults = messages[messages.length - 1].webSearch;
  }
-
  conv.messages = messages;
 
  const previousContent = isContinue