Skip to content

Commit e33ba4a

Browse files
junaidbinfarooqjunaidbinfarooq
committed
feat(platform): Standardize token usage
- Makes token usage explicit by adding a dedicated dto for token information - Populates the token usage dto inside different token output processors and then adds it to metadata object
1 parent 9b73ba6 commit e33ba4a

File tree

6 files changed

+145
-48
lines changed

6 files changed

+145
-48
lines changed

‎examples/openai/token-metadata.php

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
use Symfony\AI\Platform\Bridge\OpenAi\TokenOutputProcessor;
1616
use Symfony\AI\Platform\Message\Message;
1717
use Symfony\AI\Platform\Message\MessageBag;
18+
use Symfony\AI\Platform\Result\Metadata\TokenUsage;
1819

1920
require_once dirname(__DIR__).'/bootstrap.php';
2021

@@ -33,8 +34,11 @@
3334
]);
3435

3536
$metadata = $result->getMetadata();
37+
$tokenUsage = $metadata->get('token_usage');
3638

37-
echo 'Utilized Tokens: '.$metadata['total_tokens'].\PHP_EOL;
38-
echo '-- Prompt Tokens: '.$metadata['prompt_tokens'].\PHP_EOL;
39-
echo '-- Completion Tokens: '.$metadata['completion_tokens'].\PHP_EOL;
40-
echo 'Remaining Tokens: '.$metadata['remaining_tokens'].\PHP_EOL;
39+
assert($tokenUsage instanceof TokenUsage);
40+
41+
echo 'Utilized Tokens: '.$tokenUsage->totalTokens.\PHP_EOL;
42+
echo '-- Prompt Tokens: '.$tokenUsage->promptTokens.\PHP_EOL;
43+
echo '-- Completion Tokens: '.$tokenUsage->completionTokens.\PHP_EOL;
44+
echo 'Remaining Tokens: '.$tokenUsage->remainingTokens.\PHP_EOL;

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

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

1414
use Symfony\AI\Agent\Output;
1515
use Symfony\AI\Agent\OutputProcessorInterface;
16+
use Symfony\AI\Platform\Result\Metadata\TokenUsage;
1617
use Symfony\AI\Platform\Result\StreamResult;
1718
use Symfony\Contracts\HttpClient\ResponseInterface;
1819

@@ -34,25 +35,29 @@ public function processOutput(Output $output): void
3435
}
3536

3637
$metadata = $output->result->getMetadata();
38+
$headers = $rawResponse->getHeaders(false);
3739

38-
$metadata->add(
39-
'remaining_tokens_minute',
40-
(int) $rawResponse->getHeaders(false)['x-ratelimit-limit-tokens-minute'][0],
41-
);
42-
43-
$metadata->add(
44-
'remaining_tokens_month',
45-
(int) $rawResponse->getHeaders(false)['x-ratelimit-limit-tokens-month'][0],
40+
$remainingTokensMinute = $headers['x-ratelimit-limit-tokens-minute'][0] ?? null;
41+
$remainingTokensMonth = $headers['x-ratelimit-limit-tokens-month'][0] ?? null;
42+
$tokenUsage = new TokenUsage(
43+
remainingTokensMinute: null !== $remainingTokensMinute ? (int) $remainingTokensMinute : null,
44+
remainingTokensMonth: null !== $remainingTokensMonth ? (int) $remainingTokensMonth : null,
4645
);
4746

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

5049
if (!\array_key_exists('usage', $content)) {
50+
$metadata->add('token_usage', $tokenUsage);
51+
5152
return;
5253
}
5354

54-
$metadata->add('prompt_tokens', $content['usage']['prompt_tokens'] ?? null);
55-
$metadata->add('completion_tokens', $content['usage']['completion_tokens'] ?? null);
56-
$metadata->add('total_tokens', $content['usage']['total_tokens'] ?? null);
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+
61+
$metadata->add('token_usage', $tokenUsage);
5762
}
5863
}

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

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

1414
use Symfony\AI\Agent\Output;
1515
use Symfony\AI\Agent\OutputProcessorInterface;
16+
use Symfony\AI\Platform\Result\Metadata\TokenUsage;
1617
use Symfony\AI\Platform\Result\StreamResult;
1718
use Symfony\Contracts\HttpClient\ResponseInterface;
1819

@@ -35,19 +36,27 @@ public function processOutput(Output $output): void
3536

3637
$metadata = $output->result->getMetadata();
3738

38-
$metadata->add(
39-
'remaining_tokens',
40-
(int) $rawResponse->getHeaders(false)['x-ratelimit-remaining-tokens'][0],
39+
$remainingTokens = $rawResponse->getHeaders(false)['x-ratelimit-remaining-tokens'][0] ?? null;
40+
$tokenUsage = new TokenUsage(
41+
remainingTokens: null !== $remainingTokens ? (int) $remainingTokens : null,
4142
);
4243

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

4546
if (!\array_key_exists('usage', $content)) {
47+
$metadata->add('token_usage', $tokenUsage);
48+
4649
return;
4750
}
4851

49-
$metadata->add('prompt_tokens', $content['usage']['prompt_tokens'] ?? null);
50-
$metadata->add('completion_tokens', $content['usage']['completion_tokens'] ?? null);
51-
$metadata->add('total_tokens', $content['usage']['total_tokens'] ?? null);
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+
60+
$metadata->add('token_usage', $tokenUsage);
5261
}
5362
}

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
<?php
2+
3+
/*
4+
* This file is part of the Symfony package.
5+
*
6+
* (c) Fabien Potencier <[email protected]>
7+
*
8+
* For the full copyright and license information, please view the LICENSE
9+
* file that was distributed with this source code.
10+
*/
11+
12+
namespace Symfony\AI\Platform\Result\Metadata;
13+
14+
/**
15+
* @author Junaid Farooq <[email protected]>
16+
*/
17+
final class TokenUsage implements \JsonSerializable
18+
{
19+
public function __construct(
20+
public ?int $promptTokens = null,
21+
public ?int $completionTokens = null,
22+
public ?int $thinkingTokens = null,
23+
public ?int $cachedTokens = null,
24+
public ?int $remainingTokens = null,
25+
public ?int $remainingTokensMinute = null,
26+
public ?int $remainingTokensMonth = null,
27+
public ?int $totalTokens = null,
28+
) {
29+
}
30+
31+
/**
32+
* @return array<string, int|null>
33+
*/
34+
public function toArray(): array
35+
{
36+
return [
37+
'prompt' => $this->promptTokens,
38+
'completion' => $this->completionTokens,
39+
'thinking' => $this->thinkingTokens,
40+
'remaining' => $this->remainingTokens,
41+
'remaining_tokens_minute' => $this->remainingTokensMinute,
42+
'remaining_tokens_month' => $this->remainingTokensMonth,
43+
'total' => $this->totalTokens,
44+
];
45+
}
46+
47+
/**
48+
* @return array<string, int|null>
49+
*/
50+
public function jsonSerialize(): array
51+
{
52+
return $this->toArray();
53+
}
54+
}

‎src/platform/tests/Bridge/Mistral/TokenOutputProcessorTest.php

Lines changed: 23 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
use Symfony\AI\Platform\Message\MessageBagInterface;
2121
use Symfony\AI\Platform\Model;
2222
use Symfony\AI\Platform\Result\Metadata\Metadata;
23+
use Symfony\AI\Platform\Result\Metadata\TokenUsage;
2324
use Symfony\AI\Platform\Result\RawHttpResult;
2425
use Symfony\AI\Platform\Result\ResultInterface;
2526
use Symfony\AI\Platform\Result\StreamResult;
@@ -70,9 +71,12 @@ public function testItAddsRemainingTokensToMetadata()
7071
$processor->processOutput($output);
7172

7273
$metadata = $output->result->getMetadata();
73-
$this->assertCount(2, $metadata);
74-
$this->assertSame(1000, $metadata->get('remaining_tokens_minute'));
75-
$this->assertSame(1000000, $metadata->get('remaining_tokens_month'));
74+
$tokenUsage = $metadata->get('token_usage');
75+
76+
$this->assertCount(1, $metadata);
77+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
78+
$this->assertSame(1000, $tokenUsage->remainingTokensMinute);
79+
$this->assertSame(1000000, $tokenUsage->remainingTokensMonth);
7680
}
7781

7882
public function testItAddsUsageTokensToMetadata()
@@ -95,12 +99,14 @@ public function testItAddsUsageTokensToMetadata()
9599
$processor->processOutput($output);
96100

97101
$metadata = $output->result->getMetadata();
98-
$this->assertCount(5, $metadata);
99-
$this->assertSame(1000, $metadata->get('remaining_tokens_minute'));
100-
$this->assertSame(1000000, $metadata->get('remaining_tokens_month'));
101-
$this->assertSame(10, $metadata->get('prompt_tokens'));
102-
$this->assertSame(20, $metadata->get('completion_tokens'));
103-
$this->assertSame(30, $metadata->get('total_tokens'));
102+
$tokenUsage = $metadata->get('token_usage');
103+
104+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
105+
$this->assertSame(1000, $tokenUsage->remainingTokensMinute);
106+
$this->assertSame(1000000, $tokenUsage->remainingTokensMonth);
107+
$this->assertSame(10, $tokenUsage->promptTokens);
108+
$this->assertSame(20, $tokenUsage->completionTokens);
109+
$this->assertSame(30, $tokenUsage->totalTokens);
104110
}
105111

106112
public function testItHandlesMissingUsageFields()
@@ -122,12 +128,14 @@ public function testItHandlesMissingUsageFields()
122128
$processor->processOutput($output);
123129

124130
$metadata = $output->result->getMetadata();
125-
$this->assertCount(5, $metadata);
126-
$this->assertSame(1000, $metadata->get('remaining_tokens_minute'));
127-
$this->assertSame(1000000, $metadata->get('remaining_tokens_month'));
128-
$this->assertSame(10, $metadata->get('prompt_tokens'));
129-
$this->assertNull($metadata->get('completion_tokens'));
130-
$this->assertNull($metadata->get('total_tokens'));
131+
$tokenUsage = $metadata->get('token_usage');
132+
133+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
134+
$this->assertSame(1000, $tokenUsage->remainingTokensMinute);
135+
$this->assertSame(1000000, $tokenUsage->remainingTokensMonth);
136+
$this->assertSame(10, $tokenUsage->promptTokens);
137+
$this->assertNull($tokenUsage->completionTokens);
138+
$this->assertNull($tokenUsage->totalTokens);
131139
}
132140

133141
private function createRawResponse(array $data = []): RawHttpResult

‎src/platform/tests/Bridge/OpenAi/TokenOutputProcessorTest.php

Lines changed: 29 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
use Symfony\AI\Platform\Message\MessageBagInterface;
2121
use Symfony\AI\Platform\Model;
2222
use Symfony\AI\Platform\Result\Metadata\Metadata;
23+
use Symfony\AI\Platform\Result\Metadata\TokenUsage;
2324
use Symfony\AI\Platform\Result\RawHttpResult;
2425
use Symfony\AI\Platform\Result\ResultInterface;
2526
use Symfony\AI\Platform\Result\StreamResult;
@@ -31,6 +32,7 @@
3132
#[UsesClass(TextResult::class)]
3233
#[UsesClass(StreamResult::class)]
3334
#[UsesClass(Metadata::class)]
35+
#[UsesClass(TokenUsage::class)]
3436
#[Small]
3537
final class TokenOutputProcessorTest extends TestCase
3638
{
@@ -70,8 +72,11 @@ public function testItAddsRemainingTokensToMetadata()
7072
$processor->processOutput($output);
7173

7274
$metadata = $output->result->getMetadata();
75+
$tokenUsage = $metadata->get('token_usage');
76+
7377
$this->assertCount(1, $metadata);
74-
$this->assertSame(1000, $metadata->get('remaining_tokens'));
78+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
79+
$this->assertSame(1000, $tokenUsage->remainingTokens);
7580
}
7681

7782
public function testItAddsUsageTokensToMetadata()
@@ -83,7 +88,13 @@ public function testItAddsUsageTokensToMetadata()
8388
'usage' => [
8489
'prompt_tokens' => 10,
8590
'completion_tokens' => 20,
86-
'total_tokens' => 30,
91+
'total_tokens' => 50,
92+
'completion_tokens_details' => [
93+
'reasoning_tokens' => 20,
94+
],
95+
'prompt_tokens_details' => [
96+
'cached_tokens' => 40,
97+
],
8798
],
8899
]);
89100

@@ -94,11 +105,15 @@ public function testItAddsUsageTokensToMetadata()
94105
$processor->processOutput($output);
95106

96107
$metadata = $output->result->getMetadata();
97-
$this->assertCount(4, $metadata);
98-
$this->assertSame(1000, $metadata->get('remaining_tokens'));
99-
$this->assertSame(10, $metadata->get('prompt_tokens'));
100-
$this->assertSame(20, $metadata->get('completion_tokens'));
101-
$this->assertSame(30, $metadata->get('total_tokens'));
108+
$tokenUsage = $metadata->get('token_usage');
109+
110+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
111+
$this->assertSame(10, $tokenUsage->promptTokens);
112+
$this->assertSame(20, $tokenUsage->completionTokens);
113+
$this->assertSame(1000, $tokenUsage->remainingTokens);
114+
$this->assertSame(20, $tokenUsage->thinkingTokens);
115+
$this->assertSame(40, $tokenUsage->cachedTokens);
116+
$this->assertSame(50, $tokenUsage->totalTokens);
102117
}
103118

104119
public function testItHandlesMissingUsageFields()
@@ -120,11 +135,13 @@ public function testItHandlesMissingUsageFields()
120135
$processor->processOutput($output);
121136

122137
$metadata = $output->result->getMetadata();
123-
$this->assertCount(4, $metadata);
124-
$this->assertSame(1000, $metadata->get('remaining_tokens'));
125-
$this->assertSame(10, $metadata->get('prompt_tokens'));
126-
$this->assertNull($metadata->get('completion_tokens'));
127-
$this->assertNull($metadata->get('total_tokens'));
138+
$tokenUsage = $metadata->get('token_usage');
139+
140+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
141+
$this->assertSame(10, $tokenUsage->promptTokens);
142+
$this->assertSame(1000, $tokenUsage->remainingTokens);
143+
$this->assertNull($tokenUsage->completionTokens);
144+
$this->assertNull($tokenUsage->totalTokens);
128145
}
129146

130147
private function createRawResult(array $data = []): RawHttpResult

0 commit comments

Comments
 (0)