🪢 chore: Consolidate Pricing and Tx Imports After tx.js Module Removal (#12086)

* 🧹 chore: resolve imports due to rebase

* chore: Update model mocks in unit tests for consistency

- Consolidated model mock implementations across various test files to streamline setup and reduce redundancy.
- Removed duplicate mock definitions for `getMultiplier` and `getCacheMultiplier`, ensuring a unified approach in `recordCollectedUsage.spec.js`, `openai.spec.js`, `responses.unit.spec.js`, and `abortMiddleware.spec.js`.
- Enhanced clarity and maintainability of test files by aligning mock structures with the latest model updates.

* fix: Safeguard token credit checks in transaction tests

- Updated assertions in `transaction.spec.ts` to handle potential null values for `updatedBalance` by using optional chaining.
- Enhanced robustness of tests related to token credit calculations, ensuring they correctly account for scenarios where the balance may not be found.

* chore: transaction methods with bulk insert functionality

- Introduced `bulkInsertTransactions` method in `transaction.ts` to facilitate batch insertion of transaction documents.
- Updated test file `transactions.bulk-parity.spec.ts` to utilize new pricing function assignments and handle potential null values in calculations, improving test robustness.
- Refactored pricing function initialization for clarity and consistency.

* refactor: Enhance type definitions and introduce new utility functions for model matching

- Added `findMatchingPattern` and `matchModelName` utility functions to improve model name matching logic in transaction methods.
- Updated type definitions for `findMatchingPattern` to accept a more specific tokensMap structure, enhancing type safety.
- Refactored `dbMethods` initialization in `transactions.bulk-parity.spec.ts` to include the new utility functions, improving test clarity and functionality.

* refactor: Update database method imports and enhance transaction handling

- Refactored `abortMiddleware.js` to utilize centralized database methods for message handling and conversation retrieval, improving code consistency.
- Enhanced `bulkInsertTransactions` in `transaction.ts` to handle empty document arrays gracefully and added error logging for better debugging.
- Updated type definitions in `transactions.ts` to enforce stricter typing for token types, enhancing type safety across transaction methods.
- Improved test setup in `transactions.bulk-parity.spec.ts` by refining pricing function assignments and ensuring robust handling of potential null values.

* refactor: Update database method references and improve transaction multiplier handling

- Refactored `client.js` to update database method references for `bulkInsertTransactions` and `updateBalance`, ensuring consistency in method usage.
- Enhanced transaction multiplier calculations in `transaction.spec.ts` to provide fallback values for write and read multipliers, improving robustness in cost calculations across structured token spending tests.

api/server/controllers/agents/__tests__/openai.spec.js

@@ -79,11 +79,6 @@ jest.mock('~/server/services/ToolService', () => ({
 
 const mockGetMultiplier = jest.fn().mockReturnValue(1);
 const mockGetCacheMultiplier = jest.fn().mockReturnValue(null);
-jest.mock('~/models/tx', () => ({
-  getMultiplier: mockGetMultiplier,
-  getCacheMultiplier: mockGetCacheMultiplier,
-}));
-
 
 jest.mock('~/server/controllers/agents/callbacks', () => ({
   createToolEndCallback: jest.fn().mockReturnValue(jest.fn()),
@@ -110,6 +105,8 @@ jest.mock('~/models', () => ({
   bulkInsertTransactions: mockBulkInsertTransactions,
   spendTokens: mockSpendTokens,
   spendStructuredTokens: mockSpendStructuredTokens,
+  getMultiplier: mockGetMultiplier,
+  getCacheMultiplier: mockGetCacheMultiplier,
   getConvoFiles: jest.fn().mockResolvedValue([]),
 }));
 

api/server/controllers/agents/__tests__/responses.unit.spec.js

@@ -103,11 +103,6 @@ jest.mock('~/server/services/ToolService', () => ({
 
 const mockGetMultiplier = jest.fn().mockReturnValue(1);
 const mockGetCacheMultiplier = jest.fn().mockReturnValue(null);
-jest.mock('~/models/tx', () => ({
-  getMultiplier: mockGetMultiplier,
-  getCacheMultiplier: mockGetCacheMultiplier,
-}));
-
 
 jest.mock('~/server/controllers/agents/callbacks', () => ({
   createToolEndCallback: jest.fn().mockReturnValue(jest.fn()),
@@ -136,6 +131,8 @@ jest.mock('~/models', () => ({
   bulkInsertTransactions: mockBulkInsertTransactions,
   spendTokens: mockSpendTokens,
   spendStructuredTokens: mockSpendStructuredTokens,
+  getMultiplier: mockGetMultiplier,
+  getCacheMultiplier: mockGetCacheMultiplier,
   getConvoFiles: jest.fn().mockResolvedValue([]),
   saveConvo: jest.fn().mockResolvedValue({}),
   getConvo: jest.fn().mockResolvedValue(null),

api/server/controllers/agents/client.js

@@ -46,8 +46,6 @@ const {
   removeNullishValues,
 } = require('librechat-data-provider');
 const { encodeAndFormat } = require('~/server/services/Files/images/encode');
-const { updateBalance, bulkInsertTransactions } = require('~/models');
-const { getMultiplier, getCacheMultiplier } = require('~/models/tx');
 const { createContextHandlers } = require('~/app/clients/prompts');
 const { getMCPServerTools } = require('~/server/services/Config');
 const BaseClient = require('~/app/clients/BaseClient');
@@ -631,8 +629,8 @@ class AgentClient extends BaseClient {
       {
         spendTokens: db.spendTokens,
         spendStructuredTokens: db.spendStructuredTokens,
-        pricing: { getMultiplier, getCacheMultiplier },
-        bulkWriteOps: { insertMany: bulkInsertTransactions, updateBalance },
+        pricing: { getMultiplier: db.getMultiplier, getCacheMultiplier: db.getCacheMultiplier },
+        bulkWriteOps: { insertMany: db.bulkInsertTransactions, updateBalance: db.updateBalance },
       },
       {
         user: this.user ?? this.options.req.user?.id,

api/server/controllers/agents/openai.js

@@ -24,7 +24,6 @@ const {
 const { loadAgentTools, loadToolsForExecution } = require('~/server/services/ToolService');
 const { createToolEndCallback } = require('~/server/controllers/agents/callbacks');
 const { findAccessibleResources } = require('~/server/services/PermissionService');
-const { getMultiplier, getCacheMultiplier } = require('~/models/tx');
 const db = require('~/models');
 
 /**
@@ -494,7 +493,7 @@ const OpenAIChatCompletionController = async (req, res) => {
       {
         spendTokens: db.spendTokens,
         spendStructuredTokens: db.spendStructuredTokens,
-        pricing: { getMultiplier, getCacheMultiplier },
+        pricing: { getMultiplier: db.getMultiplier, getCacheMultiplier: db.getCacheMultiplier },
         bulkWriteOps: { insertMany: db.bulkInsertTransactions, updateBalance: db.updateBalance },
       },
       {

api/server/controllers/agents/recordCollectedUsage.spec.js

@@ -21,14 +21,8 @@ const mockRecordCollectedUsage = jest
 jest.mock('~/models', () => ({
   spendTokens: (...args) => mockSpendTokens(...args),
   spendStructuredTokens: (...args) => mockSpendStructuredTokens(...args),
-}));
-
-jest.mock('~/models/tx', () => ({
   getMultiplier: mockGetMultiplier,
   getCacheMultiplier: mockGetCacheMultiplier,
-}));
-
-jest.mock('~/models', () => ({
   updateBalance: mockUpdateBalance,
   bulkInsertTransactions: mockBulkInsertTransactions,
 }));

api/server/controllers/agents/responses.js

@@ -36,7 +36,6 @@ const {
 } = require('~/server/controllers/agents/callbacks');
 const { loadAgentTools, loadToolsForExecution } = require('~/server/services/ToolService');
 const { findAccessibleResources } = require('~/server/services/PermissionService');
-const { getMultiplier, getCacheMultiplier } = require('~/models/tx');
 const db = require('~/models');
 
 /** @type {import('@librechat/api').AppConfig | null} */
@@ -514,7 +513,7 @@ const createResponse = async (req, res) => {
         {
           spendTokens: db.spendTokens,
           spendStructuredTokens: db.spendStructuredTokens,
-          pricing: { getMultiplier, getCacheMultiplier },
+          pricing: { getMultiplier: db.getMultiplier, getCacheMultiplier: db.getCacheMultiplier },
           bulkWriteOps: { insertMany: db.bulkInsertTransactions, updateBalance: db.updateBalance },
         },
         {
@@ -668,7 +667,7 @@ const createResponse = async (req, res) => {
         {
           spendTokens: db.spendTokens,
           spendStructuredTokens: db.spendStructuredTokens,
-          pricing: { getMultiplier, getCacheMultiplier },
+          pricing: { getMultiplier: db.getMultiplier, getCacheMultiplier: db.getCacheMultiplier },
           bulkWriteOps: { insertMany: db.bulkInsertTransactions, updateBalance: db.updateBalance },
         },
         {

api/server/middleware/abortMiddleware.js

@@ -8,13 +8,11 @@ const {
   recordCollectedUsage,
   sanitizeMessageForTransmit,
 } = require('@librechat/api');
-const { isAssistantsEndpoint, ErrorTypes } = require('librechat-data-provider');
-const { saveMessage, getConvo, updateBalance, bulkInsertTransactions } = require('~/models');
 const { truncateText, smartTruncateText } = require('~/app/clients/prompts');
-const { getMultiplier, getCacheMultiplier } = require('~/models/tx');
 const clearPendingReq = require('~/cache/clearPendingReq');
 const { sendError } = require('~/server/middleware/error');
 const { abortRun } = require('./abortRun');
+const db = require('~/models');
 
 /**
  * Spend tokens for all models from collected usage.
@@ -44,10 +42,10 @@ async function spendCollectedUsage({
 
   await recordCollectedUsage(
     {
-      spendTokens,
-      spendStructuredTokens,
-      pricing: { getMultiplier, getCacheMultiplier },
-      bulkWriteOps: { insertMany: bulkInsertTransactions, updateBalance },
+      spendTokens: db.spendTokens,
+      spendStructuredTokens: db.spendStructuredTokens,
+      pricing: { getMultiplier: db.getMultiplier, getCacheMultiplier: db.getCacheMultiplier },
+      bulkWriteOps: { insertMany: db.bulkInsertTransactions, updateBalance: db.updateBalance },
     },
     {
       user: userId,
@@ -123,13 +121,13 @@ async function abortMessage(req, res) {
     });
   } else {
     // Fallback: no collected usage, use text-based token counting for primary model only
-    await spendTokens(
+    await db.spendTokens(
       { ...responseMessage, context: 'incomplete', user: userId },
       { promptTokens, completionTokens },
     );
   }
 
-  await saveMessage(
+  await db.saveMessage(
     {
       userId: req?.user?.id,
       isTemporary: req?.body?.isTemporary,
@@ -140,7 +138,7 @@ async function abortMessage(req, res) {
   );
 
   // Get conversation for title
-  const conversation = await getConvo(userId, conversationId);
+  const conversation = await db.getConvo(userId, conversationId);
 
   const finalEvent = {
     title: conversation && !conversation.title ? null : conversation?.title || 'New Chat',

api/server/middleware/abortMiddleware.spec.js

@@ -20,8 +20,6 @@ const mockRecordCollectedUsage = jest
 const mockGetMultiplier = jest.fn().mockReturnValue(1);
 const mockGetCacheMultiplier = jest.fn().mockReturnValue(null);
 
-
-
 jest.mock('@librechat/data-schemas', () => ({
   logger: {
     debug: jest.fn(),
@@ -65,6 +63,10 @@ jest.mock('~/models', () => ({
   getConvo: jest.fn().mockResolvedValue({ title: 'Test Chat' }),
   updateBalance: mockUpdateBalance,
   bulkInsertTransactions: mockBulkInsertTransactions,
+  spendTokens: (...args) => mockSpendTokens(...args),
+  spendStructuredTokens: (...args) => mockSpendStructuredTokens(...args),
+  getMultiplier: mockGetMultiplier,
+  getCacheMultiplier: mockGetCacheMultiplier,
 }));
 
 jest.mock('./abortRun', () => ({

api/server/services/Files/process.spec.js

@@ -30,11 +30,6 @@ jest.mock('~/server/controllers/assistants/v2', () => ({
   deleteResourceFileId: jest.fn(),
 }));
 
-jest.mock('~/models/Agent', () => ({
-  addAgentResourceFile: jest.fn().mockResolvedValue({}),
-  removeAgentResourceFiles: jest.fn(),
-}));
-
 jest.mock('~/server/controllers/assistants/helpers', () => ({
   getOpenAIClient: jest.fn(),
 }));
@@ -47,6 +42,8 @@ jest.mock('~/models', () => ({
   createFile: jest.fn().mockResolvedValue({ file_id: 'created-file-id' }),
   updateFileUsage: jest.fn(),
   deleteFiles: jest.fn(),
+  addAgentResourceFile: jest.fn().mockResolvedValue({}),
+  removeAgentResourceFiles: jest.fn(),
 }));
 
 jest.mock('~/server/utils/getFileStrategy', () => ({