From a88946b6359e295b2d91dfdf95b6027d0c0cfb83 Mon Sep 17 00:00:00 2001 From: Adnan Hussain Date: Tue, 17 Dec 2024 12:43:10 +0500 Subject: [PATCH 1/3] feat: enable file search option for azure v2 assistant --- api/server/controllers/assistants/v2.js | 86 ++++++++++++++++++ api/server/services/Files/VectorStore/crud.js | 89 +++++++++++++++++++ .../services/Files/VectorStore/index.js | 5 ++ api/server/services/Files/process.js | 35 ++++++-- api/server/services/Files/strategies.js | 9 ++ client/src/common/assistants-types.ts | 2 + .../SidePanel/Builder/AssistantSelect.tsx | 14 ++- .../SidePanel/Builder/CapabilitiesForm.tsx | 23 ++++- .../SidePanel/Builder/CodeFiles.tsx | 15 +++- packages/data-provider/src/config.ts | 5 +- packages/data-provider/src/types/files.ts | 1 + 11 files changed, 266 insertions(+), 18 deletions(-) create mode 100644 api/server/services/Files/VectorStore/crud.js create mode 100644 api/server/services/Files/VectorStore/index.js diff --git a/api/server/controllers/assistants/v2.js b/api/server/controllers/assistants/v2.js index 54f9a6fbc6b..b7916f16309 100644 --- a/api/server/controllers/assistants/v2.js +++ b/api/server/controllers/assistants/v2.js @@ -231,6 +231,90 @@ const deleteResourceFileId = async ({ req, openai, assistant_id, tool_resource, }); }; +/** + * Modifies an assistant with the resource vector store ID. + * @param {object} params + * @param {Express.Request} params.req + * @param {OpenAIClient} params.openai + * @param {string} params.assistant_id + * @param {string} params.tool_resource + * @param {string} params.vector_store_id + * @returns {Promise} The updated assistant. + */ +const addResourceVectorId = async ({ + req, + openai, + assistant_id, + tool_resource, + vector_store_id, +}) => { + const assistant = await openai.beta.assistants.retrieve(assistant_id); + const { tool_resources = {} } = assistant; + + if (tool_resources[tool_resource]) { + // Replace the vector_store_id if already exists + tool_resources[tool_resource].vector_store_ids.push(vector_store_id); + } else { + // Initialize with the new vector_store_id + tool_resources[tool_resource] = { vector_store_ids: [vector_store_id] }; + } + + delete assistant.id; + return await updateAssistant({ + req, + openai, + assistant_id, + updateData: { tools: assistant.tools, tool_resources }, + }); +}; + +/** + * Deletes a vector store ID from an assistant's resource. + * @param {object} params + * @param {Express.Request} params.req + * @param {OpenAIClient} params.openai + * @param {string} params.assistant_id + * @param {string} [params.tool_resource] + * @param {string} params.vector_store_id + * @param {AssistantUpdateParams} params.updateData + * @returns {Promise} The updated assistant. + */ +const deleteResourceVectorId = async ({ + req, + openai, + assistant_id, + tool_resource, + vector_store_id, +}) => { + const assistant = await openai.beta.assistants.retrieve(assistant_id); + const { tool_resources = {} } = assistant; + + if (tool_resource && tool_resources[tool_resource]) { + const resource = tool_resources[tool_resource]; + const index = resource.vector_store_ids?.indexOf(vector_store_id); + if (index !== -1) { + resource.vector_store_ids.splice(index, 1); + } + } else { + for (const resourceKey in tool_resources) { + const resource = tool_resources[resourceKey]; + const index = resource.vector_store_ids?.indexOf(vector_store_id); + if (index !== -1) { + resource.vector_store_ids.splice(index, 1); + break; + } + } + } + + delete assistant.id; + return await updateAssistant({ + req, + openai, + assistant_id, + updateData: { tools: assistant.tools, tool_resources }, + }); +}; + /** * Modifies an assistant. * @route PATCH /assistants/:id @@ -260,4 +344,6 @@ module.exports = { updateAssistant, addResourceFileId, deleteResourceFileId, + addResourceVectorId, + deleteResourceVectorId, }; diff --git a/api/server/services/Files/VectorStore/crud.js b/api/server/services/Files/VectorStore/crud.js new file mode 100644 index 00000000000..c20e07c2676 --- /dev/null +++ b/api/server/services/Files/VectorStore/crud.js @@ -0,0 +1,89 @@ +const fs = require('fs'); +const { FilePurpose } = require('librechat-data-provider'); +const axios = require('axios'); +const { getOpenAIClient } = require('../../../controllers/assistants/helpers'); +const { logger } = require('~/config'); + +/** + * + * @param {OpenAIClient} openai - The initialized OpenAI client. + * @returns + */ +async function createVectorStore(openai) { + try { + const response = await openai.beta.vectorStores.create({ + name: 'Financial Statements', + }); + return response.id; + } catch (error) { + logger.error('[createVectorStore] Error creating vector store:', error.message); + throw error; + } +} + +/** + * Uploads a file to Azure OpenAI Vector Store for file search. + * + * @param {Object} params - The parameters for the upload. + * @param {Express.Multer.File} params.file - The file uploaded to the server via multer. + * @param {OpenAIClient} params.openai - The initialized OpenAI client. + * @param {string} [params.vectorStoreId] - The ID of the vector store. + * @returns {Promise} The response from Azure containing the file details. + */ +async function uploadToVectorStore({ openai, file, vectorStoreId }) { + try { + const filePath = file.path; + const fileStreams = [fs.createReadStream(filePath)]; + const response = await openai.beta.vectorStores.fileBatches.uploadAndPoll(vectorStoreId, { + files: fileStreams, + }); + logger.debug( + `[uploadToVectorStore] Successfully uploaded file to Azure Vector Store: ${response.id}`, + ); + return { + id: response.vector_store_id, + }; + } catch (error) { + logger.error('[uploadToVectorStore] Error uploading file:', error.message); + throw new Error(`Failed to upload file to Vector Store: ${error.message}`); + } +} + +/** + * Deletes a file from Azure OpenAI Vector Store. + * + * @param {string} file_id - The ID of the file to delete. + * @param {string} vectorStoreId - The ID of the vector store. + * @returns {Promise} + */ +async function deleteFromVectorStore(file_id, vectorStoreId) { + try { + // Get OpenAI client directly + const { openai } = await getOpenAIClient(); + const azureOpenAIEndpoint = openai.baseURL; + const azureOpenAIKey = openai.apiKey; + + const response = await axios.delete( + `${azureOpenAIEndpoint}/vector_stores/${vectorStoreId}/files/${file_id}?api-version=2024-10-01-preview`, + { + headers: { + 'api-key': azureOpenAIKey, + 'Content-Type': 'application/json', + }, + }, + ); + + if (!response.data.deleted) { + throw new Error(`Failed to delete file ${file_id} from Azure Vector Store`); + } + + logger.debug( + `[deleteFromVectorStore] Successfully deleted file ${file_id} from Azure Vector Store`, + ); + } catch (error) { + logger.error('[deleteFromVectorStore] Error deleting file:', error.message); + throw error; + } +} + +module.exports = { uploadToVectorStore, deleteFromVectorStore, createVectorStore }; diff --git a/api/server/services/Files/VectorStore/index.js b/api/server/services/Files/VectorStore/index.js new file mode 100644 index 00000000000..647813ec240 --- /dev/null +++ b/api/server/services/Files/VectorStore/index.js @@ -0,0 +1,5 @@ +const crud = require('./crud'); + +module.exports = { + ...crud, +}; diff --git a/api/server/services/Files/process.js b/api/server/services/Files/process.js index ab401420f1f..56ee7365ad7 100644 --- a/api/server/services/Files/process.js +++ b/api/server/services/Files/process.js @@ -18,7 +18,12 @@ const { isAssistantsEndpoint, } = require('librechat-data-provider'); const { EnvVar } = require('@librechat/agents'); -const { addResourceFileId, deleteResourceFileId } = require('~/server/controllers/assistants/v2'); +const { + addResourceFileId, + deleteResourceFileId, + addResourceVectorId, + deleteResourceVectorId, +} = require('~/server/controllers/assistants/v2'); const { convertImage, resizeAndConvert } = require('~/server/services/Files/images'); const { addAgentResourceFile, removeAgentResourceFiles } = require('~/models/Agent'); const { getOpenAIClient } = require('~/server/controllers/assistants/helpers'); @@ -375,14 +380,21 @@ const processFileUpload = async ({ req, res, metadata }) => { const isAssistantUpload = isAssistantsEndpoint(metadata.endpoint); const assistantSource = metadata.endpoint === EModelEndpoint.azureAssistants ? FileSources.azure : FileSources.openai; - const source = isAssistantUpload ? assistantSource : FileSources.vectordb; - const { handleFileUpload } = getStrategyFunctions(source); + const fileSource = + isAssistantUpload && metadata.tool_resource === EToolResources.file_search + ? FileSources.vector_store + : assistantSource; + const source = isAssistantUpload ? fileSource : FileSources.vectordb; + const { handleFileUpload, createStore } = getStrategyFunctions(source); const { file_id, temp_file_id } = metadata; - /** @type {OpenAI | undefined} */ - let openai; - if (checkOpenAIStorage(source)) { - ({ openai } = await getOpenAIClient({ req })); + /** @type {{ openai: OpenAIClient }} */ + let { openai } = await getOpenAIClient({ req, res }); + + let vector_id; + + if (source === FileSources.vector_store && typeof createStore === 'function') { + vector_id = await createStore(openai); } const { file } = req; @@ -399,12 +411,21 @@ const processFileUpload = async ({ req, res, metadata }) => { file, file_id, openai, + ...(source === FileSources.vector_store && { vectorStoreId: vector_id }), }); if (isAssistantUpload && !metadata.message_file && !metadata.tool_resource) { await openai.beta.assistants.files.create(metadata.assistant_id, { file_id: id, }); + } else if (isAssistantUpload && metadata.tool_resource === EToolResources.file_search) { + await addResourceVectorId({ + req, + openai, + vector_store_id: id, + assistant_id: metadata.assistant_id, + tool_resource: metadata.tool_resource, + }); } else if (isAssistantUpload && !metadata.message_file) { await addResourceFileId({ req, diff --git a/api/server/services/Files/strategies.js b/api/server/services/Files/strategies.js index ddfdd574690..5eb28847d89 100644 --- a/api/server/services/Files/strategies.js +++ b/api/server/services/Files/strategies.js @@ -24,6 +24,7 @@ const { const { uploadOpenAIFile, deleteOpenAIFile, getOpenAIFileStream } = require('./OpenAI'); const { getCodeOutputDownloadStream, uploadCodeEnvFile } = require('./Code'); const { uploadVectors, deleteVectors } = require('./VectorDB'); +const { uploadToVectorStore, deleteFromVectorStore, createVectorStore } = require('./VectorStore'); /** * Firebase Storage Strategy Functions @@ -103,6 +104,12 @@ const openAIStrategy = () => ({ getDownloadStream: getOpenAIFileStream, }); +const vectorStoreStrategy = () => ({ + createStore: createVectorStore, + handleFileUpload: uploadToVectorStore, + deleteFile: deleteFromVectorStore, +}); + /** * Code Output Strategy Functions * @@ -139,6 +146,8 @@ const getStrategyFunctions = (fileSource) => { return openAIStrategy(); } else if (fileSource === FileSources.vectordb) { return vectorStrategy(); + } else if (fileSource === FileSources.vector_store) { + return vectorStoreStrategy(); } else if (fileSource === FileSources.execute_code) { return codeOutputStrategy(); } else { diff --git a/client/src/common/assistants-types.ts b/client/src/common/assistants-types.ts index f54a8416909..d75f2bac762 100644 --- a/client/src/common/assistants-types.ts +++ b/client/src/common/assistants-types.ts @@ -10,12 +10,14 @@ export type TAssistantOption = Assistant & { files?: Array<[string, ExtendedFile]>; code_files?: Array<[string, ExtendedFile]>; + search_files?: Array<[string, ExtendedFile]>; }); export type Actions = { [Capabilities.code_interpreter]: boolean; [Capabilities.image_vision]: boolean; [Capabilities.retrieval]: boolean; + [Capabilities.file_search]: boolean; }; export type AssistantForm = { diff --git a/client/src/components/SidePanel/Builder/AssistantSelect.tsx b/client/src/components/SidePanel/Builder/AssistantSelect.tsx index d3202cdd2c6..8e6ce734588 100644 --- a/client/src/components/SidePanel/Builder/AssistantSelect.tsx +++ b/client/src/components/SidePanel/Builder/AssistantSelect.tsx @@ -1,5 +1,5 @@ import { Plus } from 'lucide-react'; -import { useCallback, useEffect, useRef } from 'react'; +import { useCallback, useEffect, useRef, useMemo } from 'react'; import { Tools, FileSources, @@ -78,6 +78,9 @@ export default function AssistantSelect({ code_files: _assistant.tool_resources?.code_interpreter?.file_ids ? ([] as Array<[string, ExtendedFile]>) : undefined, + search_files: _assistant.tool_resources?.file_search?.vector_store_ids + ? ([] as Array<[string, ExtendedFile]>) + : undefined, }; const handleFile = (file_id: string, list?: Array<[string, ExtendedFile]>) => { @@ -124,6 +127,12 @@ export default function AssistantSelect({ ); } + if (assistant.search_files && _assistant.tool_resources?.file_search?.vector_store_ids) { + _assistant.tool_resources.file_search.vector_store_ids.forEach((file_id) => + handleFile(file_id, assistant.search_files), + ); + } + const assistantDoc = documentsMap?.get(_assistant.id); /* If no user updates, use the latest assistant docs */ if (assistantDoc) { @@ -159,6 +168,7 @@ export default function AssistantSelect({ const actions: Actions = { [Capabilities.code_interpreter]: false, [Capabilities.image_vision]: false, + [Capabilities.file_search]: false, [Capabilities.retrieval]: false, }; @@ -167,7 +177,7 @@ export default function AssistantSelect({ .map((tool) => tool.function?.name || tool.type) .forEach((tool) => { if (tool === Tools.file_search) { - actions[Capabilities.retrieval] = true; + actions[Capabilities.file_search] = true; } actions[tool] = true; }); diff --git a/client/src/components/SidePanel/Builder/CapabilitiesForm.tsx b/client/src/components/SidePanel/Builder/CapabilitiesForm.tsx index efd7227cf04..4dee0448834 100644 --- a/client/src/components/SidePanel/Builder/CapabilitiesForm.tsx +++ b/client/src/components/SidePanel/Builder/CapabilitiesForm.tsx @@ -1,5 +1,5 @@ import { useMemo } from 'react'; -import { Capabilities } from 'librechat-data-provider'; +import { Capabilities, EToolResources } from 'librechat-data-provider'; import { useFormContext, useWatch } from 'react-hook-form'; import type { TConfig, AssistantsEndpoint } from 'librechat-data-provider'; import type { AssistantForm } from '~/common'; @@ -32,7 +32,14 @@ export default function CapabilitiesForm({ if (typeof assistant === 'string') { return []; } - return assistant.code_files; + return assistant.code_files ?? []; + }, [assistant]); + + const fileSearch = useMemo(() => { + if (typeof assistant === 'string') { + return []; + } + return assistant.search_files ?? []; }, [assistant]); const retrievalModels = useMemo( @@ -55,6 +62,15 @@ export default function CapabilitiesForm({
{codeEnabled && } + {codeEnabled && version && ( + + )} {retrievalEnabled && ( )} @@ -64,7 +80,8 @@ export default function CapabilitiesForm({ assistant_id={assistant_id} version={version} endpoint={endpoint} - files={files} + files={fileSearch} + tool_resource={EToolResources.file_search} /> )}
diff --git a/client/src/components/SidePanel/Builder/CodeFiles.tsx b/client/src/components/SidePanel/Builder/CodeFiles.tsx index 22be0a5dcde..c77a18c19ed 100644 --- a/client/src/components/SidePanel/Builder/CodeFiles.tsx +++ b/client/src/components/SidePanel/Builder/CodeFiles.tsx @@ -12,17 +12,19 @@ import { useFileHandling } from '~/hooks/Files'; import useLocalize from '~/hooks/useLocalize'; import { useChatContext } from '~/Providers'; -const tool_resource = EToolResources.code_interpreter; +const default_tool_resource = EToolResources.code_interpreter; export default function CodeFiles({ endpoint, assistant_id, files: _files, + tool_resource = default_tool_resource, }: { version: number | string; endpoint: AssistantsEndpoint; assistant_id: string; files?: [string, ExtendedFile][]; + tool_resource?: EToolResources; }) { const localize = useLocalize(); const { setFilesLoading } = useChatContext(); @@ -61,9 +63,9 @@ export default function CodeFiles({ return (
-
+ {/*
{localize('com_assistants_code_interpreter_files')} -
+
*/} - {localize('com_ui_upload_files')} + + {tool_resource === EToolResources.code_interpreter + ? localize('com_ui_upload_code_files') + : tool_resource === EToolResources.file_search + ? localize('com_ui_upload_file_search') + : null}
diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 506fe6b3fd6..a9ae8013065 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -135,6 +135,7 @@ export enum Capabilities { code_interpreter = 'code_interpreter', image_vision = 'image_vision', retrieval = 'retrieval', + file_search = 'file_search', actions = 'actions', tools = 'tools', } @@ -150,7 +151,7 @@ export enum AgentCapabilities { export const defaultAssistantsVersion = { [EModelEndpoint.assistants]: 2, - [EModelEndpoint.azureAssistants]: 1, + [EModelEndpoint.azureAssistants]: 2, }; export const baseEndpointSchema = z.object({ @@ -707,7 +708,7 @@ export const EndpointURLs: { [key in EModelEndpoint]: string } = { [EModelEndpoint.gptPlugins]: `/api/ask/${EModelEndpoint.gptPlugins}`, [EModelEndpoint.azureOpenAI]: `/api/ask/${EModelEndpoint.azureOpenAI}`, [EModelEndpoint.chatGPTBrowser]: `/api/ask/${EModelEndpoint.chatGPTBrowser}`, - [EModelEndpoint.azureAssistants]: '/api/assistants/v1/chat', + [EModelEndpoint.azureAssistants]: '/api/assistants/v2/chat', [EModelEndpoint.assistants]: '/api/assistants/v2/chat', [EModelEndpoint.agents]: `/api/${EModelEndpoint.agents}/chat`, [EModelEndpoint.bedrock]: `/api/${EModelEndpoint.bedrock}/chat`, diff --git a/packages/data-provider/src/types/files.ts b/packages/data-provider/src/types/files.ts index 5985096f4c1..fa2ceade40a 100644 --- a/packages/data-provider/src/types/files.ts +++ b/packages/data-provider/src/types/files.ts @@ -8,6 +8,7 @@ export enum FileSources { s3 = 's3', vectordb = 'vectordb', execute_code = 'execute_code', + vector_store = "vector_store", } export const checkOpenAIStorage = (source: string) => From f6d684bb90545e47360f0bf6880ab28886f9f06a Mon Sep 17 00:00:00 2001 From: Adnan Hussain Date: Tue, 17 Dec 2024 14:34:39 +0500 Subject: [PATCH 2/3] chore: remove unused file --- api/server/services/Files/process.js | 1 - 1 file changed, 1 deletion(-) diff --git a/api/server/services/Files/process.js b/api/server/services/Files/process.js index 56ee7365ad7..42e58db4097 100644 --- a/api/server/services/Files/process.js +++ b/api/server/services/Files/process.js @@ -22,7 +22,6 @@ const { addResourceFileId, deleteResourceFileId, addResourceVectorId, - deleteResourceVectorId, } = require('~/server/controllers/assistants/v2'); const { convertImage, resizeAndConvert } = require('~/server/services/Files/images'); const { addAgentResourceFile, removeAgentResourceFiles } = require('~/models/Agent'); From 5a2bb03410ebcc58da01da36c69cb13a30365878 Mon Sep 17 00:00:00 2001 From: Adnan Hussain Date: Wed, 18 Dec 2024 15:55:40 +0500 Subject: [PATCH 3/3] chore: add file search in tools --- client/src/components/SidePanel/Builder/AssistantPanel.tsx | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/client/src/components/SidePanel/Builder/AssistantPanel.tsx b/client/src/components/SidePanel/Builder/AssistantPanel.tsx index 63f2587da07..765e7d818a9 100644 --- a/client/src/components/SidePanel/Builder/AssistantPanel.tsx +++ b/client/src/components/SidePanel/Builder/AssistantPanel.tsx @@ -82,7 +82,7 @@ export default function AssistantPanel({ [assistantsConfig], ); const retrievalEnabled = useMemo( - () => assistantsConfig?.capabilities?.includes(Capabilities.retrieval), + () => assistantsConfig?.capabilities?.includes(Capabilities.file_search), [assistantsConfig], ); const codeEnabled = useMemo( @@ -155,6 +155,9 @@ export default function AssistantPanel({ if (data.code_interpreter) { tools.push({ type: Tools.code_interpreter }); } + if (data.file_search) { + tools.push({ type: Tools.file_search }); + } if (data.retrieval) { tools.push({ type: version == 2 ? Tools.file_search : Tools.retrieval }); }