🪙 fix: Use Fallback Token Transaction if No Collected Usage (#8503)

This commit is contained in:
Danny Avila 2025-07-16 17:58:15 -04:00 committed by GitHub
parent 1dabe96404
commit 35d8ef50f4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 56 additions and 5 deletions

View file

@ -108,12 +108,15 @@ class BaseClient {
/** /**
* Abstract method to record token usage. Subclasses must implement this method. * Abstract method to record token usage. Subclasses must implement this method.
* If a correction to the token usage is needed, the method should return an object with the corrected token counts. * If a correction to the token usage is needed, the method should return an object with the corrected token counts.
* Should only be used if `recordCollectedUsage` was not used instead.
* @param {string} [model]
* @param {number} promptTokens * @param {number} promptTokens
* @param {number} completionTokens * @param {number} completionTokens
* @returns {Promise<void>} * @returns {Promise<void>}
*/ */
async recordTokenUsage({ promptTokens, completionTokens }) { async recordTokenUsage({ model, promptTokens, completionTokens }) {
logger.debug('[BaseClient] `recordTokenUsage` not implemented.', { logger.debug('[BaseClient] `recordTokenUsage` not implemented.', {
model,
promptTokens, promptTokens,
completionTokens, completionTokens,
}); });
@ -741,9 +744,13 @@ class BaseClient {
} else { } else {
responseMessage.tokenCount = this.getTokenCountForResponse(responseMessage); responseMessage.tokenCount = this.getTokenCountForResponse(responseMessage);
completionTokens = responseMessage.tokenCount; completionTokens = responseMessage.tokenCount;
await this.recordTokenUsage({
usage,
promptTokens,
completionTokens,
model: responseMessage.model,
});
} }
await this.recordTokenUsage({ promptTokens, completionTokens, usage });
} }
if (userMessagePromise) { if (userMessagePromise) {

View file

@ -1135,8 +1135,52 @@ class AgentClient extends BaseClient {
} }
} }
/** Silent method, as `recordCollectedUsage` is used instead */ /**
async recordTokenUsage() {} * @param {object} params
* @param {number} params.promptTokens
* @param {number} params.completionTokens
* @param {OpenAIUsageMetadata} [params.usage]
* @param {string} [params.model]
* @param {string} [params.context='message']
* @returns {Promise<void>}
*/
async recordTokenUsage({ model, promptTokens, completionTokens, usage, context = 'message' }) {
try {
await spendTokens(
{
model,
context,
conversationId: this.conversationId,
user: this.user ?? this.options.req.user?.id,
endpointTokenConfig: this.options.endpointTokenConfig,
},
{ promptTokens, completionTokens },
);
if (
usage &&
typeof usage === 'object' &&
'reasoning_tokens' in usage &&
typeof usage.reasoning_tokens === 'number'
) {
await spendTokens(
{
model,
context: 'reasoning',
conversationId: this.conversationId,
user: this.user ?? this.options.req.user?.id,
endpointTokenConfig: this.options.endpointTokenConfig,
},
{ completionTokens: usage.reasoning_tokens },
);
}
} catch (error) {
logger.error(
'[api/server/controllers/agents/client.js #recordTokenUsage] Error recording token usage',
error,
);
}
}
getEncoding() { getEncoding() {
return 'o200k_base'; return 'o200k_base';