Skip to content

Commit

Permalink
📂 feat: RAG Improvements (#2169)
Browse files Browse the repository at this point in the history
* feat: new vector file processing strategy

* chore: remove unused client files

* chore: remove more unused client files

* chore: remove more unused client files and move used to new dir

* chore(DataIcon): add className

* WIP: Model Endpoint Settings Update, draft additional context settings

* feat: improve parsing for augmented prompt, add full context option

* chore: remove volume mounting from rag.yml as no longer necessary
  • Loading branch information
danny-avila committed Mar 22, 2024
1 parent f427ad7 commit 45a95ac
Show file tree
Hide file tree
Showing 40 changed files with 716 additions and 2,047 deletions.
149 changes: 94 additions & 55 deletions api/app/clients/prompts/createContextHandlers.js
Original file line number Diff line number Diff line change
@@ -1,4 +1,15 @@
const axios = require('axios');
const { isEnabled } = require('~/server/utils');

const footer = `Use the context as your learned knowledge to better answer the user.
In your response, remember to follow these guidelines:
- If you don't know the answer, simply say that you don't know.
- If you are unsure how to answer, ask for clarification.
- Avoid mentioning that you obtained the information from the context.
Answer appropriately in the user's language.
`;

function createContextHandlers(req, userMessageContent) {
if (!process.env.RAG_API_URL) {
Expand All @@ -9,25 +20,37 @@ function createContextHandlers(req, userMessageContent) {
const processedFiles = [];
const processedIds = new Set();
const jwtToken = req.headers.authorization.split(' ')[1];
const useFullContext = isEnabled(process.env.RAG_USE_FULL_CONTEXT);

const query = async (file) => {
if (useFullContext) {
return axios.get(`${process.env.RAG_API_URL}/documents/${file.file_id}/context`, {
headers: {
Authorization: `Bearer ${jwtToken}`,
},
});
}

return axios.post(
`${process.env.RAG_API_URL}/query`,
{
file_id: file.file_id,
query: userMessageContent,
k: 4,
},
{
headers: {
Authorization: `Bearer ${jwtToken}`,
'Content-Type': 'application/json',
},
},
);
};

const processFile = async (file) => {
if (file.embedded && !processedIds.has(file.file_id)) {
try {
const promise = axios.post(
`${process.env.RAG_API_URL}/query`,
{
file_id: file.file_id,
query: userMessageContent,
k: 4,
},
{
headers: {
Authorization: `Bearer ${jwtToken}`,
'Content-Type': 'application/json',
},
},
);

const promise = query(file);
queryPromises.push(promise);
processedFiles.push(file);
processedIds.add(file.file_id);
Expand All @@ -43,67 +66,83 @@ function createContextHandlers(req, userMessageContent) {
return '';
}

const oneFile = processedFiles.length === 1;
const header = `The user has attached ${oneFile ? 'a' : processedFiles.length} file${
!oneFile ? 's' : ''
} to the conversation:`;

const files = `${
oneFile
? ''
: `
<files>`
}${processedFiles
.map(
(file) => `
<file>
<filename>${file.filename}</filename>
<type>${file.type}</type>
</file>`,
)
.join('')}${
oneFile
? ''
: `
</files>`
}`;

const resolvedQueries = await Promise.all(queryPromises);

const context = resolvedQueries
.map((queryResult, index) => {
const file = processedFiles[index];
const contextItems = queryResult.data
let contextItems = queryResult.data;

const generateContext = (currentContext) =>
`
<file>
<filename>${file.filename}</filename>
<context>${currentContext}
</context>
</file>`;

if (useFullContext) {
return generateContext(`\n${contextItems}`);
}

contextItems = queryResult.data
.map((item) => {
const pageContent = item[0].page_content;
return `
<contextItem>
<![CDATA[${pageContent}]]>
</contextItem>
`;
<![CDATA[${pageContent?.trim()}]]>
</contextItem>`;
})
.join('');

return `
<file>
<filename>${file.filename}</filename>
<context>
${contextItems}
</context>
</file>
`;
return generateContext(contextItems);
})
.join('');

const template = `The user has attached ${
processedFiles.length === 1 ? 'a' : processedFiles.length
} file${processedFiles.length !== 1 ? 's' : ''} to the conversation:
<files>
${processedFiles
.map(
(file) => `
<file>
<filename>${file.filename}</filename>
<type>${file.type}</type>
</file>
`,
)
.join('')}
</files>
if (useFullContext) {
const prompt = `${header}
${context}
${footer}`;

return prompt;
}

const prompt = `${header}
${files}
A semantic search was executed with the user's message as the query, retrieving the following context inside <context></context> XML tags.
<context>
${context}
<context>${context}
</context>
Use the context as your learned knowledge to better answer the user.
In your response, remember to follow these guidelines:
- If you don't know the answer, simply say that you don't know.
- If you are unsure how to answer, ask for clarification.
- Avoid mentioning that you obtained the information from the context.
Answer appropriately in the user's language.
`;
${footer}`;

return template;
return prompt;
} catch (error) {
console.error('Error creating context:', error);
throw error; // Re-throw the error to propagate it to the caller
Expand Down
96 changes: 96 additions & 0 deletions api/server/services/Files/VectorDB/crud.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
const fs = require('fs');
const axios = require('axios');
const FormData = require('form-data');
const { FileSources } = require('librechat-data-provider');
const { logger } = require('~/config');

/**
* Deletes a file from the vector database. This function takes a file object, constructs the full path, and
* verifies the path's validity before deleting the file. If the path is invalid, an error is thrown.
*
* @param {Express.Request} req - The request object from Express. It should have an `app.locals.paths` object with
* a `publicPath` property.
* @param {MongoFile} file - The file object to be deleted. It should have a `filepath` property that is
* a string representing the path of the file relative to the publicPath.
*
* @returns {Promise<void>}
* A promise that resolves when the file has been successfully deleted, or throws an error if the
* file path is invalid or if there is an error in deletion.
*/
const deleteVectors = async (req, file) => {
if (file.embedded && process.env.RAG_API_URL) {
const jwtToken = req.headers.authorization.split(' ')[1];
axios.delete(`${process.env.RAG_API_URL}/documents`, {
headers: {
Authorization: `Bearer ${jwtToken}`,
'Content-Type': 'application/json',
accept: 'application/json',
},
data: [file.file_id],
});
}
};

/**
* Uploads a file to the configured Vector database
*
* @param {Object} params - The params object.
* @param {Object} params.req - The request object from Express. It should have a `user` property with an `id`
* representing the user, and an `app.locals.paths` object with an `uploads` path.
* @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.
*
* @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 }) {
if (!process.env.RAG_API_URL) {
throw new Error('RAG_API_URL not defined');
}

try {
const jwtToken = req.headers.authorization.split(' ')[1];
const formData = new FormData();
formData.append('file_id', file_id);
formData.append('file', fs.createReadStream(file.path));

const formHeaders = formData.getHeaders(); // Automatically sets the correct Content-Type

const response = await axios.post(`${process.env.RAG_API_URL}/embed`, formData, {
headers: {
Authorization: `Bearer ${jwtToken}`,
accept: 'application/json',
...formHeaders,
},
});

const responseData = response.data;
logger.debug('Response from embedding file', responseData);

if (responseData.known_type === false) {
throw new Error(`File embedding failed. The filetype ${file.mimetype} is not supported`);
}

if (!responseData.status) {
throw new Error('File embedding failed.');
}

return {
bytes: file.size,
filename: file.originalname,
filepath: FileSources.vectordb,
embedded: Boolean(responseData.known_type),
};
} catch (error) {
logger.error('Error embedding file', error);
throw new Error(error.message || 'An error occurred during file upload.');
}
}

module.exports = {
deleteVectors,
uploadVectors,
};
5 changes: 5 additions & 0 deletions api/server/services/Files/VectorDB/index.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
const crud = require('./crud');

module.exports = {
...crud,
};
43 changes: 7 additions & 36 deletions api/server/services/Files/process.js
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
const path = require('path');
const { v4 } = require('uuid');
const axios = require('axios');
const mime = require('mime/lite');
const {
isUUID,
Expand Down Expand Up @@ -265,50 +264,22 @@ const uploadImageBuffer = async ({ req, context }) => {
*/
const processFileUpload = async ({ req, res, file, metadata }) => {
const isAssistantUpload = metadata.endpoint === EModelEndpoint.assistants;
const source = isAssistantUpload ? FileSources.openai : req.app.locals.fileStrategy;
const source = isAssistantUpload ? FileSources.openai : FileSources.vectordb;
const { handleFileUpload } = getStrategyFunctions(source);
const { file_id, temp_file_id } = metadata;

let embedded = false;
if (process.env.RAG_API_URL) {
try {
const jwtToken = req.headers.authorization.split(' ')[1];
const filepath = `./uploads/temp/${file.path.split('uploads/temp/')[1]}`;
const response = await axios.post(
`${process.env.RAG_API_URL}/embed`,
{
filename: file.originalname,
file_content_type: file.mimetype,
filepath,
file_id,
},
{
headers: {
Authorization: `Bearer ${jwtToken}`,
'Content-Type': 'application/json',
},
},
);

if (response.status === 200) {
embedded = true;
}
} catch (error) {
logger.error('Error embedding file', error);
throw new Error(error);
}
} else if (!isAssistantUpload) {
logger.error('RAG_API_URL not set, cannot support process file upload');
throw new Error('RAG_API_URL not set, cannot support process file upload');
}

/** @type {OpenAI | undefined} */
let openai;
if (source === FileSources.openai) {
({ openai } = await initializeClient({ req }));
}

const { id, bytes, filename, filepath } = await handleFileUpload({ req, file, file_id, openai });
const { id, bytes, filename, filepath, embedded } = await handleFileUpload({
req,
file,
file_id,
openai,
});

if (isAssistantUpload && !metadata.message_file) {
await openai.beta.assistants.files.create(metadata.assistant_id, {
Expand Down
Loading

0 comments on commit 45a95ac

Please sign in to comment.