Skip to content

Commit e694758

Browse files
chr-hertelchr-hertel
committed
Implement token usage aggregation
1 parent 917cab3 commit e694758

18 files changed

+405
-202
lines changed

‎examples/bootstrap.php‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -107,15 +107,15 @@ function print_token_usage(Metadata $metadata): void
107107
$table = new Table(output());
108108
$table->setHeaderTitle('Token Usage');
109109
$table->setRows([
110-
['Prompt tokens', $tokenUsage->promptTokens ?? $na],
111-
['Completion tokens', $tokenUsage->completionTokens ?? $na],
112-
['Thinking tokens', $tokenUsage->thinkingTokens ?? $na],
113-
['Tool tokens', $tokenUsage->toolTokens ?? $na],
114-
['Cached tokens', $tokenUsage->cachedTokens ?? $na],
115-
['Remaining tokens minute', $tokenUsage->remainingTokensMinute ?? $na],
116-
['Remaining tokens month', $tokenUsage->remainingTokensMonth ?? $na],
117-
['Remaining tokens', $tokenUsage->remainingTokens ?? $na],
118-
['Utilized tokens', $tokenUsage->totalTokens ?? $na],
110+
['Prompt tokens', $tokenUsage->getPromptTokens() ?? $na],
111+
['Completion tokens', $tokenUsage->getCompletionTokens() ?? $na],
112+
['Thinking tokens', $tokenUsage->getThinkingTokens() ?? $na],
113+
['Tool tokens', $tokenUsage->getToolTokens() ?? $na],
114+
['Cached tokens', $tokenUsage->getCachedTokens() ?? $na],
115+
['Remaining tokens minute', $tokenUsage->getRemainingTokensMinute() ?? $na],
116+
['Remaining tokens month', $tokenUsage->getRemainingTokensMonth() ?? $na],
117+
['Remaining tokens', $tokenUsage->getRemainingTokens() ?? $na],
118+
['Total tokens', $tokenUsage->getTotalTokens() ?? $na],
119119
]);
120120
$table->render();
121121
}

‎src/platform/src/Bridge/Anthropic/TokenOutputProcessor.php‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -32,28 +32,25 @@ public function processOutput(Output $output): void
3232
}
3333

3434
$metadata = $output->getResult()->getMetadata();
35-
36-
$tokenUsage = new TokenUsage();
37-
3835
$content = $rawResponse->toArray(false);
36+
3937
if (!\array_key_exists('usage', $content)) {
40-
$metadata->add('token_usage', $tokenUsage);
38+
$metadata->add('token_usage', new TokenUsage());
4139

4240
return;
4341
}
4442

4543
$usage = $content['usage'];
46-
47-
$tokenUsage->promptTokens = $usage['input_tokens'] ?? null;
48-
$tokenUsage->completionTokens = $usage['output_tokens'] ?? null;
49-
$tokenUsage->toolTokens = $usage['server_tool_use']['web_search_requests'] ?? null;
50-
5144
$cachedTokens = null;
5245
if (\array_key_exists('cache_creation_input_tokens', $usage) || \array_key_exists('cache_read_input_tokens', $usage)) {
5346
$cachedTokens = ($usage['cache_creation_input_tokens'] ?? 0) + ($usage['cache_read_input_tokens'] ?? 0);
5447
}
55-
$tokenUsage->cachedTokens = $cachedTokens;
5648

57-
$metadata->add('token_usage', $tokenUsage);
49+
$metadata->add('token_usage', new TokenUsage(
50+
promptTokens: $usage['input_tokens'] ?? null,
51+
completionTokens: $usage['output_tokens'] ?? null,
52+
toolTokens: $usage['server_tool_use']['web_search_requests'] ?? null,
53+
cachedTokens: $cachedTokens,
54+
));
5855
}
5956
}

‎src/platform/src/Bridge/Gemini/TokenOutputProcessor.php‎

Lines changed: 8 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -32,24 +32,16 @@ public function processOutput(Output $output): void
3232
}
3333

3434
$metadata = $output->getResult()->getMetadata();
35-
36-
$tokenUsage = new TokenUsage();
37-
3835
$content = $rawResponse->toArray(false);
39-
if (!\array_key_exists('usageMetadata', $content)) {
40-
$metadata->add('token_usage', $tokenUsage);
41-
42-
return;
43-
}
44-
45-
$usage = $content['usageMetadata'];
4636

47-
$tokenUsage->promptTokens = $usage['promptTokenCount'] ?? null;
48-
$tokenUsage->completionTokens = $usage['candidatesTokenCount'] ?? null;
49-
$tokenUsage->thinkingTokens = $usage['thoughtsTokenCount'] ?? null;
50-
$tokenUsage->toolTokens = $usage['toolUsePromptTokenCount'] ?? null;
51-
$tokenUsage->cachedTokens = $usage['cachedContentTokenCount'] ?? null;
52-
$tokenUsage->totalTokens = $usage['totalTokenCount'] ?? null;
37+
$tokenUsage = new TokenUsage(
38+
promptTokens: $content['usageMetadata']['promptTokenCount'] ?? null,
39+
completionTokens: $content['usageMetadata']['candidatesTokenCount'] ?? null,
40+
thinkingTokens: $content['usageMetadata']['thoughtsTokenCount'] ?? null,
41+
toolTokens: $content['usageMetadata']['toolUsePromptTokenCount'] ?? null,
42+
cachedTokens: $content['usageMetadata']['cachedContentTokenCount'] ?? null,
43+
totalTokens: $content['usageMetadata']['totalTokenCount'] ?? null,
44+
);
5345

5446
$metadata->add('token_usage', $tokenUsage);
5547
}

‎src/platform/src/Bridge/Mistral/TokenOutputProcessor.php‎

Lines changed: 6 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -39,25 +39,17 @@ public function processOutput(Output $output): void
3939

4040
$remainingTokensMinute = $headers['x-ratelimit-limit-tokens-minute'][0] ?? null;
4141
$remainingTokensMonth = $headers['x-ratelimit-limit-tokens-month'][0] ?? null;
42+
43+
$content = $rawResponse->toArray(false);
44+
4245
$tokenUsage = new TokenUsage(
46+
promptTokens: $content['usage']['prompt_tokens'] ?? null,
47+
completionTokens: $content['usage']['completion_tokens'] ?? null,
4348
remainingTokensMinute: null !== $remainingTokensMinute ? (int) $remainingTokensMinute : null,
4449
remainingTokensMonth: null !== $remainingTokensMonth ? (int) $remainingTokensMonth : null,
50+
totalTokens: $content['usage']['total_tokens'] ?? null,
4551
);
4652

47-
$content = $rawResponse->toArray(false);
48-
49-
if (!\array_key_exists('usage', $content)) {
50-
$metadata->add('token_usage', $tokenUsage);
51-
52-
return;
53-
}
54-
55-
$usage = $content['usage'];
56-
57-
$tokenUsage->promptTokens = $usage['prompt_tokens'] ?? null;
58-
$tokenUsage->completionTokens = $usage['completion_tokens'] ?? null;
59-
$tokenUsage->totalTokens = $usage['total_tokens'] ?? null;
60-
6153
$metadata->add('token_usage', $tokenUsage);
6254
}
6355
}

‎src/platform/src/Bridge/OpenAi/TokenOutputProcessor.php‎

Lines changed: 6 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -35,28 +35,18 @@ public function processOutput(Output $output): void
3535
}
3636

3737
$metadata = $output->getResult()->getMetadata();
38+
$content = $rawResponse->toArray(false);
3839

3940
$remainingTokens = $rawResponse->getHeaders(false)['x-ratelimit-remaining-tokens'][0] ?? null;
4041
$tokenUsage = new TokenUsage(
42+
promptTokens: $content['usage']['prompt_tokens'] ?? null,
43+
completionTokens: $content['usage']['completion_tokens'] ?? null,
44+
thinkingTokens: $content['usage']['completion_tokens_details']['reasoning_tokens'] ?? null,
45+
cachedTokens: $content['usage']['prompt_tokens_details']['cached_tokens'] ?? null,
4146
remainingTokens: null !== $remainingTokens ? (int) $remainingTokens : null,
47+
totalTokens: $content['usage']['total_tokens'] ?? null,
4248
);
4349

44-
$content = $rawResponse->toArray(false);
45-
46-
if (!\array_key_exists('usage', $content)) {
47-
$metadata->add('token_usage', $tokenUsage);
48-
49-
return;
50-
}
51-
52-
$usage = $content['usage'];
53-
54-
$tokenUsage->promptTokens = $usage['prompt_tokens'] ?? null;
55-
$tokenUsage->completionTokens = $usage['completion_tokens'] ?? null;
56-
$tokenUsage->thinkingTokens = $usage['completion_tokens_details']['reasoning_tokens'] ?? null;
57-
$tokenUsage->cachedTokens = $usage['prompt_tokens_details']['cached_tokens'] ?? null;
58-
$tokenUsage->totalTokens = $usage['total_tokens'] ?? null;
59-
6050
$metadata->add('token_usage', $tokenUsage);
6151
}
6252
}

‎src/platform/src/Bridge/Perplexity/TokenOutputProcessor.php‎

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -36,18 +36,13 @@ public function processOutput(Output $output): void
3636

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

39-
if (!\array_key_exists('usage', $content)) {
40-
return;
41-
}
42-
4339
$metadata = $output->getResult()->getMetadata();
44-
$tokenUsage = new TokenUsage();
45-
$usage = $content['usage'];
46-
47-
$tokenUsage->promptTokens = $usage['prompt_tokens'] ?? null;
48-
$tokenUsage->completionTokens = $usage['completion_tokens'] ?? null;
49-
$tokenUsage->thinkingTokens = $usage['reasoning_tokens'] ?? null;
50-
$tokenUsage->totalTokens = $usage['total_tokens'] ?? null;
40+
$tokenUsage = new TokenUsage(
41+
promptTokens: $content['usage']['prompt_tokens'] ?? null,
42+
completionTokens: $content['usage']['completion_tokens'] ?? null,
43+
thinkingTokens: $content['usage']['reasoning_tokens'] ?? null,
44+
totalTokens: $content['usage']['total_tokens'] ?? null,
45+
);
5146

5247
$metadata->add('token_usage', $tokenUsage);
5348
}

‎src/platform/src/Bridge/VertexAi/TokenOutputProcessor.php‎

Lines changed: 10 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313

1414
use Symfony\AI\Agent\Output;
1515
use Symfony\AI\Agent\OutputProcessorInterface;
16-
use Symfony\AI\Platform\Metadata\Metadata;
1716
use Symfony\AI\Platform\Metadata\TokenUsage;
1817
use Symfony\AI\Platform\Result\StreamResult;
1918
use Symfony\Contracts\HttpClient\Exception\ClientExceptionInterface;
@@ -37,7 +36,6 @@ final class TokenOutputProcessor implements OutputProcessorInterface
3736
*/
3837
public function processOutput(Output $output): void
3938
{
40-
$tokenUsage = new TokenUsage();
4139
$metadata = $output->getResult()->getMetadata();
4240

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

5351
if ($lastChunk) {
54-
$this->extractUsageMetadata($lastChunk['usageMetadata'], $metadata, $tokenUsage);
52+
$metadata->add('token_usage', $this->extractUsageMetadata($lastChunk['usageMetadata']));
5553
}
5654

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

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

67-
if (!isset($content['usageMetadata'])) {
68-
$metadata->add('token_usage', $tokenUsage);
69-
70-
return;
71-
}
72-
73-
$this->extractUsageMetadata($content['usageMetadata'], $metadata, $tokenUsage);
65+
$metadata->add('token_usage', $this->extractUsageMetadata($content['usageMetadata'] ?? []));
7466
}
7567

7668
/**
@@ -82,14 +74,14 @@ public function processOutput(Output $output): void
8274
* totalTokenCount?: int
8375
* } $usage
8476
*/
85-
private function extractUsageMetadata(array $usage, Metadata $metadata, TokenUsage $tokenUsage): void
77+
private function extractUsageMetadata(array $usage): TokenUsage
8678
{
87-
$tokenUsage->promptTokens = $usage['promptTokenCount'] ?? null;
88-
$tokenUsage->completionTokens = $usage['candidatesTokenCount'] ?? null;
89-
$tokenUsage->thinkingTokens = $usage['thoughtsTokenCount'] ?? null;
90-
$tokenUsage->cachedTokens = $usage['cachedContentTokenCount'] ?? null;
91-
$tokenUsage->totalTokens = $usage['totalTokenCount'] ?? null;
92-
93-
$metadata->add('token_usage', $tokenUsage);
79+
return new TokenUsage(
80+
promptTokens: $usage['promptTokenCount'] ?? null,
81+
completionTokens: $usage['candidatesTokenCount'] ?? null,
82+
thinkingTokens: $usage['thoughtsTokenCount'] ?? null,
83+
cachedTokens: $usage['cachedContentTokenCount'] ?? null,
84+
totalTokens: $usage['totalTokenCount'] ?? null,
85+
);
9486
}
9587
}

‎src/platform/src/Metadata/TokenUsage.php‎

Lines changed: 53 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -14,46 +14,63 @@
1414
/**
1515
* @author Junaid Farooq <ulislam.junaid125@gmail.com>
1616
*/
17-
final class TokenUsage implements \JsonSerializable
17+
final class TokenUsage implements TokenUsageInterface
1818
{
1919
public function __construct(
20-
public ?int $promptTokens = null,
21-
public ?int $completionTokens = null,
22-
public ?int $thinkingTokens = null,
23-
public ?int $toolTokens = null,
24-
public ?int $cachedTokens = null,
25-
public ?int $remainingTokens = null,
26-
public ?int $remainingTokensMinute = null,
27-
public ?int $remainingTokensMonth = null,
28-
public ?int $totalTokens = null,
20+
private readonly ?int $promptTokens = null,
21+
private readonly ?int $completionTokens = null,
22+
private readonly ?int $thinkingTokens = null,
23+
private readonly ?int $toolTokens = null,
24+
private readonly ?int $cachedTokens = null,
25+
private readonly ?int $remainingTokens = null,
26+
private readonly ?int $remainingTokensMinute = null,
27+
private readonly ?int $remainingTokensMonth = null,
28+
private readonly ?int $totalTokens = null,
2929
) {
3030
}
3131

32-
/**
33-
* @return array{
34-
* prompt_tokens: ?int,
35-
* completion_tokens: ?int,
36-
* thinking_tokens: ?int,
37-
* tool_tokens: ?int,
38-
* cached_tokens: ?int,
39-
* remaining_tokens: ?int,
40-
* remaining_tokens_minute: ?int,
41-
* remaining_tokens_month: ?int,
42-
* total_tokens: ?int,
43-
* }
44-
*/
45-
public function jsonSerialize(): array
46-
{
47-
return [
48-
'prompt_tokens' => $this->promptTokens,
49-
'completion_tokens' => $this->completionTokens,
50-
'thinking_tokens' => $this->thinkingTokens,
51-
'tool_tokens' => $this->toolTokens,
52-
'cached_tokens' => $this->cachedTokens,
53-
'remaining_tokens' => $this->remainingTokens,
54-
'remaining_tokens_minute' => $this->remainingTokensMinute,
55-
'remaining_tokens_month' => $this->remainingTokensMonth,
56-
'total_tokens' => $this->totalTokens,
57-
];
32+
public function getPromptTokens(): ?int
33+
{
34+
return $this->promptTokens;
35+
}
36+
37+
public function getCompletionTokens(): ?int
38+
{
39+
return $this->completionTokens;
40+
}
41+
42+
public function getThinkingTokens(): ?int
43+
{
44+
return $this->thinkingTokens;
45+
}
46+
47+
public function getToolTokens(): ?int
48+
{
49+
return $this->toolTokens;
50+
}
51+
52+
public function getCachedTokens(): ?int
53+
{
54+
return $this->cachedTokens;
55+
}
56+
57+
public function getRemainingTokens(): ?int
58+
{
59+
return $this->remainingTokens;
60+
}
61+
62+
public function getRemainingTokensMinute(): ?int
63+
{
64+
return $this->remainingTokensMinute;
65+
}
66+
67+
public function getRemainingTokensMonth(): ?int
68+
{
69+
return $this->remainingTokensMonth;
70+
}
71+
72+
public function getTotalTokens(): ?int
73+
{
74+
return $this->totalTokens;
5875
}
5976
}

0 commit comments

Comments
 (0)