Skip to content

feat: capture token usage from AI providers #877

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

Merged
chubes4 merged 1 commit into from
Mar 19, 2026
Merged

feat: capture token usage from AI providers #877

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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: 48 additions & 7 deletions inc/Api/Chat/ChatOrchestrator.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,24 @@ public static function processChat(
'has_pending_tools' => ! $is_completed,
);

// Accumulate token usage across turns in session metadata.
$turn_usage = $result['usage'] ?? array();
if ( ! empty( $turn_usage ) && ( $turn_usage['total_tokens'] ?? 0 ) > 0 ) {
$existing_session = $chat_db->get_session( $session_id );
$existing_metadata = ! empty( $existing_session['metadata'] ) ? json_decode( $existing_session['metadata'], true ) : array();
$existing_usage = $existing_metadata['token_usage'] ?? array(
'prompt_tokens' => 0,
'completion_tokens' => 0,
'total_tokens' => 0,
);

$metadata['token_usage'] = array(
'prompt_tokens' => (int) $existing_usage['prompt_tokens'] + (int) ( $turn_usage['prompt_tokens'] ?? 0 ),
'completion_tokens' => (int) $existing_usage['completion_tokens'] + (int) ( $turn_usage['completion_tokens'] ?? 0 ),
'total_tokens' => (int) $existing_usage['total_tokens'] + (int) ( $turn_usage['total_tokens'] ?? 0 ),
);
}

if ( $selected_pipeline_id ) {
$metadata['selected_pipeline_id'] = $selected_pipeline_id;
}
Expand Down Expand Up @@ -312,6 +330,21 @@ public static function processContinue( string $session_id, int $user_id ): arra
'has_pending_tools' => ! $is_completed,
);

// Accumulate token usage across continuation turns.
$turn_usage = $result['usage'] ?? array();
$existing_usage = $metadata['token_usage'] ?? array(
'prompt_tokens' => 0,
'completion_tokens' => 0,
'total_tokens' => 0,
);
if ( ! empty( $turn_usage ) && ( $turn_usage['total_tokens'] ?? 0 ) > 0 ) {
$updated_metadata['token_usage'] = array(
'prompt_tokens' => (int) $existing_usage['prompt_tokens'] + (int) ( $turn_usage['prompt_tokens'] ?? 0 ),
'completion_tokens' => (int) $existing_usage['completion_tokens'] + (int) ( $turn_usage['completion_tokens'] ?? 0 ),
'total_tokens' => (int) $existing_usage['total_tokens'] + (int) ( $turn_usage['total_tokens'] ?? 0 ),
);
}

if ( $selected_pipeline_id ) {
$updated_metadata['selected_pipeline_id'] = $selected_pipeline_id;
}
Expand Down Expand Up @@ -400,16 +433,23 @@ public static function processPing( string $message, string $provider, string $m
return $result;
}

// Update session to completed with ping source.
// Update session to completed with ping source and token usage.
$ping_metadata = array(
'status' => 'completed',
'last_activity' => current_time( 'mysql', true ),
'message_count' => count( $result['messages'] ),
'source' => 'ping',
);

$ping_usage = $result['usage'] ?? array();
if ( ! empty( $ping_usage ) && ( $ping_usage['total_tokens'] ?? 0 ) > 0 ) {
$ping_metadata['token_usage'] = $ping_usage;
}

$chat_db->update_session(
$session_id,
$result['messages'],
array(
'status' => 'completed',
'last_activity' => current_time( 'mysql', true ),
'message_count' => count( $result['messages'] ),
'source' => 'ping',
),
$ping_metadata,
$provider,
$model
);
Expand Down Expand Up @@ -621,6 +661,7 @@ public static function executeConversationTurn(
'last_tool_calls' => $loop_result['last_tool_calls'] ?? array(),
'warning' => $loop_result['warning'] ?? null,
'max_turns_reached' => $loop_result['max_turns_reached'] ?? false,
'usage' => $loop_result['usage'] ?? array(),
);
} catch ( \Throwable $e ) {
do_action(
Expand Down
9 changes: 9 additions & 0 deletions inc/Core/Steps/AI/AIStep.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -256,6 +256,15 @@ protected function executeStep(): array {
return array();
}

// Store token usage in job engine_data (read-then-merge — store_engine_data is a full overwrite).
$usage = $loop_result['usage'] ?? array();
if ( ! empty( $usage ) && $this->job_id > 0 && ( $usage['total_tokens'] ?? 0 ) > 0 ) {
$jobs_db = new \DataMachine\Core\Database\Jobs\Jobs();
$existing_data = $engine_data;
$existing_data['token_usage'] = $usage;
$jobs_db->store_engine_data( $this->job_id, $existing_data );
}

// Process loop results into data packets
return self::processLoopResults( $loop_result, $this->dataPackets, $payload, $available_tools );
}
Expand Down
18 changes: 18 additions & 0 deletions inc/Engine/AI/AIConversationLoop.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ class AIConversationLoop {
* @type int $turn_count Number of turns executed
* @type bool $completed Whether loop finished naturally (no tool calls)
* @type array $last_tool_calls Last set of tool calls (if any)
* @type array $usage Accumulated token usage {prompt_tokens, completion_tokens, total_tokens}
* }
*/
public function execute(
Expand All @@ -63,6 +64,13 @@ public function execute(
$last_tool_calls = array();
$tool_execution_results = array();

// Accumulate token usage across all turns.
$total_usage = array(
'prompt_tokens' => 0,
'completion_tokens' => 0,
'total_tokens' => 0,
);

// Track which handler tools have been executed for multi-handler support.
// In pipeline mode, conversation should only complete when ALL configured
// handlers have fired, not just the first one.
Expand Down Expand Up @@ -116,12 +124,21 @@ public function execute(
'completed' => false,
'last_tool_calls' => array(),
'error' => $ai_response['error'] ?? 'AI request failed',
'usage' => $total_usage,
);
}

$tool_calls = $ai_response['data']['tool_calls'] ?? array();
$ai_content = $ai_response['data']['content'] ?? '';

// Accumulate token usage from this turn.
$turn_usage = $ai_response['data']['usage'] ?? array();
if ( ! empty( $turn_usage ) ) {
$total_usage['prompt_tokens'] += (int) ( $turn_usage['prompt_tokens'] ?? 0 );
$total_usage['completion_tokens'] += (int) ( $turn_usage['completion_tokens'] ?? 0 );
$total_usage['total_tokens'] += (int) ( $turn_usage['total_tokens'] ?? 0 );
}

// Store final content from this turn
if ( ! empty( $ai_content ) ) {
$final_content = $ai_content;
Expand Down Expand Up @@ -355,6 +372,7 @@ public function execute(
'last_tool_calls' => $last_tool_calls,
'tool_execution_results' => $tool_execution_results,
'has_pending_tools' => ! empty( $last_tool_calls ) && ! $conversation_complete,
'usage' => $total_usage,
);

if ( $turn_count >= $max_turns && ! $conversation_complete ) {
Expand Down
Loading