Generlize RAG

Author: mishig25

src/lib/server/pdfSearch.ts renamed to src/lib/server/rag/pdfchat/rag.ts

@@ -1,18 +1,20 @@
-import type { PdfSearch } from "$lib/types/PdfChat";
 import { createEmbeddings, findSimilarSentences } from "$lib/server/embeddings";
 import type { Conversation } from "$lib/types/Conversation";
 import type { MessageUpdate } from "$lib/types/MessageUpdate";
-import { downloadPdfEmbeddings } from "./files/downloadFile";
+import { downloadPdfEmbeddings } from "../../files/downloadFile";
 import { Tensor } from "@xenova/transformers";
+import type { RAG } from "../RAG";
+import type { RagContext } from "$lib/types/rag";
 
 // todo: embed the prompt, download the embeddings, serialize them, and find the closest sentences, and get their texts, lets go
-export async function runPdfSearch(
+async function runPdfSearch(
 	conv: Conversation,
 	prompt: string,
 	updatePad: (upd: MessageUpdate) => void
 ) {
-	const pdfSearch: PdfSearch = {
+	const pdfSearch: RagContext = {
 		context: "",
+		type: "pdfchat",
 		createdAt: new Date(),
 		updatedAt: new Date(),
 	};
@@ -50,3 +52,8 @@ export async function runPdfSearch(
 
 	return pdfSearch;
 }
+
+export const ragPdfchat: RAG = {
+	type: "pdfchat", 
+	retrieveRagContext: runPdfSearch,
+}

src/lib/server/rag/rag.ts

@@ -0,0 +1,25 @@
+import type { Conversation } from "$lib/types/Conversation";
+import type { MessageUpdate } from "$lib/types/MessageUpdate";
+import type { RagContext } from "$lib/types/rag";
+import { ragPdfchat } from "./pdfchat/rag";
+import { ragWebsearch } from "./websearch/rag";
+
+type RetrieveRagContext<T=RagContext> = (conv: Conversation, prompt: string, updatePad: (upd: MessageUpdate) => void) => Promise<T>;
+
+export type RAGType = "websearch" | "pdfchat";
+
+export interface RAG {
+    type: RAGType, // should be one of websearch or somethign, make it stronger typed
+    retrieveRagContext: RetrieveRagContext,
+    buildPrompt?: string,
+}
+
+// list of all rags
+export const RAGs: {
+	[Key in RAGType]: RAG;
+} = {
+	websearch: ragWebsearch,
+	pdfchat: ragPdfchat,
+};
+
+export default RAGs;

src/lib/server/websearch/generateQuery.ts renamed to src/lib/server/rag/websearch/generateQuery.ts

@@ -1,6 +1,6 @@
 import type { Message } from "$lib/types/Message";
 import { format } from "date-fns";
-import { generateFromDefaultEndpoint } from "../generateFromDefaultEndpoint";
+import { generateFromDefaultEndpoint } from "../../generateFromDefaultEndpoint";
 import { WEBSEARCH_ALLOWLIST, WEBSEARCH_BLOCKLIST } from "$env/static/private";
 import { z } from "zod";
 import JSON5 from "json5";

src/lib/server/websearch/parseWeb.ts renamed to src/lib/server/rag/websearch/parseWeb.ts

@@ -1,21 +1,22 @@
-import { searchWeb } from "$lib/server/websearch/searchWeb";
+import { searchWeb } from "$lib/server/rag/websearch/searchWeb";
 import type { Message } from "$lib/types/Message";
 import type { WebSearch, WebSearchSource } from "$lib/types/WebSearch";
-import { generateQuery } from "$lib/server/websearch/generateQuery";
-import { parseWeb } from "$lib/server/websearch/parseWeb";
+import { generateQuery } from "$lib/server/rag/websearch/generateQuery";
+import { parseWeb } from "$lib/server/rag/websearch/parseWeb";
 import { chunk } from "$lib/utils/chunk";
 import { findSimilarSentences } from "$lib/server/sentenceSimilarity";
 import type { Conversation } from "$lib/types/Conversation";
 import type { MessageUpdate } from "$lib/types/MessageUpdate";
 import { getWebSearchProvider } from "./searchWeb";
 import { defaultEmbeddingModel, embeddingModels } from "$lib/server/embeddingModels";
+import type { RAG } from "../RAG";
 
 const MAX_N_PAGES_SCRAPE = 10 as const;
 const MAX_N_PAGES_EMBED = 5 as const;
 
 const DOMAIN_BLOCKLIST = ["youtube.com", "twitter.com"];
 
-export async function runWebSearch(
+async function runWebSearch(
 	conv: Conversation,
 	prompt: string,
 	updatePad: (upd: MessageUpdate) => void
@@ -25,6 +26,7 @@ export async function runWebSearch(
 	})() satisfies Message[];
 
 	const webSearch: WebSearch = {
+		type: "websearch",
 		prompt: prompt,
 		searchQuery: "",
 		results: [],
@@ -130,3 +132,8 @@ export async function runWebSearch(
 
 	return webSearch;
 }
+
+export const ragWebsearch: RAG = {
+	type: "websearch", 
+	retrieveRagContext: runWebSearch,
+}

src/lib/server/websearch/runWebSearch.ts renamed to src/lib/server/rag/websearch/rag.ts

@@ -1,5 +1,5 @@
-import type { YouWebSearch } from "../../types/WebSearch";
-import { WebSearchProvider } from "../../types/WebSearch";
+import type { YouWebSearch } from "../../../types/WebSearch";
+import { WebSearchProvider } from "../../../types/WebSearch";
 import {
 	SERPAPI_KEY,
 	SERPER_API_KEY,

src/lib/server/websearch/searchWeb.ts renamed to src/lib/server/rag/websearch/searchWeb.ts

@@ -1,11 +1,3 @@
-import type { ObjectId } from "mongodb";
-import type { Timestamps } from "./Timestamps";
-
-export interface PdfSearch extends Timestamps {
-	_id?: ObjectId;
-	context: string;
-}
-
 /* eslint-disable no-shadow */
 export enum PdfUploadStatus {
 	Ready = "Ready",

src/lib/server/websearch/searchWebLocal.ts renamed to src/lib/server/rag/websearch/searchWebLocal.ts

@@ -1,16 +1,9 @@
-import type { ObjectId } from "mongodb";
-import type { Conversation } from "./Conversation";
-import type { Timestamps } from "./Timestamps";
-
-export interface WebSearch extends Timestamps {
-	_id?: ObjectId;
-	convId?: Conversation["_id"];
+import type { RagContext } from "./rag";
 
+export interface WebSearch extends RagContext {
 	prompt: string;
-
 	searchQuery: string;
 	results: WebSearchSource[];
-	context: string;
 	contextSources: WebSearchSource[];
 }
 

src/lib/types/PdfChat.ts

@@ -0,0 +1,11 @@
+import type { ObjectId } from "mongodb";
+import type { Conversation } from "./Conversation";
+import type { Timestamps } from "./Timestamps";
+import type { RAGType } from "$lib/server/rag/RAG";
+
+export interface RagContext extends Timestamps {
+	_id?: ObjectId;
+	convId?: Conversation["_id"];
+	type: RAGType
+	context: string;
+}

src/lib/types/WebSearch.ts

@@ -8,14 +8,13 @@ import { error } from "@sveltejs/kit";
 import { ObjectId } from "mongodb";
 import { z } from "zod";
 import type { MessageUpdate } from "$lib/types/MessageUpdate";
-import { runWebSearch } from "$lib/server/websearch/runWebSearch";
-import { runPdfSearch } from "$lib/server/pdfSearch";
 import type { WebSearch } from "$lib/types/WebSearch";
 import type { PdfSearch } from "$lib/types/PdfChat";
 import { abortedGenerations } from "$lib/server/abortedGenerations";
 import { summarize } from "$lib/server/summarize";
 import { uploadImgFile } from "$lib/server/files/uploadFile";
 import sizeof from "image-size";
+import RAGs from "$lib/server/rag/RAG";
 
 export async function POST({ request, locals, params, getClientAddress }) {
 	const id = z.string().parse(params.id);
@@ -237,15 +236,15 @@ export async function POST({ request, locals, params, getClientAddress }) {
 			let webSearchResults: WebSearch | undefined;
 
 			if (webSearch) {
-				webSearchResults = await runWebSearch(conv, newPrompt, update);
+				webSearchResults = await RAGs["websearch"].retrieveRagContext(conv, newPrompt, update);
 			}
 
 			messages[messages.length - 1].webSearch = webSearchResults;
 
 			let pdfSearchResults: PdfSearch | undefined;
 			const pdfSearch = await collections.files.findOne({ filename: `${convId.toString()}-pdf` });
 			if (pdfSearch) {
-				pdfSearchResults = await runPdfSearch(conv, newPrompt, update);
+				pdfSearchResults = await RAGs["websearch"].retrieveRagContext(conv, newPrompt, update);
 			}
 
 			messages[messages.length - 1].pdfSearch = pdfSearchResults;