mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-16 16:30:15 +01:00
➿ fix: createFileSearchTool to return tuples for error messages (#10547)
This commit is contained in:
Danny Avila
GitHub
parent
1b2f1ff09b
commit
49c57b27fd
2 changed files with 203 additions and 57 deletions
|
|
@ -78,11 +78,11 @@ const createFileSearchTool = async ({ userId, files, entity_id, fileCitations =
|
|||
return tool(
|
||||
async ({ query }) => {
|
||||
if (files.length === 0) {
|
||||
return 'No files to search. Instruct the user to add files for the search.';
|
||||
return ['No files to search. Instruct the user to add files for the search.', undefined];
|
||||
}
|
||||
const jwtToken = generateShortLivedToken(userId);
|
||||
if (!jwtToken) {
|
||||
return 'There was an error authenticating the file search request.';
|
||||
return ['There was an error authenticating the file search request.', undefined];
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -122,7 +122,7 @@ const createFileSearchTool = async ({ userId, files, entity_id, fileCitations =
|
|||
const validResults = results.filter((result) => result !== null);
|
||||
|
||||
if (validResults.length === 0) {
|
||||
return 'No results found or errors occurred while searching the files.';
|
||||
return ['No results found or errors occurred while searching the files.', undefined];
|
||||
}
|
||||
|
||||
const formattedResults = validResults
|
||||
|
|
|
|||
|
|
@ -1,17 +1,11 @@
|
|||
const { createFileSearchTool } = require('../../../../../app/clients/tools/util/fileSearch');
|
||||
const axios = require('axios');
|
||||
|
||||
// Mock dependencies
|
||||
jest.mock('../../../../../models', () => ({
|
||||
Files: {
|
||||
find: jest.fn(),
|
||||
},
|
||||
jest.mock('axios');
|
||||
jest.mock('@librechat/api', () => ({
|
||||
generateShortLivedToken: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('../../../../../server/services/Files/VectorDB/crud', () => ({
|
||||
queryVectors: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('../../../../../config', () => ({
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: {
|
||||
warn: jest.fn(),
|
||||
error: jest.fn(),
|
||||
|
|
@ -19,68 +13,220 @@ jest.mock('../../../../../config', () => ({
|
|||
},
|
||||
}));
|
||||
|
||||
const { queryVectors } = require('../../../../../server/services/Files/VectorDB/crud');
|
||||
jest.mock('~/models/File', () => ({
|
||||
getFiles: jest.fn().mockResolvedValue([]),
|
||||
}));
|
||||
|
||||
describe('fileSearch.js - test only new file_id and page additions', () => {
|
||||
jest.mock('~/server/services/Files/permissions', () => ({
|
||||
filterFilesByAgentAccess: jest.fn((options) => Promise.resolve(options.files)),
|
||||
}));
|
||||
|
||||
const { createFileSearchTool } = require('~/app/clients/tools/util/fileSearch');
|
||||
const { generateShortLivedToken } = require('@librechat/api');
|
||||
|
||||
describe('fileSearch.js - tuple return validation', () => {
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
process.env.RAG_API_URL = 'http://localhost:8000';
|
||||
});
|
||||
|
||||
// Test only the specific changes: file_id and page metadata additions
|
||||
it('should add file_id and page to search result format', async () => {
|
||||
const mockFiles = [{ file_id: 'test-file-123' }];
|
||||
const mockResults = [
|
||||
{
|
||||
describe('error cases should return tuple with undefined as second value', () => {
|
||||
it('should return tuple when no files provided', async () => {
|
||||
const fileSearchTool = await createFileSearchTool({
|
||||
userId: 'user1',
|
||||
files: [],
|
||||
});
|
||||
|
||||
const result = await fileSearchTool.func({ query: 'test query' });
|
||||
|
||||
expect(Array.isArray(result)).toBe(true);
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toBe('No files to search. Instruct the user to add files for the search.');
|
||||
expect(result[1]).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should return tuple when JWT token generation fails', async () => {
|
||||
generateShortLivedToken.mockReturnValue(null);
|
||||
|
||||
const fileSearchTool = await createFileSearchTool({
|
||||
userId: 'user1',
|
||||
files: [{ file_id: 'file-1', filename: 'test.pdf' }],
|
||||
});
|
||||
|
||||
const result = await fileSearchTool.func({ query: 'test query' });
|
||||
|
||||
expect(Array.isArray(result)).toBe(true);
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toBe('There was an error authenticating the file search request.');
|
||||
expect(result[1]).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should return tuple when no valid results found', async () => {
|
||||
generateShortLivedToken.mockReturnValue('mock-jwt-token');
|
||||
axios.post.mockRejectedValue(new Error('API Error'));
|
||||
|
||||
const fileSearchTool = await createFileSearchTool({
|
||||
userId: 'user1',
|
||||
files: [{ file_id: 'file-1', filename: 'test.pdf' }],
|
||||
});
|
||||
|
||||
const result = await fileSearchTool.func({ query: 'test query' });
|
||||
|
||||
expect(Array.isArray(result)).toBe(true);
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]).toBe('No results found or errors occurred while searching the files.');
|
||||
expect(result[1]).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('success cases should return tuple with artifact object', () => {
|
||||
it('should return tuple with formatted results and sources artifact', async () => {
|
||||
generateShortLivedToken.mockReturnValue('mock-jwt-token');
|
||||
|
||||
const mockApiResponse = {
|
||||
data: [
|
||||
[
|
||||
{
|
||||
page_content: 'test content',
|
||||
metadata: { source: 'test.pdf', page: 1 },
|
||||
page_content: 'This is test content from the document',
|
||||
metadata: { source: '/path/to/test.pdf', page: 1 },
|
||||
},
|
||||
0.3,
|
||||
0.2,
|
||||
],
|
||||
[
|
||||
{
|
||||
page_content: 'Additional relevant content',
|
||||
metadata: { source: '/path/to/test.pdf', page: 2 },
|
||||
},
|
||||
0.35,
|
||||
],
|
||||
],
|
||||
},
|
||||
];
|
||||
};
|
||||
|
||||
queryVectors.mockResolvedValue(mockResults);
|
||||
axios.post.mockResolvedValue(mockApiResponse);
|
||||
|
||||
const fileSearchTool = await createFileSearchTool({
|
||||
userId: 'user1',
|
||||
files: mockFiles,
|
||||
entity_id: 'agent-123',
|
||||
const fileSearchTool = await createFileSearchTool({
|
||||
userId: 'user1',
|
||||
files: [{ file_id: 'file-123', filename: 'test.pdf' }],
|
||||
entity_id: 'agent-456',
|
||||
});
|
||||
|
||||
const result = await fileSearchTool.func({ query: 'test query' });
|
||||
|
||||
expect(Array.isArray(result)).toBe(true);
|
||||
expect(result).toHaveLength(2);
|
||||
|
||||
const [formattedString, artifact] = result;
|
||||
|
||||
expect(typeof formattedString).toBe('string');
|
||||
expect(formattedString).toContain('File: test.pdf');
|
||||
expect(formattedString).toContain('Relevance:');
|
||||
expect(formattedString).toContain('This is test content from the document');
|
||||
expect(formattedString).toContain('Additional relevant content');
|
||||
|
||||
expect(artifact).toBeDefined();
|
||||
expect(artifact).toHaveProperty('file_search');
|
||||
expect(artifact.file_search).toHaveProperty('sources');
|
||||
expect(artifact.file_search).toHaveProperty('fileCitations', false);
|
||||
expect(Array.isArray(artifact.file_search.sources)).toBe(true);
|
||||
expect(artifact.file_search.sources.length).toBe(2);
|
||||
|
||||
const source = artifact.file_search.sources[0];
|
||||
expect(source).toMatchObject({
|
||||
type: 'file',
|
||||
fileId: 'file-123',
|
||||
fileName: 'test.pdf',
|
||||
content: expect.any(String),
|
||||
relevance: expect.any(Number),
|
||||
pages: [1],
|
||||
pageRelevance: { 1: expect.any(Number) },
|
||||
});
|
||||
});
|
||||
|
||||
// Mock the tool's function to return the formatted result
|
||||
fileSearchTool.func = jest.fn().mockImplementation(async () => {
|
||||
// Simulate the new format with file_id and page
|
||||
const formattedResults = [
|
||||
{
|
||||
filename: 'test.pdf',
|
||||
content: 'test content',
|
||||
distance: 0.3,
|
||||
file_id: 'test-file-123', // NEW: added file_id
|
||||
page: 1, // NEW: added page
|
||||
},
|
||||
];
|
||||
it('should include file citations in description when enabled', async () => {
|
||||
generateShortLivedToken.mockReturnValue('mock-jwt-token');
|
||||
|
||||
// NEW: Internal data section for processAgentResponse
|
||||
const internalData = formattedResults
|
||||
.map(
|
||||
(result) =>
|
||||
`File: ${result.filename}\nFile_ID: ${result.file_id}\nRelevance: ${(1.0 - result.distance).toFixed(4)}\nPage: ${result.page || 'N/A'}\nContent: ${result.content}\n`,
|
||||
)
|
||||
.join('\n---\n');
|
||||
const mockApiResponse = {
|
||||
data: [
|
||||
[
|
||||
{
|
||||
page_content: 'Content with citations',
|
||||
metadata: { source: '/path/to/doc.pdf', page: 3 },
|
||||
},
|
||||
0.15,
|
||||
],
|
||||
],
|
||||
};
|
||||
|
||||
return `File: test.pdf\nRelevance: 0.7000\nContent: test content\n\n<!-- INTERNAL_DATA_START -->\n${internalData}\n<!-- INTERNAL_DATA_END -->`;
|
||||
axios.post.mockResolvedValue(mockApiResponse);
|
||||
|
||||
const fileSearchTool = await createFileSearchTool({
|
||||
userId: 'user1',
|
||||
files: [{ file_id: 'file-789', filename: 'doc.pdf' }],
|
||||
fileCitations: true,
|
||||
});
|
||||
|
||||
const result = await fileSearchTool.func({ query: 'test query' });
|
||||
|
||||
expect(Array.isArray(result)).toBe(true);
|
||||
expect(result).toHaveLength(2);
|
||||
|
||||
const [formattedString, artifact] = result;
|
||||
|
||||
expect(formattedString).toContain('Anchor:');
|
||||
expect(formattedString).toContain('\\ue202turn0file0');
|
||||
expect(artifact.file_search.fileCitations).toBe(true);
|
||||
});
|
||||
|
||||
const result = await fileSearchTool.func('test');
|
||||
it('should handle multiple files correctly', async () => {
|
||||
generateShortLivedToken.mockReturnValue('mock-jwt-token');
|
||||
|
||||
// Verify the new additions
|
||||
expect(result).toContain('File_ID: test-file-123');
|
||||
expect(result).toContain('Page: 1');
|
||||
expect(result).toContain('<!-- INTERNAL_DATA_START -->');
|
||||
expect(result).toContain('<!-- INTERNAL_DATA_END -->');
|
||||
const mockResponse1 = {
|
||||
data: [
|
||||
[
|
||||
{
|
||||
page_content: 'Content from file 1',
|
||||
metadata: { source: '/path/to/file1.pdf', page: 1 },
|
||||
},
|
||||
0.25,
|
||||
],
|
||||
],
|
||||
};
|
||||
|
||||
const mockResponse2 = {
|
||||
data: [
|
||||
[
|
||||
{
|
||||
page_content: 'Content from file 2',
|
||||
metadata: { source: '/path/to/file2.pdf', page: 1 },
|
||||
},
|
||||
0.15,
|
||||
],
|
||||
],
|
||||
};
|
||||
|
||||
axios.post.mockResolvedValueOnce(mockResponse1).mockResolvedValueOnce(mockResponse2);
|
||||
|
||||
const fileSearchTool = await createFileSearchTool({
|
||||
userId: 'user1',
|
||||
files: [
|
||||
{ file_id: 'file-1', filename: 'file1.pdf' },
|
||||
{ file_id: 'file-2', filename: 'file2.pdf' },
|
||||
],
|
||||
});
|
||||
|
||||
const result = await fileSearchTool.func({ query: 'test query' });
|
||||
|
||||
expect(Array.isArray(result)).toBe(true);
|
||||
expect(result).toHaveLength(2);
|
||||
|
||||
const [formattedString, artifact] = result;
|
||||
|
||||
expect(formattedString).toContain('file1.pdf');
|
||||
expect(formattedString).toContain('file2.pdf');
|
||||
expect(artifact.file_search.sources).toHaveLength(2);
|
||||
// Results are sorted by distance (ascending), so file-2 (0.15) comes before file-1 (0.25)
|
||||
expect(artifact.file_search.sources[0].fileId).toBe('file-2');
|
||||
expect(artifact.file_search.sources[1].fileId).toBe('file-1');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue