Skip to content

[Platform] Implement token usage aggregation #1051

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

Draft
chr-hertel wants to merge 1 commit into
base: main
Choose a base branch
from
+446 −217
Draft

[Platform] Implement token usage aggregation #1051

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
18 changes: 9 additions & 9 deletions examples/bootstrap.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,15 +107,15 @@ function print_token_usage(Metadata $metadata): void
$table = new Table(output());
$table->setHeaderTitle('Token Usage');
$table->setRows([
['Prompt tokens', $tokenUsage->promptTokens ?? $na],
['Completion tokens', $tokenUsage->completionTokens ?? $na],
['Thinking tokens', $tokenUsage->thinkingTokens ?? $na],
['Tool tokens', $tokenUsage->toolTokens ?? $na],
['Cached tokens', $tokenUsage->cachedTokens ?? $na],
['Remaining tokens minute', $tokenUsage->remainingTokensMinute ?? $na],
['Remaining tokens month', $tokenUsage->remainingTokensMonth ?? $na],
['Remaining tokens', $tokenUsage->remainingTokens ?? $na],
['Utilized tokens', $tokenUsage->totalTokens ?? $na],
['Prompt tokens', $tokenUsage->getPromptTokens() ?? $na],
['Completion tokens', $tokenUsage->getCompletionTokens() ?? $na],
['Thinking tokens', $tokenUsage->getThinkingTokens() ?? $na],
['Tool tokens', $tokenUsage->getToolTokens() ?? $na],
['Cached tokens', $tokenUsage->getCachedTokens() ?? $na],
['Remaining tokens minute', $tokenUsage->getRemainingTokensMinute() ?? $na],
['Remaining tokens month', $tokenUsage->getRemainingTokensMonth() ?? $na],
['Remaining tokens', $tokenUsage->getRemainingTokens() ?? $na],
['Total tokens', $tokenUsage->getTotalTokens() ?? $na],
]);
$table->render();
}
Expand Down
19 changes: 8 additions & 11 deletions src/platform/src/Bridge/Anthropic/TokenOutputProcessor.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,28 +32,25 @@ public function processOutput(Output $output): void
}

$metadata = $output->getResult()->getMetadata();

$tokenUsage = new TokenUsage();

$content = $rawResponse->toArray(false);

if (!\array_key_exists('usage', $content)) {
$metadata->add('token_usage', $tokenUsage);
$metadata->add('token_usage', new TokenUsage());

return;
}

$usage = $content['usage'];

$tokenUsage->promptTokens = $usage['input_tokens'] ?? null;
$tokenUsage->completionTokens = $usage['output_tokens'] ?? null;
$tokenUsage->toolTokens = $usage['server_tool_use']['web_search_requests'] ?? null;

$cachedTokens = null;
if (\array_key_exists('cache_creation_input_tokens', $usage) || \array_key_exists('cache_read_input_tokens', $usage)) {
$cachedTokens = ($usage['cache_creation_input_tokens'] ?? 0) + ($usage['cache_read_input_tokens'] ?? 0);
}
$tokenUsage->cachedTokens = $cachedTokens;

$metadata->add('token_usage', $tokenUsage);
$metadata->add('token_usage', new TokenUsage(
promptTokens: $usage['input_tokens'] ?? null,
completionTokens: $usage['output_tokens'] ?? null,
toolTokens: $usage['server_tool_use']['web_search_requests'] ?? null,
cachedTokens: $cachedTokens,
));
}
}
21 changes: 6 additions & 15 deletions src/platform/src/Bridge/DeepSeek/TokenOutputProcessor.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,23 +35,14 @@ public function processOutput(Output $output): void
}

$metadata = $output->getResult()->getMetadata();

$tokenUsage = new TokenUsage();

$content = $rawResponse->toArray(false);

if (!\array_key_exists('usage', $content)) {
$metadata->add('token_usage', $tokenUsage);

return;
}

$usage = $content['usage'];

$tokenUsage->promptTokens = $usage['prompt_tokens'] ?? null;
$tokenUsage->completionTokens = $usage['completion_tokens'] ?? null;
$tokenUsage->cachedTokens = $usage['prompt_cache_hit_tokens'] ?? null;
$tokenUsage->totalTokens = $usage['total_tokens'] ?? null;
$tokenUsage = new TokenUsage(
promptTokens: $content['usage']['prompt_tokens'] ?? null,
completionTokens: $content['usage']['completion_tokens'] ?? null,
cachedTokens: $content['usage']['prompt_cache_hit_tokens'] ?? null,
totalTokens: $content['usage']['total_tokens'] ?? null,
);

$metadata->add('token_usage', $tokenUsage);
}
Expand Down
24 changes: 8 additions & 16 deletions src/platform/src/Bridge/Gemini/TokenOutputProcessor.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,24 +32,16 @@ public function processOutput(Output $output): void
}

$metadata = $output->getResult()->getMetadata();

$tokenUsage = new TokenUsage();

$content = $rawResponse->toArray(false);
if (!\array_key_exists('usageMetadata', $content)) {
$metadata->add('token_usage', $tokenUsage);

return;
}

$usage = $content['usageMetadata'];

$tokenUsage->promptTokens = $usage['promptTokenCount'] ?? null;
$tokenUsage->completionTokens = $usage['candidatesTokenCount'] ?? null;
$tokenUsage->thinkingTokens = $usage['thoughtsTokenCount'] ?? null;
$tokenUsage->toolTokens = $usage['toolUsePromptTokenCount'] ?? null;
$tokenUsage->cachedTokens = $usage['cachedContentTokenCount'] ?? null;
$tokenUsage->totalTokens = $usage['totalTokenCount'] ?? null;
$tokenUsage = new TokenUsage(
promptTokens: $content['usageMetadata']['promptTokenCount'] ?? null,
completionTokens: $content['usageMetadata']['candidatesTokenCount'] ?? null,
thinkingTokens: $content['usageMetadata']['thoughtsTokenCount'] ?? null,
toolTokens: $content['usageMetadata']['toolUsePromptTokenCount'] ?? null,
cachedTokens: $content['usageMetadata']['cachedContentTokenCount'] ?? null,
totalTokens: $content['usageMetadata']['totalTokenCount'] ?? null,
);

$metadata->add('token_usage', $tokenUsage);
}
Expand Down
20 changes: 6 additions & 14 deletions src/platform/src/Bridge/Mistral/TokenOutputProcessor.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,25 +39,17 @@ public function processOutput(Output $output): void

$remainingTokensMinute = $headers['x-ratelimit-limit-tokens-minute'][0] ?? null;
$remainingTokensMonth = $headers['x-ratelimit-limit-tokens-month'][0] ?? null;

$content = $rawResponse->toArray(false);

$tokenUsage = new TokenUsage(
promptTokens: $content['usage']['prompt_tokens'] ?? null,
completionTokens: $content['usage']['completion_tokens'] ?? null,
remainingTokensMinute: null !== $remainingTokensMinute ? (int) $remainingTokensMinute : null,
remainingTokensMonth: null !== $remainingTokensMonth ? (int) $remainingTokensMonth : null,
totalTokens: $content['usage']['total_tokens'] ?? null,
);

$content = $rawResponse->toArray(false);

if (!\array_key_exists('usage', $content)) {
$metadata->add('token_usage', $tokenUsage);

return;
}

$usage = $content['usage'];

$tokenUsage->promptTokens = $usage['prompt_tokens'] ?? null;
$tokenUsage->completionTokens = $usage['completion_tokens'] ?? null;
$tokenUsage->totalTokens = $usage['total_tokens'] ?? null;

$metadata->add('token_usage', $tokenUsage);
}
}
22 changes: 6 additions & 16 deletions src/platform/src/Bridge/OpenAi/TokenOutputProcessor.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,28 +35,18 @@ public function processOutput(Output $output): void
}

$metadata = $output->getResult()->getMetadata();
$content = $rawResponse->toArray(false);

$remainingTokens = $rawResponse->getHeaders(false)['x-ratelimit-remaining-tokens'][0] ?? null;
$tokenUsage = new TokenUsage(
promptTokens: $content['usage']['prompt_tokens'] ?? null,
completionTokens: $content['usage']['completion_tokens'] ?? null,
thinkingTokens: $content['usage']['completion_tokens_details']['reasoning_tokens'] ?? null,
cachedTokens: $content['usage']['prompt_tokens_details']['cached_tokens'] ?? null,
remainingTokens: null !== $remainingTokens ? (int) $remainingTokens : null,
totalTokens: $content['usage']['total_tokens'] ?? null,
);

$content = $rawResponse->toArray(false);

if (!\array_key_exists('usage', $content)) {
$metadata->add('token_usage', $tokenUsage);

return;
}

$usage = $content['usage'];

$tokenUsage->promptTokens = $usage['prompt_tokens'] ?? null;
$tokenUsage->completionTokens = $usage['completion_tokens'] ?? null;
$tokenUsage->thinkingTokens = $usage['completion_tokens_details']['reasoning_tokens'] ?? null;
$tokenUsage->cachedTokens = $usage['prompt_tokens_details']['cached_tokens'] ?? null;
$tokenUsage->totalTokens = $usage['total_tokens'] ?? null;

$metadata->add('token_usage', $tokenUsage);
}
}
17 changes: 6 additions & 11 deletions src/platform/src/Bridge/Perplexity/TokenOutputProcessor.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,18 +36,13 @@ public function processOutput(Output $output): void

$content = $rawResponse->toArray(false);

if (!\array_key_exists('usage', $content)) {
return;
}

$metadata = $output->getResult()->getMetadata();
$tokenUsage = new TokenUsage();
$usage = $content['usage'];

$tokenUsage->promptTokens = $usage['prompt_tokens'] ?? null;
$tokenUsage->completionTokens = $usage['completion_tokens'] ?? null;
$tokenUsage->thinkingTokens = $usage['reasoning_tokens'] ?? null;
$tokenUsage->totalTokens = $usage['total_tokens'] ?? null;
$tokenUsage = new TokenUsage(
promptTokens: $content['usage']['prompt_tokens'] ?? null,
completionTokens: $content['usage']['completion_tokens'] ?? null,
thinkingTokens: $content['usage']['reasoning_tokens'] ?? null,
totalTokens: $content['usage']['total_tokens'] ?? null,
);

$metadata->add('token_usage', $tokenUsage);
}
Expand Down
28 changes: 10 additions & 18 deletions src/platform/src/Bridge/VertexAi/TokenOutputProcessor.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@

use Symfony\AI\Agent\Output;
use Symfony\AI\Agent\OutputProcessorInterface;
use Symfony\AI\Platform\Metadata\Metadata;
use Symfony\AI\Platform\Metadata\TokenUsage;
use Symfony\AI\Platform\Result\StreamResult;
use Symfony\Contracts\HttpClient\Exception\ClientExceptionInterface;
Expand All @@ -37,7 +36,6 @@ final class TokenOutputProcessor implements OutputProcessorInterface
*/
public function processOutput(Output $output): void
{
$tokenUsage = new TokenUsage();
$metadata = $output->getResult()->getMetadata();

if ($output->getResult() instanceof StreamResult) {
Expand All @@ -51,7 +49,7 @@ public function processOutput(Output $output): void
}

if ($lastChunk) {
$this->extractUsageMetadata($lastChunk['usageMetadata'], $metadata, $tokenUsage);
$metadata->add('token_usage', $this->extractUsageMetadata($lastChunk['usageMetadata']));
}

return;
Expand All @@ -64,13 +62,7 @@ public function processOutput(Output $output): void

$content = $rawResponse->toArray(false);

if (!isset($content['usageMetadata'])) {
$metadata->add('token_usage', $tokenUsage);

return;
}

$this->extractUsageMetadata($content['usageMetadata'], $metadata, $tokenUsage);
$metadata->add('token_usage', $this->extractUsageMetadata($content['usageMetadata'] ?? []));
}

/**
Expand All @@ -82,14 +74,14 @@ public function processOutput(Output $output): void
* totalTokenCount?: int
* } $usage
*/
private function extractUsageMetadata(array $usage, Metadata $metadata, TokenUsage $tokenUsage): void
private function extractUsageMetadata(array $usage): TokenUsage
{
$tokenUsage->promptTokens = $usage['promptTokenCount'] ?? null;
$tokenUsage->completionTokens = $usage['candidatesTokenCount'] ?? null;
$tokenUsage->thinkingTokens = $usage['thoughtsTokenCount'] ?? null;
$tokenUsage->cachedTokens = $usage['cachedContentTokenCount'] ?? null;
$tokenUsage->totalTokens = $usage['totalTokenCount'] ?? null;

$metadata->add('token_usage', $tokenUsage);
return new TokenUsage(
promptTokens: $usage['promptTokenCount'] ?? null,
completionTokens: $usage['candidatesTokenCount'] ?? null,
thinkingTokens: $usage['thoughtsTokenCount'] ?? null,
cachedTokens: $usage['cachedContentTokenCount'] ?? null,
totalTokens: $usage['totalTokenCount'] ?? null,
);
}
}
89 changes: 53 additions & 36 deletions src/platform/src/Metadata/TokenUsage.php
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,46 +14,63 @@
/**
* @author Junaid Farooq <ulislam.junaid125@gmail.com>
*/
final class TokenUsage implements \JsonSerializable
final class TokenUsage implements TokenUsageInterface
{
public function __construct(
public ?int $promptTokens = null,
public ?int $completionTokens = null,
public ?int $thinkingTokens = null,
public ?int $toolTokens = null,
public ?int $cachedTokens = null,
public ?int $remainingTokens = null,
public ?int $remainingTokensMinute = null,
public ?int $remainingTokensMonth = null,
public ?int $totalTokens = null,
private readonly ?int $promptTokens = null,
private readonly ?int $completionTokens = null,
private readonly ?int $thinkingTokens = null,
private readonly ?int $toolTokens = null,
private readonly ?int $cachedTokens = null,
private readonly ?int $remainingTokens = null,
private readonly ?int $remainingTokensMinute = null,
private readonly ?int $remainingTokensMonth = null,
private readonly ?int $totalTokens = null,
) {
}

/**
* @return array{
* prompt_tokens: ?int,
* completion_tokens: ?int,
* thinking_tokens: ?int,
* tool_tokens: ?int,
* cached_tokens: ?int,
* remaining_tokens: ?int,
* remaining_tokens_minute: ?int,
* remaining_tokens_month: ?int,
* total_tokens: ?int,
* }
*/
public function jsonSerialize(): array
{
return [
'prompt_tokens' => $this->promptTokens,
'completion_tokens' => $this->completionTokens,
'thinking_tokens' => $this->thinkingTokens,
'tool_tokens' => $this->toolTokens,
'cached_tokens' => $this->cachedTokens,
'remaining_tokens' => $this->remainingTokens,
'remaining_tokens_minute' => $this->remainingTokensMinute,
'remaining_tokens_month' => $this->remainingTokensMonth,
'total_tokens' => $this->totalTokens,
];
public function getPromptTokens(): ?int
{
return $this->promptTokens;
}

public function getCompletionTokens(): ?int
{
return $this->completionTokens;
}

public function getThinkingTokens(): ?int
{
return $this->thinkingTokens;
}

public function getToolTokens(): ?int
{
return $this->toolTokens;
}

public function getCachedTokens(): ?int
{
return $this->cachedTokens;
}

public function getRemainingTokens(): ?int
{
return $this->remainingTokens;
}

public function getRemainingTokensMinute(): ?int
{
return $this->remainingTokensMinute;
}

public function getRemainingTokensMonth(): ?int
{
return $this->remainingTokensMonth;
}

public function getTotalTokens(): ?int
{
return $this->totalTokens;
}
}
Loading