🛠️ fix: getStreamUsage Method in OpenAIClient (#4133)

This commit is contained in:
Danny Avila 2024-09-19 18:07:49 -04:00 committed by GitHub
parent 5f28682314
commit 3ea2d908e0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 69 additions and 1 deletions

View file

@ -923,7 +923,9 @@ ${convo}
this.usage &&
typeof this.usage === 'object' &&
'completion_tokens_details' in this.usage &&
typeof this.usage.completion_tokens_details === 'object'
this.usage.completion_tokens_details &&
typeof this.usage.completion_tokens_details === 'object' &&
'reasoning_tokens' in this.usage.completion_tokens_details
) {
const outputTokens = Math.abs(
this.usage.completion_tokens_details.reasoning_tokens - this.usage[this.outputTokensKey],

View file

@ -701,4 +701,70 @@ describe('OpenAIClient', () => {
expect(client.modelOptions.stop).toBeUndefined();
});
});
describe('getStreamUsage', () => {
it('should return this.usage when completion_tokens_details is null', () => {
const client = new OpenAIClient('test-api-key', defaultOptions);
client.usage = {
completion_tokens_details: null,
prompt_tokens: 10,
completion_tokens: 20,
};
client.inputTokensKey = 'prompt_tokens';
client.outputTokensKey = 'completion_tokens';
const result = client.getStreamUsage();
expect(result).toEqual(client.usage);
});
it('should return this.usage when completion_tokens_details is missing reasoning_tokens', () => {
const client = new OpenAIClient('test-api-key', defaultOptions);
client.usage = {
completion_tokens_details: {
other_tokens: 5,
},
prompt_tokens: 10,
completion_tokens: 20,
};
client.inputTokensKey = 'prompt_tokens';
client.outputTokensKey = 'completion_tokens';
const result = client.getStreamUsage();
expect(result).toEqual(client.usage);
});
it('should calculate output tokens correctly when completion_tokens_details is present with reasoning_tokens', () => {
const client = new OpenAIClient('test-api-key', defaultOptions);
client.usage = {
completion_tokens_details: {
reasoning_tokens: 30,
other_tokens: 5,
},
prompt_tokens: 10,
completion_tokens: 20,
};
client.inputTokensKey = 'prompt_tokens';
client.outputTokensKey = 'completion_tokens';
const result = client.getStreamUsage();
expect(result).toEqual({
reasoning_tokens: 30,
other_tokens: 5,
prompt_tokens: 10,
completion_tokens: 10, // |30 - 20| = 10
});
});
it('should return this.usage when it is undefined', () => {
const client = new OpenAIClient('test-api-key', defaultOptions);
client.usage = undefined;
const result = client.getStreamUsage();
expect(result).toBeUndefined();
});
});
});