Generlize RAG (#689)

* Generlize RAG

* wip

* fix casting

Author: Mishig

src/lib/buildPrompt.ts

@@ -1,80 +1,32 @@
 import type { BackendModel } from "./server/models";
 import type { Message } from "./types/Message";
-import { format } from "date-fns";
-import type { WebSearch } from "./types/WebSearch";
-import type { PdfSearch } from "./types/PdfChat";
 import { downloadImgFile } from "./server/files/downloadFile";
 import type { Conversation } from "./types/Conversation";
+import RAGs from "./server/rag/rag";
+import type { RagContext } from "./types/rag";
+
+export type BuildPromptMessage = Pick<Message, "from" | "content" | "files">;
 
 interface buildPromptOptions {
-	messages: Pick<Message, "from" | "content" | "files">[];
+	messages: BuildPromptMessage[];
 	id?: Conversation["_id"];
 	model: BackendModel;
 	locals?: App.Locals;
-	webSearch?: WebSearch;
-	pdfSearch?: PdfSearch;
+	ragContext?: RagContext;
 	preprompt?: string;
 	files?: File[];
 }
 
 export async function buildPrompt({
 	messages,
 	model,
-	webSearch,
-	pdfSearch,
+	ragContext,
 	preprompt,
 	id,
 }: buildPromptOptions): Promise<string> {
-	if (webSearch && webSearch.context) {
-		const lastMsg = messages.slice(-1)[0];
-		const messagesWithoutLastUsrMsg = messages.slice(0, -1);
-		const previousUserMessages = messages.filter((el) => el.from === "user").slice(0, -1);
-
-		const previousQuestions =
-			previousUserMessages.length > 0
-				? `Previous questions: \n${previousUserMessages
-						.map(({ content }) => `- ${content}`)
-						.join("\n")}`
-				: "";
-		const currentDate = format(new Date(), "MMMM d, yyyy");
-		messages = [
-			...messagesWithoutLastUsrMsg,
-			{
-				from: "user",
-				content: `I searched the web using the query: ${webSearch.searchQuery}. Today is ${currentDate} and here are the results:
-				=====================
-				${webSearch.context}
-				=====================
-				${previousQuestions}
-				Answer the question: ${lastMsg.content} 
-				`,
-			},
-		];
-	} else if (pdfSearch && pdfSearch.context) {
-		const lastMsg = messages.slice(-1)[0];
-		const messagesWithoutLastUsrMsg = messages.slice(0, -1);
-		const previousUserMessages = messages.filter((el) => el.from === "user").slice(0, -1);
-
-		const previousQuestions =
-			previousUserMessages.length > 0
-				? `Previous questions: \n${previousUserMessages
-						.map(({ content }) => `- ${content}`)
-						.join("\n")}`
-				: "";
-
-		messages = [
-			...messagesWithoutLastUsrMsg,
-			{
-				from: "user",
-				content: `Below are the information I extracted from a PDF file that might be useful:
-				=====================
-				${pdfSearch.context}
-				=====================
-				${previousQuestions}
-				Answer the question: ${lastMsg.content} 
-				`,
-			},
-		];
+	if (ragContext) {
+		const { type: ragType } = ragContext;
+		messages = RAGs[ragType].buildPrompt(messages, ragContext);
 	}
 
 	// section to handle potential files input

src/lib/components/OpenPdfSearchResults.svelte

@@ -1,5 +1,6 @@
 <script lang="ts">
-	import type { WebSearchUpdate } from "$lib/types/MessageUpdate";
+	import type { RAGUpdate } from "$lib/types/MessageUpdate";
+	import type { RAGType } from "$lib/types/rag";
 	import CarbonCaretRight from "~icons/carbon/caret-right";
 
 	import CarbonCheckmark from "~icons/carbon/checkmark-filled";
@@ -9,11 +10,22 @@
 
 	export let loading = false;
 	export let classNames = "";
-	export let webSearchMessages: WebSearchUpdate[] = [];
+	export let ragUpdates: RAGUpdate[] = [];
+
+	const TITLE_MAPPING: Record<RAGType, string> = {
+		webSearch: "Web search",
+		pdfChat: "PDF Chat",
+	};
 
 	let detailsOpen: boolean;
 	let error: boolean;
-	$: error = webSearchMessages[webSearchMessages.length - 1]?.messageType === "error";
+	$: error = ragUpdates[ragUpdates.length - 1]?.messageType === "error";
+
+	$: ragType = ragUpdates[0].type;
+
+	$: loading =
+		ragUpdates.length > 0 &&
+		!["sources", "done"].includes(ragUpdates[ragUpdates.length - 1].messageType);
 </script>
 
 <details
@@ -31,21 +43,21 @@
 			<CarbonCheckmark class="my-auto text-gray-500" />
 		{/if}
 		<span class="px-2 font-medium" class:text-red-700={error} class:dark:text-red-500={error}>
-			Web search
+			{TITLE_MAPPING[ragType]}
 		</span>
 		<div class="my-auto transition-all" class:rotate-90={detailsOpen}>
 			<CarbonCaretRight />
 		</div>
 	</summary>
 
 	<div class="content px-5 pb-5 pt-4">
-		{#if webSearchMessages.length === 0}
+		{#if ragUpdates.length === 0}
 			<div class="mx-auto w-fit">
 				<EosIconsLoading class="mb-3 h-4 w-4" />
 			</div>
 		{:else}
 			<ol>
-				{#each webSearchMessages as message}
+				{#each ragUpdates as message}
 					{#if message.messageType === "update"}
 						<li class="group border-l pb-6 last:!border-transparent last:pb-0 dark:border-gray-800">
 							<div class="flex items-start">

src/lib/components/OpenWebSearchResults.svelte renamed to src/lib/components/OpenRAGResults.svelte

@@ -16,9 +16,9 @@
 	import { PUBLIC_SEP_TOKEN } from "$lib/constants/publicSepToken";
 	import type { Model } from "$lib/types/Model";
 
-	import OpenWebSearchResults from "../OpenWebSearchResults.svelte";
-	import OpenPdfSearchResults from "../OpenPdfSearchResults.svelte";
-	import type { RAGUpdate, WebSearchUpdate, PdfSearchUpdate } from "$lib/types/MessageUpdate";
+	import OpenRAGResults from "../OpenRAGResults.svelte";
+	import type { RAGUpdate } from "$lib/types/MessageUpdate";
+	import { ragTypes } from "$lib/types/rag";
 
 	function sanitizeMd(md: string) {
 		let ret = md
@@ -61,6 +61,7 @@
 	let loadingEl: IconLoading;
 	let pendingTimeout: ReturnType<typeof setTimeout>;
 	let isCopied = false;
+	let ragIsLoading = false;
 
 	const renderer = new marked.Renderer();
 	// For code blocks with simple backticks
@@ -107,29 +108,17 @@
 		}
 	});
 
-	let searchUpdates: WebSearchUpdate[] = [];
+	let ragUpdates: RAGUpdate[] = [];
 
-	$: searchUpdates = ((RAGMessages.filter(({ type }) => type === "webSearch").length > 0
-		? RAGMessages.filter(({ type }) => type === "webSearch")
-		: message.updates?.filter(({ type }) => type === "webSearch")) ?? []) as WebSearchUpdate[];
-
-	let pdfUpdates: PdfSearchUpdate[] = [];
-
-	$: pdfUpdates = ((RAGMessages.filter(({ type }) => type === "pdfSearch").length > 0
-		? RAGMessages.filter(({ type }) => type === "pdfSearch")
-		: message.updates?.filter(({ type }) => type === "pdfSearch")) ?? []) as PdfSearchUpdate[];
+	$: ragUpdates = ((RAGMessages.length > 0
+		? RAGMessages
+		: message.updates?.filter(({ type }) => ragTypes.includes(type))) ?? []) as RAGUpdate[];
 
 	$: downloadLink =
 		message.from === "user" ? `${$page.url.pathname}/message/${message.id}/prompt` : undefined;
 
-	let webSearchIsDone = true;
-
-	$: webSearchIsDone =
-		searchUpdates.length > 0 && searchUpdates[searchUpdates.length - 1].messageType === "sources";
-
 	$: webSearchSources =
-		searchUpdates &&
-		searchUpdates?.filter(({ messageType }) => messageType === "sources")?.[0]?.sources;
+		ragUpdates && ragUpdates?.filter(({ messageType }) => messageType === "sources")?.[0]?.sources;
 
 	$: if (isCopied) {
 		setTimeout(() => {
@@ -153,21 +142,14 @@
 		<div
 			class="relative min-h-[calc(2rem+theme(spacing[3.5])*2)] min-w-[60px] break-words rounded-2xl border border-gray-100 bg-gradient-to-br from-gray-50 px-5 py-3.5 text-gray-600 prose-pre:my-2 dark:border-gray-800 dark:from-gray-800/40 dark:text-gray-300"
 		>
-			{#if searchUpdates && searchUpdates.length > 0}
-				<OpenWebSearchResults
-					classNames={tokens.length ? "mb-3.5" : ""}
-					webSearchMessages={searchUpdates}
-					loading={!(searchUpdates[searchUpdates.length - 1]?.messageType === "sources")}
-				/>
-			{/if}
-			{#if pdfUpdates && pdfUpdates.length > 0}
-				<OpenPdfSearchResults
+			{#if ragUpdates && ragUpdates.length > 0}
+				<OpenRAGResults
 					classNames={tokens.length ? "mb-3.5" : ""}
-					pdfSearchMessages={pdfUpdates}
-					loading={!(pdfUpdates[pdfUpdates.length - 1]?.messageType === "done")}
+					{ragUpdates}
+					bind:loading={ragIsLoading}
 				/>
 			{/if}
-			{#if !message.content && (webSearchIsDone || (RAGMessages && RAGMessages.length === 0))}
+			{#if !message.content && (!ragIsLoading || (RAGMessages && RAGMessages.length === 0))}
 				<IconLoading />
 			{/if}
 

src/lib/components/chat/ChatMessage.svelte

@@ -39,8 +39,7 @@ export async function endpointAws(
 	return async ({ conversation }) => {
 		const prompt = await buildPrompt({
 			messages: conversation.messages,
-			webSearch: conversation.messages[conversation.messages.length - 1].webSearch,
-			pdfSearch: conversation.messages[conversation.messages.length - 1].pdfSearch,
+			ragContext: conversation.messages[conversation.messages.length - 1].ragContext,
 			preprompt: conversation.preprompt,
 			model,
 		});

src/lib/server/endpoints/aws/endpointAws.ts

@@ -22,8 +22,7 @@ export function endpointLlamacpp(
 	return async ({ conversation }) => {
 		const prompt = await buildPrompt({
 			messages: conversation.messages,
-			webSearch: conversation.messages[conversation.messages.length - 1].webSearch,
-			pdfSearch: conversation.messages[conversation.messages.length - 1].pdfSearch,
+			ragContext: conversation.messages[conversation.messages.length - 1].ragContext,
 			preprompt: conversation.preprompt,
 			model,
 		});

src/lib/server/endpoints/llamacpp/endpointLlamacpp.ts

@@ -17,8 +17,7 @@ export function endpointOllama(input: z.input<typeof endpointOllamaParametersSch
 	return async ({ conversation }) => {
 		const prompt = await buildPrompt({
 			messages: conversation.messages,
-			webSearch: conversation.messages[conversation.messages.length - 1].webSearch,
-			pdfSearch: conversation.messages[conversation.messages.length - 1].pdfSearch,
+			ragContext: conversation.messages[conversation.messages.length - 1].ragContext,
 			preprompt: conversation.preprompt,
 			model,
 		});

src/lib/server/endpoints/ollama/endpointOllama.ts

@@ -40,8 +40,7 @@ export async function endpointOai(
 					model: model.id ?? model.name,
 					prompt: await buildPrompt({
 						messages: conversation.messages,
-						webSearch: conversation.messages[conversation.messages.length - 1].webSearch,
-						pdfSearch: conversation.messages[conversation.messages.length - 1].pdfSearch,
+						ragContext: conversation.messages[conversation.messages.length - 1].ragContext,
 						preprompt: conversation.preprompt,
 						model,
 					}),

src/lib/server/endpoints/openai/endpointOai.ts

@@ -18,8 +18,7 @@ export function endpointTgi(input: z.input<typeof endpointTgiParametersSchema>):
 	return async ({ conversation }) => {
 		const prompt = await buildPrompt({
 			messages: conversation.messages,
-			webSearch: conversation.messages[conversation.messages.length - 1].webSearch,
-			pdfSearch: conversation.messages[conversation.messages.length - 1].pdfSearch,
+			ragContext: conversation.messages[conversation.messages.length - 1].ragContext,
 			preprompt: conversation.preprompt,
 			model,
 			id: conversation._id,