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');
}
async getTokenCountForResponse(response) {
if (this.options.debug) {
console.debug('`recordTokenUsage` not implemented.', response);
}
}
async recordTokenUsage({ promptTokens, completionTokens }) {
if (this.options.debug) {
console.debug('`recordTokenUsage` not implemented.', { promptTokens, completionTokens });
@ -455,11 +461,16 @@ class BaseClient {
promptTokens,
};
if (tokenCountMap && this.getTokenCount) {
responseMessage.tokenCount = this.getTokenCount(completion);
responseMessage.completionTokens = responseMessage.tokenCount;
if (
tokenCountMap &&
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);
delete responseMessage.tokenCount;
return responseMessage;