🔍 feat: Add Entity ID Support for File Search Shared Resources (danny…

…-avila#5028)
openwengo · Dec 20, 2024 · bca9a18 · bca9a18
1 parent 80814d1
commit bca9a18
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 22 deletions.
diff --git a/api/app/clients/tools/util/fileSearch.js b/api/app/clients/tools/util/fileSearch.js
@@ -50,9 +50,10 @@ const primeFiles = async (options) => {
  * @param {Object} options
  * @param {ServerRequest} options.req
  * @param {Array<{ file_id: string; filename: string }>} options.files
+ * @param {string} [options.entity_id]
  * @returns
  */
-const createFileSearchTool = async ({ req, files }) => {
+const createFileSearchTool = async ({ req, files, entity_id }) => {
   return tool(
     async ({ query }) => {
       if (files.length === 0) {
@@ -62,27 +63,36 @@ const createFileSearchTool = async ({ req, files }) => {
       if (!jwtToken) {
         return 'There was an error authenticating the file search request.';
       }
+
+      /**
+       *
+       * @param {import('librechat-data-provider').TFile} file
+       * @returns {{ file_id: string, query: string, k: number, entity_id?: string }}
+       */
+      const createQueryBody = (file) => {
+        const body = {
+          file_id: file.file_id,
+          query,
+          k: 5,
+        };
+        if (!entity_id) {
+          return body;
+        }
+        body.entity_id = entity_id;
+        logger.debug(`[${Tools.file_search}] RAG API /query body`, body);
+        return body;
+      };
+
       const queryPromises = files.map((file) =>
         axios
-          .post(
-            `${process.env.RAG_API_URL}/query`,
-            {
-              file_id: file.file_id,
-              query,
-              k: 5,
-            },
-            {
-              headers: {
-                Authorization: `Bearer ${jwtToken}`,
-                'Content-Type': 'application/json',
-              },
+          .post(`${process.env.RAG_API_URL}/query`, createQueryBody(file), {
+            headers: {
+              Authorization: `Bearer ${jwtToken}`,
+              'Content-Type': 'application/json',
             },
-          )
+          })
           .catch((error) => {
-            logger.error(
-              `Error encountered in \`file_search\` while querying file_id ${file._id}:`,
-              error,
-            );
+            logger.error('Error encountered in `file_search` while querying file:', error);
             return null;
           }),
       );

diff --git a/api/app/clients/tools/util/handleTools.js b/api/app/clients/tools/util/handleTools.js
@@ -256,7 +256,7 @@ const loadTools = async ({
         if (toolContext) {
           toolContextMap[tool] = toolContext;
         }
-        return createFileSearchTool({ req: options.req, files });
+        return createFileSearchTool({ req: options.req, files, entity_id: agent?.id });
       };
       continue;
     } else if (mcpToolPattern.test(tool)) {

diff --git a/api/server/services/Files/VectorDB/crud.js b/api/server/services/Files/VectorDB/crud.js
@@ -50,13 +50,14 @@ const deleteVectors = async (req, file) => {
  * @param {Express.Multer.File} params.file - The file object, which is part of the request. The file object should
  *                                     have a `path` property that points to the location of the uploaded file.
  * @param {string} params.file_id - The file ID.
+ * @param {string} [params.entity_id] - The entity ID for shared resources.
  *
  * @returns {Promise<{ filepath: string, bytes: number }>}
  *          A promise that resolves to an object containing:
  *            - filepath: The path where the file is saved.
  *            - bytes: The size of the file in bytes.
  */
-async function uploadVectors({ req, file, file_id }) {
+async function uploadVectors({ req, file, file_id, entity_id }) {
   if (!process.env.RAG_API_URL) {
     throw new Error('RAG_API_URL not defined');
   }
@@ -66,8 +67,11 @@ async function uploadVectors({ req, file, file_id }) {
     const formData = new FormData();
     formData.append('file_id', file_id);
     formData.append('file', fs.createReadStream(file.path));
+    if (entity_id != null && entity_id) {
+      formData.append('entity_id', entity_id);
+    }
 
-    const formHeaders = formData.getHeaders(); // Automatically sets the correct Content-Type
+    const formHeaders = formData.getHeaders();
 
     const response = await axios.post(`${process.env.RAG_API_URL}/embed`, formData, {
       headers: {

diff --git a/api/server/services/Files/process.js b/api/server/services/Files/process.js
@@ -479,6 +479,7 @@ const processAgentFileUpload = async ({ req, res, metadata }) => {
   }
 
   let fileInfoMetadata;
+  const entity_id = messageAttachment === true ? undefined : agent_id;
   if (tool_resource === EToolResources.execute_code) {
     const { handleFileUpload: uploadCodeEnvFile } = getStrategyFunctions(FileSources.execute_code);
     const result = await loadAuthValues({ userId: req.user.id, authFields: [EnvVar.CODE_API_KEY] });
@@ -488,7 +489,7 @@ const processAgentFileUpload = async ({ req, res, metadata }) => {
       stream,
       filename: file.originalname,
       apiKey: result[EnvVar.CODE_API_KEY],
-      entity_id: messageAttachment === true ? undefined : agent_id,
+      entity_id,
     });
     fileInfoMetadata = { fileIdentifier };
   }
@@ -512,6 +513,7 @@ const processAgentFileUpload = async ({ req, res, metadata }) => {
     req,
     file,
     file_id,
+    entity_id,
   });
 
   let filepath = _filepath;