Merge branch 'agent-paper-main' of https://github.com/brown-dash/Dash-Web into joanne-tutorialagent

1 files changed, 10 insertions, 3 deletions

diff --git a/src/client/views/nodes/chatbot/tools/RAGTool.ts b/src/client/views/nodes/chatbot/tools/RAGTool.ts
index ef374ed22..af44de520 100644
--- a/src/client/views/nodes/chatbot/tools/RAGTool.ts
+++ b/src/client/views/nodes/chatbot/tools/RAGTool.ts
@@ -3,6 +3,7 @@ import { Observation, RAGChunk } from '../types/types';
 import { ParametersType, ToolInfo } from '../types/tool_types';
 import { Vectorstore } from '../vectorstore/Vectorstore';
 import { BaseTool } from './BaseTool';
+import { DocumentMetadataTool } from './DocumentMetadataTool';
 
 const ragToolParams = [
     {
@@ -11,13 +12,19 @@ const ragToolParams = [
         description: "A detailed prompt representing an ideal chunk to embed and compare against document vectors to retrieve the most relevant content for answering the user's query.",
         required: true,
     },
+    {
+        name: 'doc_ids',
+        type: 'string[]',
+        description: 'An optional array of document IDs to retrieve chunks from. If you want to retrieve chunks from all documents, leave this as an empty array: [] (DO NOT LEAVE THIS EMPTY).',
+        required: false,
+    },
 ] as const;
 
 type RAGToolParamsType = typeof ragToolParams;
 
 const ragToolInfo: ToolInfo<RAGToolParamsType> = {
     name: 'rag',
-    description: 'Performs a RAG (Retrieval-Augmented Generation) search on user documents and returns a set of document chunks (text or images) to provide a grounded response based on user documents.',
+    description: `Performs a RAG (Retrieval-Augmented Generation) search on user documents (only PDF, audio, and video are supported—for information about other document types, use the ${DocumentMetadataTool.name} tool) and returns a set of document chunks (text or images) to provide a grounded response based on user documents.`,
     citationRules: `When using the RAG tool, the structure must adhere to the format described in the ReAct prompt. Below are additional guidelines specifically for RAG-based responses:
 
             1. **Grounded Text Guidelines**:
@@ -68,14 +75,14 @@ export class RAGTool extends BaseTool<RAGToolParamsType> {
     }
 
     async execute(args: ParametersType<RAGToolParamsType>): Promise<Observation[]> {
-        const relevantChunks = await this.vectorstore.retrieve(args.hypothetical_document_chunk);
+        const relevantChunks = await this.vectorstore.retrieve(args.hypothetical_document_chunk, undefined, args.doc_ids ?? undefined);
         const formattedChunks = await this.getFormattedChunks(relevantChunks);
         return formattedChunks;
     }
 
     async getFormattedChunks(relevantChunks: RAGChunk[]): Promise<Observation[]> {
         try {
-            const { formattedChunks } = await Networking.PostToServer('/formatChunks', { relevantChunks }) as { formattedChunks: Observation[]}
+            const { formattedChunks } = (await Networking.PostToServer('/formatChunks', { relevantChunks })) as { formattedChunks: Observation[] };
 
             if (!formattedChunks) {
                 throw new Error('Failed to format chunks');