refactor: add back getTokenCountForResponse for slightly more accurate mapping of responses token counts (#1067)

This commit is contained in:
Danny Avila 2023-10-17 06:42:58 -04:00 committed by GitHub
parent 6d8aed7ef8
commit 377f2c7c19
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 26 additions and 6 deletions

View file

@ -40,6 +40,12 @@ class BaseClient {
throw new Error('Subclasses attempted to call summarizeMessages without implementing it'); throw new Error('Subclasses attempted to call summarizeMessages without implementing it');
} }
async getTokenCountForResponse(response) {
if (this.options.debug) {
console.debug('`recordTokenUsage` not implemented.', response);
}
}
async recordTokenUsage({ promptTokens, completionTokens }) { async recordTokenUsage({ promptTokens, completionTokens }) {
if (this.options.debug) { if (this.options.debug) {
console.debug('`recordTokenUsage` not implemented.', { promptTokens, completionTokens }); console.debug('`recordTokenUsage` not implemented.', { promptTokens, completionTokens });
@ -455,11 +461,16 @@ class BaseClient {
promptTokens, promptTokens,
}; };
if (tokenCountMap && this.getTokenCount) { if (
responseMessage.tokenCount = this.getTokenCount(completion); tokenCountMap &&
responseMessage.completionTokens = responseMessage.tokenCount; this.recordTokenUsage &&
this.getTokenCountForResponse &&
this.getTokenCount
) {
responseMessage.tokenCount = this.getTokenCountForResponse(responseMessage);
const completionTokens = this.getTokenCount(completion);
await this.recordTokenUsage({ promptTokens, completionTokens });
} }
await this.recordTokenUsage(responseMessage);
await this.saveMessageToDatabase(responseMessage, saveOptions, user); await this.saveMessageToDatabase(responseMessage, saveOptions, user);
delete responseMessage.tokenCount; delete responseMessage.tokenCount;
return responseMessage; return responseMessage;

View file

@ -653,6 +653,13 @@ ${convo}
{ promptTokens, completionTokens }, { promptTokens, completionTokens },
); );
} }
getTokenCountForResponse(response) {
return this.getTokenCountForMessage({
role: 'assistant',
content: response.text,
});
}
} }
module.exports = OpenAIClient; module.exports = OpenAIClient;

View file

@ -230,13 +230,15 @@ If your reverse proxy is compatible to OpenAI specs in every other way, it may s
console.debug('[handleResponseMessage] Output:', { output, errorMessage, ...result }); console.debug('[handleResponseMessage] Output:', { output, errorMessage, ...result });
const { error } = responseMessage; const { error } = responseMessage;
if (!error) { if (!error) {
responseMessage.tokenCount = this.getTokenCount(responseMessage.text); responseMessage.tokenCount = this.getTokenCountForResponse(responseMessage);
responseMessage.completionTokens = responseMessage.tokenCount; responseMessage.completionTokens = this.getTokenCount(responseMessage.text);
} }
// Record usage only when completion is skipped as it is already recorded in the agent phase.
if (!this.agentOptions.skipCompletion && !error) { if (!this.agentOptions.skipCompletion && !error) {
await this.recordTokenUsage(responseMessage); await this.recordTokenUsage(responseMessage);
} }
await this.saveMessageToDatabase(responseMessage, saveOptions, user); await this.saveMessageToDatabase(responseMessage, saveOptions, user);
delete responseMessage.tokenCount; delete responseMessage.tokenCount;
return { ...responseMessage, ...result }; return { ...responseMessage, ...result };