Add allow & blocklist for websearch (#651)

* Add allow & blocklist for websearch

* fix start whitespace

Author: nsarrazin

.env

@@ -18,6 +18,9 @@ SERPER_API_KEY=#your serper.dev api key here
 SERPAPI_KEY=#your serpapi key here
 USE_LOCAL_WEBSEARCH=#set to true to parse google results yourself, overrides other API keys
 
+WEBSEARCH_ALLOWLIST=`[]` # if it's defined, allow websites from only this list.
+WEBSEARCH_BLOCKLIST=`[]` # if it's defined, block websites from this list.
+
 # Parameters to enable open id login
 OPENID_CONFIG=`{
   "PROVIDER_URL": "",

src/lib/server/websearch/generateQuery.ts

@@ -1,6 +1,18 @@
 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";
+
+const listSchema = z.array(z.string()).default([]);
+
+const allowList = listSchema.parse(JSON.parse(WEBSEARCH_ALLOWLIST));
+const blockList = listSchema.parse(JSON.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");
@@ -61,8 +73,10 @@ Current Question: Where is it being hosted ?`,
 		},
 	];
 
-	return await generateFromDefaultEndpoint({
+	const webQuery = await generateFromDefaultEndpoint({
 		messages: convQuery,
 		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();
 }