refactor: swap old ConversationCosts endpoint out for much simpler 'costs' endpoint

/costs just recieves an array of { model, endpoint }  and returns the related prompt and completion token rates, leaving the actual message processing to be done later on rather than on the backend.
This commit is contained in:
Dustin Healy 2025-09-13 12:19:26 -07:00
parent ba8c09b361
commit 0edfecf44a
2 changed files with 25 additions and 91 deletions

View file

@ -11,10 +11,11 @@ const {
} = require('~/models'); } = require('~/models');
const { findAllArtifacts, replaceArtifactContent } = require('~/server/services/Artifacts/update'); const { findAllArtifacts, replaceArtifactContent } = require('~/server/services/Artifacts/update');
const { requireJwtAuth, validateMessageReq } = require('~/server/middleware'); const { requireJwtAuth, validateMessageReq } = require('~/server/middleware');
const { tokenValues, getValueKey, defaultRate } = require('~/models/tx');
const { cleanUpPrimaryKeyValue } = require('~/lib/utils/misc'); const { cleanUpPrimaryKeyValue } = require('~/lib/utils/misc');
const { getConvosQueried } = require('~/models/Conversation'); const { getConvosQueried } = require('~/models/Conversation');
const { Message, Transaction } = require('~/db/models');
const { countTokens } = require('~/server/utils'); const { countTokens } = require('~/server/utils');
const { Message } = require('~/db/models');
const router = express.Router(); const router = express.Router();
router.use(requireJwtAuth); router.use(requireJwtAuth);
@ -160,103 +161,37 @@ router.post('/artifact/:messageId', async (req, res) => {
} }
}); });
router.get('/:conversationId/costs', validateMessageReq, async (req, res) => { /**
* POST /costs
* Get cost information for models in modelHistory array
*/
router.post('/costs', async (req, res) => {
try { try {
const user = req.user.id; const { modelHistory } = req.body;
const { conversationId } = req.params;
const [transactions, messages] = await Promise.all([ if (!Array.isArray(modelHistory)) {
Transaction.find({ return res.status(400).json({ error: 'modelHistory must be an array' });
conversationId,
user,
tokenType: { $in: ['prompt', 'completion'] },
})
.select('tokenType tokenValue createdAt')
.sort({ createdAt: 1 })
.lean(),
Message.find({ conversationId, user })
.select('messageId isCreatedByUser tokenCount createdAt')
.sort({ createdAt: 1 })
.lean(),
]);
const userMsgs = messages.filter((m) => m.isCreatedByUser);
const aiMsgs = messages.filter((m) => !m.isCreatedByUser);
const perMessageMap = new Map();
for (const msg of messages) {
perMessageMap.set(msg.messageId, {
messageId: msg.messageId,
tokenType: msg.isCreatedByUser ? 'prompt' : 'completion',
tokenCount: msg.tokenCount ?? 0,
tokenValue: 0,
usd: 0,
});
} }
let currentPrompt = 0; const modelCostTable = {};
let currentCompletion = 0;
let promptTokenValue = 0; modelHistory.forEach((modelEntry) => {
let completionTokenValue = 0; if (modelEntry && typeof modelEntry === 'object' && modelEntry.model && modelEntry.endpoint) {
const { model, endpoint } = modelEntry;
for (const tx of transactions) { const valueKey = getValueKey(model, endpoint);
const value = Math.abs(tx.tokenValue ?? 0); const pricing = tokenValues[valueKey];
if (tx.tokenType === 'prompt') {
promptTokenValue += value; modelCostTable[model] = {
const target = userMsgs[currentPrompt] ?? userMsgs[userMsgs.length - 1]; prompt: pricing?.prompt ?? defaultRate,
if (target) { completion: pricing?.completion ?? defaultRate,
const entry = perMessageMap.get(target.messageId); };
entry.tokenValue += value;
perMessageMap.set(target.messageId, entry);
if (currentPrompt < userMsgs.length - 1) {
currentPrompt++;
}
}
} else if (tx.tokenType === 'completion') {
completionTokenValue += value;
const target = aiMsgs[currentCompletion] ?? aiMsgs[aiMsgs.length - 1];
if (target) {
const entry = perMessageMap.get(target.messageId);
entry.tokenValue += value;
perMessageMap.set(target.messageId, entry);
if (currentCompletion < aiMsgs.length - 1) {
currentCompletion++;
}
}
} }
} });
const perMessage = Array.from(perMessageMap.values()).map((entry) => ({ res.status(200).json({ modelCostTable });
messageId: entry.messageId,
tokenType: entry.tokenType,
tokenCount: entry.tokenCount,
usd: entry.tokenValue / 1_000_000,
}));
const promptTokenCount = userMsgs.reduce((sum, m) => sum + (m.tokenCount ?? 0), 0);
const completionTokenCount = aiMsgs.reduce((sum, m) => sum + (m.tokenCount ?? 0), 0);
const totalTokenCount = promptTokenCount + completionTokenCount;
const totals = {
prompt: {
usd: promptTokenValue / 1_000_000,
tokenCount: promptTokenCount,
},
completion: {
usd: completionTokenValue / 1_000_000,
tokenCount: completionTokenCount,
},
total: {
usd: (promptTokenValue + completionTokenValue) / 1_000_000,
tokenCount: totalTokenCount,
},
};
const response = { conversationId, totals, perMessage };
res.status(200).json(response);
} catch (error) { } catch (error) {
logger.error('Error fetching conversation costs:', error); logger.error('Error fetching model costs:', error);
res.status(500).json({ error: 'Internal server error' }); res.status(500).json({ error: 'Internal server error' });
} }
}); });

View file

@ -66,8 +66,7 @@ export const messages = (params: q.MessagesListParams) => {
export const messagesArtifacts = (messageId: string) => `${messagesRoot}/artifacts/${messageId}`; export const messagesArtifacts = (messageId: string) => `${messagesRoot}/artifacts/${messageId}`;
export const conversationCosts = (conversationId: string) => export const costs = () => `/api/messages/costs`;
`/api/messages/${conversationId}/costs`;
const shareRoot = `${BASE_URL}/api/share`; const shareRoot = `${BASE_URL}/api/share`;
export const shareMessages = (shareId: string) => `${shareRoot}/${shareId}`; export const shareMessages = (shareId: string) => `${shareRoot}/${shareId}`;