WIP: RAG

Author: sunjay-dev

backend/package.json

@@ -41,7 +41,8 @@
     "langchain": "^1.2.29",
     "pino": "^10.3.1",
     "pino-pretty": "^13.1.3",
-    "prom-client": "^15.1.3"
+    "prom-client": "^15.1.3",
+    "zod": "^4.3.6"
   },
   "devDependencies": {
     "@types/node": "^25.3.3",

backend/src/config/groq.config.ts

@@ -33,16 +33,18 @@ export const groqChatAgent = (model: string, temperature = 0) => {
   });
 };
 
-export const createGroqAgent = (model: string, systemPrompt: string) => {
+export const createGroqAgent = (model: string, systemPrompt: string, isRag: boolean = false) => {
   const llm = getLLM(model);
   const summarizerLLM = getSummarizer();
 
   const modelConfig = MODELS.find((m) => model === m.id);
   const triggerTokens = modelConfig ? Math.floor(modelConfig.tpm * 0.5) : 3000;
 
+  const agentTools = isRag ? tools : tools.filter((t) => t.name !== "search_uploaded_documents");
+
   return createAgent({
     model: llm,
-    tools,
+    tools: agentTools,
     systemPrompt,
     checkpointer,
     middleware: [

backend/src/controllers/message.controllers.ts

@@ -6,7 +6,7 @@ import { streamText as honoStreamText } from "hono/streaming";
 import logger from "../utils/logger.utils.js";
 import { streamLoading, streamText, streamError } from "../utils/stream.utils.js";
 import { type UploadedFile } from "@app/shared/src/schemas/message.schema.js";
-import { vectorStore } from "../config/vectorStore.config.js";
+
 
 export async function handleUserMessageResponse(c: Context) {
   const requestStartTime = new Date();
@@ -33,30 +33,14 @@ export async function handleUserMessageResponse(c: Context) {
         return;
       }
 
-      if (chat.isRag && query.trim() !== "") {
-        try {
-          await streamLoading(stream, "Searching documents...");
-
-          const results = await vectorStore.similaritySearch(query, 4, {
-            chatId: chat.id,
-          });
-
-          if (results.length > 0) {
-            const context = results.map((r) => r.pageContent).join("\n\n---\n\n");
-            finalQuery = `Context Information:\n${context}\n\nUser Query: ${finalQuery}\n\nPlease answer the user query using the context information provided above if it is relevant.`;
-          }
-        } catch (err) {
-          logger.error({ message: "Vector search failed", error: err, chatId });
-        }
-      }
-
       const aiStream = generateAIResponse({
         query: finalQuery,
         threadId: chatId,
         modelName: model,
         preferences,
         memories,
         timezone,
+        isRag: chat.isRag,
       });
 
       let isFirstChunk = true;

backend/src/tools/index.ts

@@ -1,5 +1,6 @@
 import { InternetSearch } from "./internetSearch.tools.js";
 // import { getCurrentTime } from "./time.tools.js";
 import { getDailyWeatherForecast, getCurrentWeather } from "./weather.tools.js";
+import { searchUploadedDocuments } from "./ragSearch.tools.js";
 
-export default [InternetSearch, getCurrentWeather, getDailyWeatherForecast];
+export default [InternetSearch, getCurrentWeather, getDailyWeatherForecast, searchUploadedDocuments];

backend/src/tools/ragSearch.tools.ts

@@ -0,0 +1,44 @@
+import { tool } from "@langchain/core/tools";
+import { z } from "@app/shared/src/index.js";
+import { vectorStore } from "../config/vectorStore.config.js";
+import logger from "../utils/logger.utils.js";
+import { RunnableConfig } from "@langchain/core/runnables";
+
+const searchUploadedDocumentsSchema = z.object({
+    query: z.string().describe("The search query to find relevant information from the uploaded documents in the current chat."),
+});
+
+export const searchUploadedDocuments = tool(
+    async (input, config?: RunnableConfig) => {
+        try {
+            const chatId = config?.configurable?.thread_id;
+
+            if (!chatId) {
+                logger.warn("search_uploaded_documents tool called without a thread_id in config.");
+                return "Error: Could not determine the current chat context to search documents.";
+            }
+
+            logger.info({ message: "Executing AI-driven vector search", query: input.query, chatId });
+
+            const results = await vectorStore.similaritySearch(input.query, 4, {
+                chatId: chatId,
+            });
+
+            if (!results || results.length === 0) {
+                return "No relevant information found in the uploaded documents for this query.";
+            }
+
+            const context = results.map((r) => r.pageContent).join("\n\n---\n\n");
+            return `Context Information from uploaded documents: \n\n${context} `;
+
+        } catch (error) {
+            logger.error({ message: "AI vector search failed", error, query: input.query });
+            return "An error occurred while searching the documents. Please try again or inform the user.";
+        }
+    },
+    {
+        name: "search_uploaded_documents",
+        description: "Search the uploaded documents within the current chat for specific information or context. Use this tool whenever the user asks about the content of their uploaded files, PDFs, presentations, or documents.",
+        schema: searchUploadedDocumentsSchema,
+    }
+);

backend/src/utils/model.utils.ts

@@ -20,10 +20,11 @@ type Props = {
   preferences: UserPreference;
   memories: Memories[];
   timezone: string;
+  isRag: boolean;
 };
 
-export async function* generateAIResponse({ query, threadId, modelName, preferences, memories, timezone }: Props) {
-  const agent = createGroqAgent(modelName, systemPrompt(preferences, memories, timezone));
+export async function* generateAIResponse({ query, threadId, modelName, preferences, memories, timezone, isRag }: Props) {
+  const agent = createGroqAgent(modelName, systemPrompt(preferences, memories, timezone), isRag);
 
   const config = { configurable: { thread_id: threadId } };
 

frontend/src/components/containers/InputContainer.tsx

@@ -134,7 +134,7 @@ export default function InputContainer({ sendMessage, isGenerating, selectedMode
       return;
     }
 
-    if (isGlobalUploading) return;
+    if (isGlobalUploading || isRagProcessing) return;
 
     const successfulUploads = attachments.filter((a) => a.status === "success" && a.uploadData).map((a) => a.uploadData!);
 
@@ -228,10 +228,10 @@ export default function InputContainer({ sendMessage, isGenerating, selectedMode
 
   const isGlobalUploading = attachments.some((a) => a.status === "uploading");
   const hasContent = message?.trim().length > 0 || attachments.length > 0;
-  const isBusy = isGenerating || isGlobalUploading || isRagProcessing;
+  const isBusy = isGenerating || isGlobalUploading;
 
   return (
-    <div className="w-full flex justify-center items-center px-4 pb-4 bg-transparent relative">
+    <div className="w-full flex justify-center items-center px-4 pb-6 sm:pb-4 bg-transparent relative">
       {isDragging && (
         <div className="fixed inset-0 z-99 flex items-center justify-center bg-black/40 backdrop-blur-sm pointer-events-none transition-all duration-200">
           <div className="absolute inset-4 sm:inset-6 md:inset-8 border-4 border-dashed border-white/60 rounded-3xl z-0" />
@@ -246,8 +246,8 @@ export default function InputContainer({ sendMessage, isGenerating, selectedMode
       <div className="w-full max-w-svw sm:max-w-180 transition-all duration-300 ease-in-out">
         <form onSubmit={handleSubmit}>
           <div className="bg-white rounded-2xl border border-primary px-3 py-2 transition-all duration-200 ease-in-out shadow-md">
-            {attachments.length > 0 && (
-              <div className="flex flex-wrap gap-2 mb-2 pl-1">
+            {(attachments.length > 0 || isRagProcessing) && (
+              <div className="flex flex-wrap gap-2 mb-2 pl-1 items-center">
                 {attachments.map((att) => (
                   <div key={att.id} className="relative group">
                     {att.file.type.startsWith("image/") ? (
@@ -299,6 +299,13 @@ export default function InputContainer({ sendMessage, isGenerating, selectedMode
                     )}
                   </div>
                 ))}
+
+                {isRagProcessing && (
+                  <div className="flex items-center gap-2 text-xs text-primary bg-primary/5 px-3 py-1.5 rounded-lg border border-primary/10 font-medium">
+                    <Loader2 size={14} className="animate-spin" />
+                    <span>Processing document...</span>
+                  </div>
+                )}
               </div>
             )}
 
@@ -311,7 +318,7 @@ export default function InputContainer({ sendMessage, isGenerating, selectedMode
               onKeyDown={handleKeyDown}
               onPaste={handlePaste}
               disabled={isBusy}
-              placeholder={isRagProcessing ? "Document is processing..." : "Ask Anything"}
+              placeholder="Ask Anything"
               rows={1}
               className="w-full custom-scroll custom-scroll-xs py-2 px-1.5 text-primary text-sm focus:outline-none resize-none overflow-y-auto max-h-28 transition-all duration-200 ease-in-out"
             />
@@ -334,7 +341,7 @@ export default function InputContainer({ sendMessage, isGenerating, selectedMode
                     }`}
                   title="Attach files"
                 >
-                  <Paperclip size={16} strokeWidth={1.8} />
+                  <Paperclip size={isMobile ? 20 : 16} strokeWidth={1.8} />
                 </label>
               </div>
 
@@ -343,7 +350,7 @@ export default function InputContainer({ sendMessage, isGenerating, selectedMode
 
                 <button
                   type="submit"
-                  disabled={(!hasContent && !isGenerating) || isGlobalUploading}
+                  disabled={(!hasContent && !isGenerating) || isGlobalUploading || isRagProcessing}
                   className={`
                   flex items-center justify-center h-8 w-8 rounded-full transition-all duration-200
                   ${isBusy || hasContent
@@ -352,7 +359,7 @@ export default function InputContainer({ sendMessage, isGenerating, selectedMode
                     }
                 `}
                 >
-                  {isGenerating || isGlobalUploading ? <Loader2 size={16} className="animate-spin" /> : <ArrowUp size={18} strokeWidth={2.5} />}
+                  {isGenerating || isGlobalUploading || isRagProcessing ? <Loader2 size={isMobile ? 20 : 16} className="animate-spin" /> : <ArrowUp size={isMobile ? 22 : 18} strokeWidth={2.5} />}
                 </button>
               </div>
             </div>

frontend/src/components/containers/MessagesContainer.tsx

@@ -50,7 +50,7 @@ export default function MessagesContainer({ messages, onResend, isGenerating, is
         const isCopied = copiedId === message.id;
 
         return (
-          <div className="px-4 mx-auto w-full max-w-svw sm:max-w-180 py-3" style={{ contain: "content" }}>
+          <div className="px-5 sm:px-4 mx-auto w-full max-w-svw sm:max-w-180 py-4 sm:py-3" style={{ contain: "content" }}>
             {message.role === "ASSISTANT" ? (
               <ModelMessage
                 key={message.id}

frontend/src/components/messages/ModelMessage.tsx

@@ -59,7 +59,7 @@ type Props = {
 export default memo(function ModelMessage({ text, status, isCopied, onCopy, onResend, hideToolTip, isGenerating }: Props) {
 
   return (
-    <div className="max-w-full text-base md:text-sm py-4 rounded-xl">
+    <div className="max-w-full text-[15px] sm:text-sm py-2 sm:py-4 rounded-xl">
       {isGenerating && text.length === 0 && !status ? (
         <span className="inline-block animate-pulse font-bold">▍</span>
       ) : (

frontend/src/components/messages/UserMessage.tsx

@@ -37,7 +37,7 @@ export default memo(function UserMessage({ message, isCopied, onCopy }: Props) {
   return (
     <div className="flex flex-col gap-2 justify-end items-end group">
       {hasFiles && (
-        <div className="mb-1 w-full max-w-xs sm:max-w-lg flex flex-wrap justify-end gap-2">
+        <div className="mb-1 w-full max-w-[85%] sm:max-w-lg flex flex-wrap justify-end gap-2">
           {messageFiles.map((fileItem, index) => {
             const { fileUrl, fileName, fileType } = fileItem;
             const isImage = fileType?.startsWith("image/") || (fileUrl ? isImageUrl(fileUrl) : false);
@@ -71,7 +71,7 @@ export default memo(function UserMessage({ message, isCopied, onCopy }: Props) {
       )}
 
       {text && (
-        <div className="user-message-color text-primary border border-gray-300/20 px-3 py-2.5 rounded-tr-none rounded-2xl max-w-xs sm:max-w-lg relative transition-all">
+        <div className="user-message-color text-primary border border-gray-300/20 px-4 py-3 sm:px-3 sm:py-2.5 rounded-tr-none rounded-2xl max-w-[85%] sm:max-w-lg relative transition-all">
           {isLongText && (
             <button
               onClick={() => setIsExpanded(!isExpanded)}
@@ -89,7 +89,7 @@ export default memo(function UserMessage({ message, isCopied, onCopy }: Props) {
             </button>
           )}
 
-          <p className="text-base md:text-sm leading-normal whitespace-pre-wrap break-words">{displayedText}</p>
+          <p className="text-[15px] sm:text-sm leading-relaxed sm:leading-normal whitespace-pre-wrap break-words">{displayedText}</p>
         </div>
       )}