Skip to content

Expose usage counts in OpenAI streamed responses (Fixes #2003) #2016

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Expose usage counts in OpenAI streamed responses (Fixes #2003) #2016

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
6b46184
Update AI.js
Nihhaar0002 Nov 23, 2025
e4482b6
Update GeminiImageGenerationService.js
Nihhaar0002 Nov 23, 2025
331d25d
Add streaming usage totals to AI chunks
Nihhaar0002 Nov 24, 2025
044c702
Merge pull request #1 from Nihhaar0002/codex/implement-token-usage-tr…
Nihhaar0002 Nov 24, 2025
eb900be
Enable usage in OpenAI streamed chat completions
Nov 24, 2025
44d31c2
Enable usage in OpenAI streamed chat completions
Nihhaar0002 Nov 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 51 additions & 4 deletions src/backend/src/modules/puterai/ClaudeService.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -262,6 +262,50 @@ class ClaudeService extends BaseService {
const init_chat_stream = async ({ chatStream }) => {
const completion = await anthropic.messages.stream(sdk_params);
const usageSum = {};
const runningUsage = {
Copy link
Collaborator

@Salazareo Salazareo Nov 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
const runningUsage = {
const runningUsage = this.usageFormatterUtil({});

should match actual claude usages to make it more visible

input_tokens: 0,
output_tokens: 0,
total_tokens: 0,
};

// Each emitted content block now carries an incremental usage object
// ({ input_tokens, output_tokens, total_tokens }) for live metering.
const getUsage = () => ({
Copy link
Collaborator

@Salazareo Salazareo Nov 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can just spread op the data to copy it

{...runningUsage}

input_tokens: runningUsage.input_tokens,
output_tokens: runningUsage.output_tokens,
total_tokens: runningUsage.total_tokens,
});

const enhanceTextBlock = (block) => {
block.addText = (text) => {
const payload = {
type: 'text',
text,
usage: getUsage(),
Copy link
Collaborator

@Salazareo Salazareo Nov 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
usage: getUsage(),
usage: { ...runningUsage },

};
block.chatStream.stream.write(JSON.stringify(payload) + '\n');
};
return block;
};

const enhanceToolBlock = (block) => {
const originalAddPartialJSON = block.addPartialJSON?.bind(block);
if ( originalAddPartialJSON ) {
block.addPartialJSON = (partial_json) => originalAddPartialJSON(partial_json);
}
block.end = () => {
const buffer = (block.buffer || '').trim() === '' ? '{}' : block.buffer;
const payload = {
...block.contentBlock,
input: JSON.parse(buffer),
...(block.contentBlock?.text ? {} : { text: '' }),
type: 'tool_use',
Copy link
Collaborator

@Salazareo Salazareo Nov 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This needs to go at top of the block, as stream block.contentBlock might override it, to match existing method

usage: getUsage(),
Copy link
Collaborator

@Salazareo Salazareo Nov 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
usage: getUsage(),
usage: { ...runningUsage },

};
block.chatStream.stream.write(JSON.stringify(payload) + '\n');
};
return block;
};

let message, contentBlock;
for await ( const event of completion ) {
Expand All @@ -272,6 +316,9 @@ class ClaudeService extends BaseService {
if ( ! usageSum[key] ) usageSum[key] = 0;
usageSum[key] += meteredData[key];
});
runningUsage.input_tokens += meteredData.input_tokens || 0;
Copy link
Collaborator

@Salazareo Salazareo Nov 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
runningUsage.input_tokens += meteredData.input_tokens || 0;
for ( const usageType in runningUsage ) {
runningUsage[usageType] += meteredData[usageType];
}

runningUsage.output_tokens += meteredData.output_tokens || 0;
runningUsage.total_tokens = runningUsage.input_tokens + runningUsage.output_tokens;

if ( event.type === 'message_start' ) {
message = chatStream.message();
Expand All @@ -285,16 +332,16 @@ class ClaudeService extends BaseService {

if ( event.type === 'content_block_start' ) {
if ( event.content_block.type === 'tool_use' ) {
contentBlock = message.contentBlock({
contentBlock = enhanceToolBlock(message.contentBlock({
type: event.content_block.type,
id: event.content_block.id,
name: event.content_block.name,
});
}));
continue;
}
contentBlock = message.contentBlock({
contentBlock = enhanceTextBlock(message.contentBlock({
type: event.content_block.type,
});
}));
continue;
}

Expand Down