Andreas/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2026-04-04 23:07:19 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

📦 refactor: Consolidate DB models, encapsulating Mongoose usage in data-schemas (#11830)

Browse source 
* chore: move database model methods to /packages/data-schemas

* chore: add TypeScript ESLint rule to warn on unused variables

* refactor: model imports to streamline access

- Consolidated model imports across various files to improve code organization and reduce redundancy.
- Updated imports for models such as Assistant, Message, Conversation, and others to a unified import path.
- Adjusted middleware and service files to reflect the new import structure, ensuring functionality remains intact.
- Enhanced test files to align with the new import paths, maintaining test coverage and integrity.

* chore: migrate database models to packages/data-schemas and refactor all direct Mongoose Model usage outside of data-schemas

* test: update agent model mocks in unit tests

- Added `getAgent` mock to `client.test.js` to enhance test coverage for agent-related functionality.
- Removed redundant `getAgent` and `getAgents` mocks from `openai.spec.js` and `responses.unit.spec.js` to streamline test setup and reduce duplication.
- Ensured consistency in agent mock implementations across test files.

* fix: update types in data-schemas

* refactor: enhance type definitions in transaction and spending methods

- Updated type definitions in `checkBalance.ts` to use specific request and response types.
- Refined `spendTokens.ts` to utilize a new `SpendTxData` interface for better clarity and type safety.
- Improved transaction handling in `transaction.ts` by introducing `TransactionResult` and `TxData` interfaces, ensuring consistent data structures across methods.
- Adjusted unit tests in `transaction.spec.ts` to accommodate new type definitions and enhance robustness.

* refactor: streamline model imports and enhance code organization

- Consolidated model imports across various controllers and services to a unified import path, improving code clarity and reducing redundancy.
- Updated multiple files to reflect the new import structure, ensuring all functionalities remain intact.
- Enhanced overall code organization by removing duplicate import statements and optimizing the usage of model methods.

* feat: implement loadAddedAgent and refactor agent loading logic

- Introduced `loadAddedAgent` function to handle loading agents from added conversations, supporting multi-convo parallel execution.
- Created a new `load.ts` file to encapsulate agent loading functionalities, including `loadEphemeralAgent` and `loadAgent`.
- Updated the `index.ts` file to export the new `load` module instead of the deprecated `loadAgent`.
- Enhanced type definitions and improved error handling in the agent loading process.
- Adjusted unit tests to reflect changes in the agent loading structure and ensure comprehensive coverage.

* refactor: enhance balance handling with new update interface

- Introduced `IBalanceUpdate` interface to streamline balance update operations across the codebase.
- Updated `upsertBalanceFields` method signatures in `balance.ts`, `transaction.ts`, and related tests to utilize the new interface for improved type safety.
- Adjusted type imports in `balance.spec.ts` to include `IBalanceUpdate`, ensuring consistency in balance management functionalities.
- Enhanced overall code clarity and maintainability by refining type definitions related to balance operations.

* feat: add unit tests for loadAgent functionality and enhance agent loading logic

- Introduced comprehensive unit tests for the `loadAgent` function, covering various scenarios including null and empty agent IDs, loading of ephemeral agents, and permission checks.
- Enhanced the `initializeClient` function by moving `getConvoFiles` to the correct position in the database method exports, ensuring proper functionality.
- Improved test coverage for agent loading, including handling of non-existent agents and user permissions.

* chore: reorder memory method exports for consistency

- Moved `deleteAllUserMemories` to the correct position in the exported memory methods, ensuring a consistent and logical order of method exports in `memory.ts`.
This commit is contained in:
Danny Avila 2026-02-17 18:23:44 -05:00
parent 58f128bee7
commit 8ba2bde5c1
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
182 changed files with 8675 additions and 8457 deletions
Download patch file Download diff file Expand all files Collapse all files

8
packages/api/src/middleware/access.spec.ts
View file

@ -216,12 +216,12 @@ describe('access middleware', () => {
defaultParams.getRoleByName.mockResolvedValue(mockRole);
const checkObject = {};
const checkObject = { id: 'agent123' };
const result = await checkAccess({
...defaultParams,
permissions: [Permissions.USE, Permissions.SHARE],
bodyProps: {} as Record<Permissions, string[]>,
bodyProps: { [Permissions.SHARE]: ['id'] } as Record<Permissions, string[]>,
checkObject,
});
expect(result).toBe(true);
@ -333,12 +333,12 @@ describe('access middleware', () => {
} as unknown as IRole;
mockGetRoleByName.mockResolvedValue(mockRole);
mockReq.body = {};
mockReq.body = { id: 'agent123' };
const middleware = generateCheckAccess({
permissionType: PermissionTypes.AGENTS,
permissions: [Permissions.USE, Permissions.CREATE, Permissions.SHARE],
bodyProps: {} as Record<Permissions, string[]>,
bodyProps: { [Permissions.SHARE]: ['id'] } as Record<Permissions, string[]>,
getRoleByName: mockGetRoleByName,
});

84
packages/api/src/middleware/balance.spec.ts
View file

@ -2,7 +2,7 @@ import mongoose from 'mongoose';
import { MongoMemoryServer } from 'mongodb-memory-server';
import { logger, balanceSchema } from '@librechat/data-schemas';
import type { NextFunction, Request as ServerRequest, Response as ServerResponse } from 'express';
import type { IBalance } from '@librechat/data-schemas';
import type { IBalance, IBalanceUpdate } from '@librechat/data-schemas';
import { createSetBalanceConfig } from './balance';
jest.mock('@librechat/data-schemas', () => ({
@ -15,6 +15,16 @@ jest.mock('@librechat/data-schemas', () => ({
let mongoServer: MongoMemoryServer;
let Balance: mongoose.Model<IBalance>;
const findBalanceByUser = (userId: string) =>
Balance.findOne({ user: userId }).lean() as Promise<IBalance | null>;
const upsertBalanceFields = (userId: string, fields: IBalanceUpdate) =>
Balance.findOneAndUpdate(
{ user: userId },
{ $set: fields },
{ upsert: true, new: true },
).lean() as Promise<IBalance | null>;
beforeAll(async () => {
mongoServer = await MongoMemoryServer.create();
const mongoUri = mongoServer.getUri();
@ -64,7 +74,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -95,7 +106,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -120,7 +132,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -149,7 +162,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -178,7 +192,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = {} as ServerRequest;
@ -219,7 +234,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -271,7 +287,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -315,7 +332,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -346,7 +364,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -392,7 +411,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -434,21 +454,20 @@ describe('createSetBalanceConfig', () => {
},
});
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
});
const req = createMockRequest(userId);
const res = createMockResponse();
// Spy on Balance.findOneAndUpdate to verify it's not called
const updateSpy = jest.spyOn(Balance, 'findOneAndUpdate');
const upsertSpy = jest.fn();
const spiedMiddleware = createSetBalanceConfig({
getAppConfig,
findBalanceByUser,
upsertBalanceFields: upsertSpy,
});
await middleware(req as ServerRequest, res as ServerResponse, mockNext);
await spiedMiddleware(req as ServerRequest, res as ServerResponse, mockNext);
expect(mockNext).toHaveBeenCalled();
expect(updateSpy).not.toHaveBeenCalled();
expect(upsertSpy).not.toHaveBeenCalled();
});
test('should set tokenCredits for user with null tokenCredits', async () => {
@ -470,7 +489,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -498,16 +518,12 @@ describe('createSetBalanceConfig', () => {
});
const dbError = new Error('Database error');
// Mock Balance.findOne to throw an error
jest.spyOn(Balance, 'findOne').mockImplementationOnce((() => {
return {
lean: jest.fn().mockRejectedValue(dbError),
};
}) as unknown as mongoose.Model<IBalance>['findOne']);
const failingFindBalance = () => Promise.reject(dbError);
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser: failingFindBalance,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -526,7 +542,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -556,7 +573,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -590,7 +608,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);
@ -635,7 +654,8 @@ describe('createSetBalanceConfig', () => {
const middleware = createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
});
const req = createMockRequest(userId);

24
packages/api/src/middleware/balance.ts
View file

@ -1,13 +1,20 @@
import { logger } from '@librechat/data-schemas';
import type {
IBalanceUpdate,
BalanceConfig,
AppConfig,
ObjectId,
IBalance,
IUser,
} from '@librechat/data-schemas';
import type { NextFunction, Request as ServerRequest, Response as ServerResponse } from 'express';
import type { IBalance, IUser, BalanceConfig, ObjectId, AppConfig } from '@librechat/data-schemas';
import type { Model } from 'mongoose';
import type { BalanceUpdateFields } from '~/types';
import { getBalanceConfig } from '~/app/config';
export interface BalanceMiddlewareOptions {
getAppConfig: (options?: { role?: string; refresh?: boolean }) => Promise<AppConfig>;
Balance: Model<IBalance>;
findBalanceByUser: (userId: string) => Promise<IBalance | null>;
upsertBalanceFields: (userId: string, fields: IBalanceUpdate) => Promise<IBalance | null>;
}
/**
@ -75,7 +82,8 @@ function buildUpdateFields(
*/
export function createSetBalanceConfig({
getAppConfig,
Balance,
findBalanceByUser,
upsertBalanceFields,
}: BalanceMiddlewareOptions): (
req: ServerRequest,
res: ServerResponse,
@ -97,18 +105,14 @@ export function createSetBalanceConfig({
return next();
}
const userId = typeof user._id === 'string' ? user._id : user._id.toString();
const userBalanceRecord = await Balance.findOne({ user: userId }).lean();
const userBalanceRecord = await findBalanceByUser(userId);
const updateFields = buildUpdateFields(balanceConfig, userBalanceRecord, userId);
if (Object.keys(updateFields).length === 0) {
return next();
}
await Balance.findOneAndUpdate(
{ user: userId },
{ $set: updateFields },
{ upsert: true, new: true },
);
await upsertBalanceFields(userId, updateFields);
next();
} catch (error) {

168
packages/api/src/middleware/checkBalance.ts Normal file
View file

@ -0,0 +1,168 @@
import { logger } from '@librechat/data-schemas';
import { ViolationTypes } from 'librechat-data-provider';
import type { ServerRequest } from '~/types/http';
import type { Response } from 'express';
type TimeUnit = 'seconds' | 'minutes' | 'hours' | 'days' | 'weeks' | 'months';
interface BalanceRecord {
tokenCredits: number;
autoRefillEnabled?: boolean;
refillAmount?: number;
lastRefill?: Date;
refillIntervalValue?: number;
refillIntervalUnit?: TimeUnit;
}
interface TxData {
user: string;
model?: string;
endpoint?: string;
valueKey?: string;
tokenType?: string;
amount: number;
endpointTokenConfig?: unknown;
generations?: unknown[];
}
export interface CheckBalanceDeps {
findBalanceByUser: (user: string) => Promise<BalanceRecord | null>;
getMultiplier: (params: Record<string, unknown>) => number;
createAutoRefillTransaction: (
data: Record<string, unknown>,
) => Promise<{ balance: number } | undefined>;
logViolation: (
req: unknown,
res: unknown,
type: string,
errorMessage: Record<string, unknown>,
score: number,
) => Promise<void>;
}
function addIntervalToDate(date: Date, value: number, unit: TimeUnit): Date {
const result = new Date(date);
switch (unit) {
case 'seconds':
result.setSeconds(result.getSeconds() + value);
break;
case 'minutes':
result.setMinutes(result.getMinutes() + value);
break;
case 'hours':
result.setHours(result.getHours() + value);
break;
case 'days':
result.setDate(result.getDate() + value);
break;
case 'weeks':
result.setDate(result.getDate() + value * 7);
break;
case 'months':
result.setMonth(result.getMonth() + value);
break;
default:
break;
}
return result;
}
/** Checks a user's balance record and handles auto-refill if needed. */
async function checkBalanceRecord(
txData: TxData,
deps: CheckBalanceDeps,
): Promise<{ canSpend: boolean; balance: number; tokenCost: number }> {
const { user, model, endpoint, valueKey, tokenType, amount, endpointTokenConfig } = txData;
const multiplier = deps.getMultiplier({
valueKey,
tokenType,
model,
endpoint,
endpointTokenConfig,
});
const tokenCost = amount * multiplier;
const record = await deps.findBalanceByUser(user);
if (!record) {
logger.debug('[Balance.check] No balance record found for user', { user });
return { canSpend: false, balance: 0, tokenCost };
}
let balance = record.tokenCredits;
logger.debug('[Balance.check] Initial state', {
user,
model,
endpoint,
valueKey,
tokenType,
amount,
balance,
multiplier,
endpointTokenConfig: !!endpointTokenConfig,
});
if (
balance - tokenCost <= 0 &&
record.autoRefillEnabled &&
record.refillAmount &&
record.refillAmount > 0
) {
const lastRefillDate = new Date(record.lastRefill ?? 0);
const now = new Date();
if (
isNaN(lastRefillDate.getTime()) ||
now >=
addIntervalToDate(
lastRefillDate,
record.refillIntervalValue ?? 0,
record.refillIntervalUnit ?? 'days',
)
) {
try {
const result = await deps.createAutoRefillTransaction({
user,
tokenType: 'credits',
context: 'autoRefill',
rawAmount: record.refillAmount,
});
if (result) {
balance = result.balance;
}
} catch (error) {
logger.error('[Balance.check] Failed to record transaction for auto-refill', error);
}
}
}
logger.debug('[Balance.check] Token cost', { tokenCost });
return { canSpend: balance >= tokenCost, balance, tokenCost };
}
/**
* Checks balance for a user and logs a violation if they cannot spend.
* Throws an error with the balance info if insufficient funds.
*/
export async function checkBalance(
{ req, res, txData }: { req: ServerRequest; res: Response; txData: TxData },
deps: CheckBalanceDeps,
): Promise<boolean> {
const { canSpend, balance, tokenCost } = await checkBalanceRecord(txData, deps);
if (canSpend) {
return true;
}
const type = ViolationTypes.TOKEN_BALANCE;
const errorMessage: Record<string, unknown> = {
type,
balance,
tokenCost,
promptTokens: txData.amount,
};
if (txData.generations && txData.generations.length > 0) {
errorMessage.generations = txData.generations;
}
await deps.logViolation(req, res, type, errorMessage, 0);
throw new Error(JSON.stringify(errorMessage));
}

1
packages/api/src/middleware/index.ts
View file

@ -5,3 +5,4 @@ export * from './notFound';
export * from './balance';
export * from './json';
export * from './concurrency';
export * from './checkBalance';