mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-03 22:32:42 +00:00
🧹 chore: Optimize provisioning — single credential load, deferred DB writes
- loadCodeApiKey: load CODE_API_KEY once per request, pass to both checkSessionsAlive and provisionToCodeEnv (was N+1 lookups) - provisionToCodeEnv/provisionToVectorDB now return fileUpdate objects instead of writing to DB immediately - primeResources batches all DB updates via Promise.allSettled after provisioning completes - Remove updateFile import from provision.js (no longer writes directly)
This commit is contained in:
Danny Avila
parent
f3f10e9806
commit
4080f65f3c
4 changed files with 109 additions and 48 deletions
|
|
@ -40,6 +40,7 @@ const {
|
|||
enrichLoadedToolsWithAgentContext,
|
||||
} = require('./skillDeps');
|
||||
const {
|
||||
loadCodeApiKey,
|
||||
provisionToCodeEnv,
|
||||
provisionToVectorDB,
|
||||
checkSessionsAlive,
|
||||
|
|
@ -409,6 +410,8 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
provisionToCodeEnv,
|
||||
provisionToVectorDB,
|
||||
checkSessionsAlive,
|
||||
loadCodeApiKey,
|
||||
updateFile: db.updateFile,
|
||||
},
|
||||
);
|
||||
|
||||
|
|
@ -488,6 +491,8 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
provisionToCodeEnv,
|
||||
provisionToVectorDB,
|
||||
checkSessionsAlive,
|
||||
loadCodeApiKey,
|
||||
updateFile: db.updateFile,
|
||||
},
|
||||
// The callback fires during BFS, before the helper prunes agents
|
||||
// whose edges end up filtered. Don't populate `agentConfigs` here —
|
||||
|
|
|
|||
|
|
@ -12,10 +12,21 @@ const { logger } = require('@librechat/data-schemas');
|
|||
const { FileSources } = require('librechat-data-provider');
|
||||
const { loadAuthValues } = require('~/server/services/Tools/credentials');
|
||||
const { getStrategyFunctions } = require('./strategies');
|
||||
const { updateFile } = require('~/models');
|
||||
|
||||
const axios = createAxiosInstance();
|
||||
|
||||
/**
|
||||
* Loads the CODE_API_KEY for a user. Call once per request and pass the result
|
||||
* to provisionToCodeEnv / checkSessionsAlive to avoid redundant lookups.
|
||||
*
|
||||
* @param {string} userId
|
||||
* @returns {Promise<string>} The CODE_API_KEY
|
||||
*/
|
||||
async function loadCodeApiKey(userId) {
|
||||
const result = await loadAuthValues({ userId, authFields: [EnvVar.CODE_API_KEY] });
|
||||
return result[EnvVar.CODE_API_KEY];
|
||||
}
|
||||
|
||||
/**
|
||||
* Provisions a file to the code execution environment.
|
||||
* Gets a read stream from our storage and uploads to the code env.
|
||||
|
|
@ -24,9 +35,10 @@ const axios = createAxiosInstance();
|
|||
* @param {object} params.req - Express request object (needs req.user.id)
|
||||
* @param {import('librechat-data-provider').TFile} params.file - The file record from DB
|
||||
* @param {string} [params.entity_id] - Optional entity ID (agent_id)
|
||||
* @returns {Promise<string>} The fileIdentifier from the code env
|
||||
* @param {string} [params.apiKey] - Pre-loaded CODE_API_KEY (avoids redundant loadAuthValues)
|
||||
* @returns {Promise<{ fileIdentifier: string, fileUpdate: object }>} Result with deferred DB update
|
||||
*/
|
||||
async function provisionToCodeEnv({ req, file, entity_id = '' }) {
|
||||
async function provisionToCodeEnv({ req, file, entity_id = '', apiKey }) {
|
||||
const { getDownloadStream } = getStrategyFunctions(file.source);
|
||||
if (!getDownloadStream) {
|
||||
throw new Error(
|
||||
|
|
@ -34,33 +46,26 @@ async function provisionToCodeEnv({ req, file, entity_id = '' }) {
|
|||
);
|
||||
}
|
||||
|
||||
const resolvedApiKey = apiKey ?? (await loadCodeApiKey(req.user.id));
|
||||
const { handleFileUpload: uploadCodeEnvFile } = getStrategyFunctions(FileSources.execute_code);
|
||||
const result = await loadAuthValues({ userId: req.user.id, authFields: [EnvVar.CODE_API_KEY] });
|
||||
const stream = await getDownloadStream(req, file.filepath);
|
||||
|
||||
const fileIdentifier = await uploadCodeEnvFile({
|
||||
req,
|
||||
stream,
|
||||
filename: file.filename,
|
||||
apiKey: result[EnvVar.CODE_API_KEY],
|
||||
apiKey: resolvedApiKey,
|
||||
entity_id,
|
||||
});
|
||||
|
||||
const updatedMetadata = {
|
||||
...file.metadata,
|
||||
fileIdentifier,
|
||||
};
|
||||
|
||||
await updateFile({
|
||||
file_id: file.file_id,
|
||||
metadata: updatedMetadata,
|
||||
});
|
||||
|
||||
logger.debug(
|
||||
`[provisionToCodeEnv] Provisioned file "${file.filename}" (${file.file_id}) to code env`,
|
||||
);
|
||||
|
||||
return fileIdentifier;
|
||||
return {
|
||||
fileIdentifier,
|
||||
fileUpdate: { file_id: file.file_id, metadata: { ...file.metadata, fileIdentifier } },
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -71,27 +76,31 @@ async function provisionToCodeEnv({ req, file, entity_id = '' }) {
|
|||
* @param {object} params.req - Express request object
|
||||
* @param {import('librechat-data-provider').TFile} params.file - The file record from DB
|
||||
* @param {string} [params.entity_id] - Optional entity ID (agent_id)
|
||||
* @returns {Promise<{ embedded: boolean }>} Embedding result
|
||||
* @param {import('stream').Readable} [params.existingStream] - Pre-fetched download stream (avoids duplicate storage fetch)
|
||||
* @returns {Promise<{ embedded: boolean, fileUpdate: object | null }>} Result with deferred DB update
|
||||
*/
|
||||
async function provisionToVectorDB({ req, file, entity_id }) {
|
||||
async function provisionToVectorDB({ req, file, entity_id, existingStream }) {
|
||||
if (!process.env.RAG_API_URL) {
|
||||
logger.warn('[provisionToVectorDB] RAG_API_URL not defined, skipping vector provisioning');
|
||||
return { embedded: false };
|
||||
return { embedded: false, fileUpdate: null };
|
||||
}
|
||||
|
||||
const { getDownloadStream } = getStrategyFunctions(file.source);
|
||||
if (!getDownloadStream) {
|
||||
throw new Error(
|
||||
`Cannot provision file "${file.filename}" to vector DB: storage source "${file.source}" does not support download streams`,
|
||||
);
|
||||
}
|
||||
|
||||
// The uploadVectors function expects a file-like object with a `path` property for fs.createReadStream.
|
||||
// Since we're provisioning from storage (not a multer upload), we need to stream to a temp file first.
|
||||
const tmpPath = path.join(os.tmpdir(), `provision-${file.file_id}${path.extname(file.filename)}`);
|
||||
|
||||
try {
|
||||
const stream = await getDownloadStream(req, file.filepath);
|
||||
let stream = existingStream;
|
||||
if (!stream) {
|
||||
const { getDownloadStream } = getStrategyFunctions(file.source);
|
||||
if (!getDownloadStream) {
|
||||
throw new Error(
|
||||
`Cannot provision file "${file.filename}" to vector DB: storage source "${file.source}" does not support download streams`,
|
||||
);
|
||||
}
|
||||
stream = await getDownloadStream(req, file.filepath);
|
||||
}
|
||||
|
||||
// uploadVectors expects a file-like object with a `path` property for fs.createReadStream.
|
||||
// Since we're provisioning from storage (not a multer upload), we stream to a temp file first.
|
||||
await new Promise((resolve, reject) => {
|
||||
const writeStream = fs.createWriteStream(tmpPath);
|
||||
stream.pipe(writeStream);
|
||||
|
|
@ -117,18 +126,15 @@ async function provisionToVectorDB({ req, file, entity_id }) {
|
|||
|
||||
const embedded = embeddingResult?.embedded ?? false;
|
||||
|
||||
await updateFile({
|
||||
file_id: file.file_id,
|
||||
embedded,
|
||||
});
|
||||
|
||||
logger.debug(
|
||||
`[provisionToVectorDB] Provisioned file "${file.filename}" (${file.file_id}) to vector DB, embedded=${embedded}`,
|
||||
);
|
||||
|
||||
return { embedded };
|
||||
return {
|
||||
embedded,
|
||||
fileUpdate: embedded ? { file_id: file.file_id, embedded } : null,
|
||||
};
|
||||
} finally {
|
||||
// Clean up temp file
|
||||
try {
|
||||
fs.unlinkSync(tmpPath);
|
||||
} catch {
|
||||
|
|
@ -190,13 +196,11 @@ async function checkCodeEnvFileAlive({ file, apiKey }) {
|
|||
*
|
||||
* @param {object} params
|
||||
* @param {import('librechat-data-provider').TFile[]} params.files - Files with metadata.fileIdentifier
|
||||
* @param {string} params.userId - User ID for loading CODE_API_KEY
|
||||
* @param {string} params.apiKey - Pre-loaded CODE_API_KEY
|
||||
* @param {number} [params.staleSafeWindowMs=21600000] - Skip check if file updated within this window (default 6h)
|
||||
* @returns {Promise<Set<string>>} Set of file_ids that are confirmed alive
|
||||
*/
|
||||
async function checkSessionsAlive({ files, userId, staleSafeWindowMs = 6 * 60 * 60 * 1000 }) {
|
||||
const result = await loadAuthValues({ userId, authFields: [EnvVar.CODE_API_KEY] });
|
||||
const apiKey = result[EnvVar.CODE_API_KEY];
|
||||
async function checkSessionsAlive({ files, apiKey, staleSafeWindowMs = 6 * 60 * 60 * 1000 }) {
|
||||
const aliveFileIds = new Set();
|
||||
const now = Date.now();
|
||||
|
||||
|
|
@ -263,6 +267,7 @@ async function checkSessionsAlive({ files, userId, staleSafeWindowMs = 6 * 60 *
|
|||
}
|
||||
|
||||
module.exports = {
|
||||
loadCodeApiKey,
|
||||
provisionToCodeEnv,
|
||||
provisionToVectorDB,
|
||||
checkCodeEnvFileAlive,
|
||||
|
|
|
|||
|
|
@ -32,10 +32,12 @@ import type {
|
|||
import type { LCAvailableTools, RequestScopedMCPConnectionStore } from '../mcp/types';
|
||||
import type { ResolvedManualSkill, ResolvedAlwaysApplySkill } from './skills';
|
||||
import type {
|
||||
TFileUpdate,
|
||||
TFilterFilesByAgentAccess,
|
||||
TProvisionToCodeEnv,
|
||||
TProvisionToVectorDB,
|
||||
TCheckSessionsAlive,
|
||||
TLoadCodeApiKey,
|
||||
} from './resources';
|
||||
import {
|
||||
injectSkillCatalog,
|
||||
|
|
@ -555,6 +557,10 @@ export interface InitializeAgentDbMethods extends EndpointDbMethods {
|
|||
provisionToVectorDB?: TProvisionToVectorDB;
|
||||
/** Optional: batch-check code env file liveness */
|
||||
checkSessionsAlive?: TCheckSessionsAlive;
|
||||
/** Optional: load CODE_API_KEY once per request */
|
||||
loadCodeApiKey?: TLoadCodeApiKey;
|
||||
/** Optional: persist file metadata updates after provisioning */
|
||||
updateFile?: (data: TFileUpdate) => Promise<unknown>;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -772,6 +778,8 @@ export async function initializeAgent(
|
|||
provisionToCodeEnv: db.provisionToCodeEnv,
|
||||
provisionToVectorDB: db.provisionToVectorDB,
|
||||
checkSessionsAlive: db.checkSessionsAlive,
|
||||
loadCodeApiKey: db.loadCodeApiKey,
|
||||
updateFile: db.updateFile as ((data: TFileUpdate) => Promise<unknown>) | undefined,
|
||||
});
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -5,25 +5,34 @@ import type { IMongoFile, AppConfig, IUser } from '@librechat/data-schemas';
|
|||
import type { FilterQuery, QueryOptions, ProjectionType } from 'mongoose';
|
||||
import type { Request as ServerRequest } from 'express';
|
||||
|
||||
/** Deferred DB update from provisioning (batched after all files are provisioned) */
|
||||
export type TFileUpdate = {
|
||||
file_id: string;
|
||||
metadata?: Record<string, unknown>;
|
||||
embedded?: boolean;
|
||||
};
|
||||
|
||||
/**
|
||||
* Function type for provisioning a file to the code execution environment.
|
||||
* @returns The fileIdentifier from the code env
|
||||
* @returns The fileIdentifier and a deferred DB update object
|
||||
*/
|
||||
export type TProvisionToCodeEnv = (params: {
|
||||
req: ServerRequest & { user?: IUser };
|
||||
file: TFile;
|
||||
entity_id?: string;
|
||||
}) => Promise<string>;
|
||||
apiKey?: string;
|
||||
}) => Promise<{ fileIdentifier: string; fileUpdate: TFileUpdate }>;
|
||||
|
||||
/**
|
||||
* Function type for provisioning a file to the vector DB for file_search.
|
||||
* @returns Object with embedded status
|
||||
* @returns Object with embedded status and a deferred DB update object
|
||||
*/
|
||||
export type TProvisionToVectorDB = (params: {
|
||||
req: ServerRequest & { user?: IUser };
|
||||
file: TFile;
|
||||
entity_id?: string;
|
||||
}) => Promise<{ embedded: boolean }>;
|
||||
existingStream?: unknown;
|
||||
}) => Promise<{ embedded: boolean; fileUpdate: TFileUpdate | null }>;
|
||||
|
||||
/**
|
||||
* Function type for batch-checking code env file liveness.
|
||||
|
|
@ -32,10 +41,13 @@ export type TProvisionToVectorDB = (params: {
|
|||
*/
|
||||
export type TCheckSessionsAlive = (params: {
|
||||
files: TFile[];
|
||||
userId: string;
|
||||
apiKey: string;
|
||||
staleSafeWindowMs?: number;
|
||||
}) => Promise<Set<string>>;
|
||||
|
||||
/** Loads CODE_API_KEY for a user. Call once per request. */
|
||||
export type TLoadCodeApiKey = (userId: string) => Promise<string>;
|
||||
|
||||
/**
|
||||
* Function type for retrieving files from the database
|
||||
* @param filter - MongoDB filter query for files
|
||||
|
|
@ -196,6 +208,8 @@ export const primeResources = async ({
|
|||
provisionToCodeEnv,
|
||||
provisionToVectorDB,
|
||||
checkSessionsAlive,
|
||||
loadCodeApiKey,
|
||||
updateFile,
|
||||
}: {
|
||||
req: ServerRequest & { user?: IUser };
|
||||
appConfig?: AppConfig;
|
||||
|
|
@ -213,6 +227,10 @@ export const primeResources = async ({
|
|||
provisionToVectorDB?: TProvisionToVectorDB;
|
||||
/** Optional callback to batch-check code env file liveness by session */
|
||||
checkSessionsAlive?: TCheckSessionsAlive;
|
||||
/** Optional callback to load CODE_API_KEY once per request */
|
||||
loadCodeApiKey?: TLoadCodeApiKey;
|
||||
/** Optional callback to persist file metadata updates after provisioning */
|
||||
updateFile?: (data: TFileUpdate) => Promise<unknown>;
|
||||
}): Promise<{
|
||||
attachments: Array<TFile | undefined> | undefined;
|
||||
requestAttachments: Array<TFile | undefined> | undefined;
|
||||
|
|
@ -386,20 +404,34 @@ export const primeResources = async ({
|
|||
enabledToolResources.has(EToolResources.file_search) && provisionToVectorDB != null;
|
||||
|
||||
if (needsCodeEnv || needsVectorDB) {
|
||||
// Load CODE_API_KEY once for all code env operations
|
||||
let codeApiKey: string | undefined;
|
||||
if (needsCodeEnv && loadCodeApiKey && req.user?.id) {
|
||||
try {
|
||||
codeApiKey = await loadCodeApiKey(req.user.id);
|
||||
} catch (error) {
|
||||
logger.error('[primeResources] Failed to load CODE_API_KEY', error);
|
||||
warnings.push('Code execution file provisioning unavailable');
|
||||
}
|
||||
}
|
||||
|
||||
// Batch staleness check: verify code env files are still alive
|
||||
let aliveFileIds: Set<string> = new Set();
|
||||
if (needsCodeEnv && checkSessionsAlive && req.user?.id) {
|
||||
if (needsCodeEnv && codeApiKey && checkSessionsAlive) {
|
||||
const filesWithIdentifiers = attachments.filter(
|
||||
(f) => f?.metadata?.fileIdentifier && f.file_id,
|
||||
);
|
||||
if (filesWithIdentifiers.length > 0) {
|
||||
aliveFileIds = await checkSessionsAlive({
|
||||
files: filesWithIdentifiers as TFile[],
|
||||
userId: req.user.id,
|
||||
apiKey: codeApiKey,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Collect deferred DB updates from provisioning
|
||||
const pendingUpdates: TFileUpdate[] = [];
|
||||
|
||||
// Provision files in parallel
|
||||
const provisionResults = await Promise.allSettled(
|
||||
attachments.map(async (file) => {
|
||||
|
|
@ -413,6 +445,7 @@ export const primeResources = async ({
|
|||
// Code env provisioning (with staleness check)
|
||||
if (
|
||||
needsCodeEnv &&
|
||||
codeApiKey &&
|
||||
!processedResourceFiles.has(`${EToolResources.execute_code}:${file.file_id}`)
|
||||
) {
|
||||
const hasFileIdentifier = !!file.metadata?.fileIdentifier;
|
||||
|
|
@ -428,12 +461,14 @@ export const primeResources = async ({
|
|||
}
|
||||
|
||||
try {
|
||||
const fileIdentifier = await provisionToCodeEnv({
|
||||
const { fileIdentifier, fileUpdate } = await provisionToCodeEnv({
|
||||
req: typedReq,
|
||||
file,
|
||||
entity_id: agentId,
|
||||
apiKey: codeApiKey,
|
||||
});
|
||||
file.metadata = { ...file.metadata, fileIdentifier };
|
||||
pendingUpdates.push(fileUpdate);
|
||||
addFileToResource({
|
||||
file,
|
||||
resourceType: EToolResources.execute_code,
|
||||
|
|
@ -471,6 +506,9 @@ export const primeResources = async ({
|
|||
});
|
||||
if (result.embedded) {
|
||||
file.embedded = true;
|
||||
if (result.fileUpdate) {
|
||||
pendingUpdates.push(result.fileUpdate);
|
||||
}
|
||||
addFileToResource({
|
||||
file,
|
||||
resourceType: EToolResources.file_search,
|
||||
|
|
@ -493,6 +531,11 @@ export const primeResources = async ({
|
|||
logger.error('[primeResources] Unexpected provisioning rejection', result.reason);
|
||||
}
|
||||
}
|
||||
|
||||
// Batch DB updates after all provisioning completes
|
||||
if (pendingUpdates.length > 0 && updateFile) {
|
||||
await Promise.allSettled(pendingUpdates.map((update) => updateFile(update)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue