mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-02-19 08:58:09 +01:00
📦 refactor: Consolidate DB models, encapsulating Mongoose usage in data-schemas (#11830)
* 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
parent
a85e99ff45
commit
a6fb257bcf
182 changed files with 8548 additions and 8105 deletions
145
packages/data-schemas/src/methods/spendTokens.ts
Normal file
145
packages/data-schemas/src/methods/spendTokens.ts
Normal file
|
|
@ -0,0 +1,145 @@
|
|||
import logger from '~/config/winston';
|
||||
import type { TxData, TransactionResult } from './transaction';
|
||||
|
||||
/** Base transaction context passed by callers — does not include fields added internally */
|
||||
export interface SpendTxData {
|
||||
user: string | import('mongoose').Types.ObjectId;
|
||||
conversationId?: string;
|
||||
model?: string;
|
||||
context?: string;
|
||||
endpointTokenConfig?: Record<string, Record<string, number>> | null;
|
||||
balance?: { enabled?: boolean };
|
||||
transactions?: { enabled?: boolean };
|
||||
valueKey?: string;
|
||||
}
|
||||
|
||||
export function createSpendTokensMethods(
|
||||
_mongoose: typeof import('mongoose'),
|
||||
transactionMethods: {
|
||||
createTransaction: (txData: TxData) => Promise<TransactionResult | undefined>;
|
||||
createStructuredTransaction: (txData: TxData) => Promise<TransactionResult | undefined>;
|
||||
},
|
||||
) {
|
||||
/**
|
||||
* Creates up to two transactions to record the spending of tokens.
|
||||
*/
|
||||
async function spendTokens(
|
||||
txData: SpendTxData,
|
||||
tokenUsage: { promptTokens?: number; completionTokens?: number },
|
||||
) {
|
||||
const { promptTokens, completionTokens } = tokenUsage;
|
||||
logger.debug(
|
||||
`[spendTokens] conversationId: ${txData.conversationId}${
|
||||
txData?.context ? ` | Context: ${txData?.context}` : ''
|
||||
} | Token usage: `,
|
||||
{ promptTokens, completionTokens },
|
||||
);
|
||||
let prompt: TransactionResult | undefined, completion: TransactionResult | undefined;
|
||||
const normalizedPromptTokens = Math.max(promptTokens ?? 0, 0);
|
||||
try {
|
||||
if (promptTokens !== undefined) {
|
||||
prompt = await transactionMethods.createTransaction({
|
||||
...txData,
|
||||
tokenType: 'prompt',
|
||||
rawAmount: promptTokens === 0 ? 0 : -normalizedPromptTokens,
|
||||
inputTokenCount: normalizedPromptTokens,
|
||||
});
|
||||
}
|
||||
|
||||
if (completionTokens !== undefined) {
|
||||
completion = await transactionMethods.createTransaction({
|
||||
...txData,
|
||||
tokenType: 'completion',
|
||||
rawAmount: completionTokens === 0 ? 0 : -Math.max(completionTokens, 0),
|
||||
inputTokenCount: normalizedPromptTokens,
|
||||
});
|
||||
}
|
||||
|
||||
if (prompt || completion) {
|
||||
logger.debug('[spendTokens] Transaction data record against balance:', {
|
||||
user: txData.user,
|
||||
prompt: prompt?.prompt,
|
||||
promptRate: prompt?.rate,
|
||||
completion: completion?.completion,
|
||||
completionRate: completion?.rate,
|
||||
balance: completion?.balance ?? prompt?.balance,
|
||||
});
|
||||
} else {
|
||||
logger.debug('[spendTokens] No transactions incurred against balance');
|
||||
}
|
||||
} catch (err) {
|
||||
logger.error('[spendTokens]', err);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates transactions to record the spending of structured tokens.
|
||||
*/
|
||||
async function spendStructuredTokens(
|
||||
txData: SpendTxData,
|
||||
tokenUsage: {
|
||||
promptTokens?: { input?: number; write?: number; read?: number };
|
||||
completionTokens?: number;
|
||||
},
|
||||
) {
|
||||
const { promptTokens, completionTokens } = tokenUsage;
|
||||
logger.debug(
|
||||
`[spendStructuredTokens] conversationId: ${txData.conversationId}${
|
||||
txData?.context ? ` | Context: ${txData?.context}` : ''
|
||||
} | Token usage: `,
|
||||
{ promptTokens, completionTokens },
|
||||
);
|
||||
let prompt: TransactionResult | undefined, completion: TransactionResult | undefined;
|
||||
try {
|
||||
if (promptTokens) {
|
||||
const input = Math.max(promptTokens.input ?? 0, 0);
|
||||
const write = Math.max(promptTokens.write ?? 0, 0);
|
||||
const read = Math.max(promptTokens.read ?? 0, 0);
|
||||
const totalInputTokens = input + write + read;
|
||||
prompt = await transactionMethods.createStructuredTransaction({
|
||||
...txData,
|
||||
tokenType: 'prompt',
|
||||
inputTokens: -input,
|
||||
writeTokens: -write,
|
||||
readTokens: -read,
|
||||
inputTokenCount: totalInputTokens,
|
||||
});
|
||||
}
|
||||
|
||||
if (completionTokens) {
|
||||
const totalInputTokens = promptTokens
|
||||
? Math.max(promptTokens.input ?? 0, 0) +
|
||||
Math.max(promptTokens.write ?? 0, 0) +
|
||||
Math.max(promptTokens.read ?? 0, 0)
|
||||
: undefined;
|
||||
completion = await transactionMethods.createTransaction({
|
||||
...txData,
|
||||
tokenType: 'completion',
|
||||
rawAmount: -Math.max(completionTokens, 0),
|
||||
inputTokenCount: totalInputTokens,
|
||||
});
|
||||
}
|
||||
|
||||
if (prompt || completion) {
|
||||
logger.debug('[spendStructuredTokens] Transaction data record against balance:', {
|
||||
user: txData.user,
|
||||
prompt: prompt?.prompt,
|
||||
promptRate: prompt?.rate,
|
||||
completion: completion?.completion,
|
||||
completionRate: completion?.rate,
|
||||
balance: completion?.balance ?? prompt?.balance,
|
||||
});
|
||||
} else {
|
||||
logger.debug('[spendStructuredTokens] No transactions incurred against balance');
|
||||
}
|
||||
} catch (err) {
|
||||
logger.error('[spendStructuredTokens]', err);
|
||||
}
|
||||
|
||||
return { prompt, completion };
|
||||
}
|
||||
|
||||
return { spendTokens, spendStructuredTokens };
|
||||
}
|
||||
|
||||
export type SpendTokensMethods = ReturnType<typeof createSpendTokensMethods>;
|
||||
Loading…
Add table
Add a link
Reference in a new issue