🔍 feat: Add Entity ID Support for File Search Shared Resources (#5028)

This commit is contained in:
Danny Avila 2024-12-17 22:11:18 -05:00 committed by GitHub
parent 18ad89be2c
commit d3cafeee96
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 38 additions and 22 deletions

View file

@ -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,
.post(`${process.env.RAG_API_URL}/query`, createQueryBody(file), {
headers: {
Authorization: `Bearer ${jwtToken}`,
'Content-Type': 'application/json',
},
{
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;
}),
);