Skip to content

Commit a071848

Browse files
OskarStarkOskarStark
committed
feature #347 [Platform][Gemini] Add TokenOutputProcessor (VincentLanglet)
This PR was squashed before being merged into the main branch. Discussion ---------- [Platform][Gemini] Add `TokenOutputProcessor` | Q | A | ------------- | --- | Bug fix? | no | New feature? | yes | Docs? | no <!-- required for new features --> | Issues | Fix #... <!-- prefix each issue number with "Fix #", no need to create an issue if none exist, explain below instead --> | License | MIT According to https://ai.google.dev/api/generate-content?hl=en#UsageMetadata This will need the `AsTokenUsageProcessor` from PR of `@junaidbinfarooq` Commits ------- 59cc045 [Platform][Gemini] Add `TokenOutputProcessor`
2 parents a9944b1 + 59cc045 commit a071848

File tree

7 files changed

+259
-7
lines changed

7 files changed

+259
-7
lines changed

‎examples/bootstrap.php‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ function print_token_usage(Metadata $metadata): void
6464
echo 'Prompt tokens: '.$tokenUsage->promptTokens.\PHP_EOL;
6565
echo 'Completion tokens: '.$tokenUsage->completionTokens.\PHP_EOL;
6666
echo 'Thinking tokens: '.$tokenUsage->thinkingTokens.\PHP_EOL;
67+
echo 'Tool tokens: '.$tokenUsage->toolTokens.\PHP_EOL;
6768
echo 'Cached tokens: '.$tokenUsage->cachedTokens.\PHP_EOL;
6869
echo 'Remaining tokens minute: '.$tokenUsage->remainingTokensMinute.\PHP_EOL;
6970
echo 'Remaining tokens month: '.$tokenUsage->remainingTokensMonth.\PHP_EOL;

‎examples/gemini/token-metadata.php‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
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+
use Symfony\AI\Agent\Agent;
13+
use Symfony\AI\Platform\Bridge\Gemini\Gemini;
14+
use Symfony\AI\Platform\Bridge\Gemini\PlatformFactory;
15+
use Symfony\AI\Platform\Bridge\Gemini\TokenOutputProcessor;
16+
use Symfony\AI\Platform\Message\Message;
17+
use Symfony\AI\Platform\Message\MessageBag;
18+
19+
require_once dirname(__DIR__).'/bootstrap.php';
20+
21+
$platform = PlatformFactory::create(env('GEMINI_API_KEY'), http_client());
22+
$model = new Gemini(Gemini::GEMINI_2_FLASH);
23+
24+
$agent = new Agent($platform, $model, outputProcessors: [new TokenOutputProcessor()], logger: logger());
25+
$messages = new MessageBag(
26+
Message::forSystem('You are a pirate and you write funny.'),
27+
Message::ofUser('What is the Symfony framework?'),
28+
);
29+
$result = $agent->call($messages);
30+
31+
$metadata = $result->getMetadata();
32+
$tokenUsage = $metadata->get('token_usage');
33+
34+
print_token_usage($result->getMetadata());

‎examples/vertexai/token-metadata.php‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
use Symfony\AI\Platform\Bridge\VertexAi\PlatformFactory;
1515
use Symfony\AI\Platform\Message\Message;
1616
use Symfony\AI\Platform\Message\MessageBag;
17-
use Symfony\AI\Platform\Result\Metadata\TokenUsage;
1817

1918
require_once __DIR__.'/bootstrap.php';
2019

@@ -31,9 +30,4 @@
3130
$metadata = $result->getMetadata();
3231
$tokenUsage = $metadata->get('token_usage');
3332

34-
assert($tokenUsage instanceof TokenUsage);
35-
36-
echo 'Prompt Tokens: '.$tokenUsage->promptTokens.\PHP_EOL;
37-
echo 'Completion Tokens: '.$tokenUsage->completionTokens.\PHP_EOL;
38-
echo 'Thinking Tokens: '.$tokenUsage->thinkingTokens.\PHP_EOL;
39-
echo 'Utilized Tokens: '.$tokenUsage->totalTokens.\PHP_EOL;
33+
print_token_usage($result->getMetadata());

‎src/ai-bundle/config/services.php‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
use Symfony\AI\AiBundle\Security\EventListener\IsGrantedToolAttributeListener;
2727
use Symfony\AI\Platform\Bridge\Anthropic\Contract\AnthropicContract;
2828
use Symfony\AI\Platform\Bridge\Gemini\Contract\GeminiContract;
29+
use Symfony\AI\Platform\Bridge\Gemini\TokenOutputProcessor as GeminiTokenOutputProcessor;
2930
use Symfony\AI\Platform\Bridge\Mistral\TokenOutputProcessor as MistralTokenOutputProcessor;
3031
use Symfony\AI\Platform\Bridge\Ollama\Contract\OllamaContract;
3132
use Symfony\AI\Platform\Bridge\OpenAi\Contract\OpenAiContract;
@@ -137,6 +138,7 @@
137138

138139
// token usage processors
139140
->set('ai.platform.token_usage_processor.mistral', MistralTokenOutputProcessor::class)
141+
->set('ai.platform.token_usage_processor.gemini', GeminiTokenOutputProcessor::class)
140142
->set('ai.platform.token_usage_processor.openai', OpenAiTokenOutputProcessor::class)
141143
->set('ai.platform.token_usage_processor.vertexai', VertexAiTokenOutputProcessor::class)
142144
;

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

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
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\Bridge\Gemini;
13+
14+
use Symfony\AI\Agent\Output;
15+
use Symfony\AI\Agent\OutputProcessorInterface;
16+
use Symfony\AI\Platform\Metadata\TokenUsage;
17+
use Symfony\AI\Platform\Result\StreamResult;
18+
use Symfony\Contracts\HttpClient\ResponseInterface;
19+
20+
final class TokenOutputProcessor implements OutputProcessorInterface
21+
{
22+
public function processOutput(Output $output): void
23+
{
24+
if ($output->result instanceof StreamResult) {
25+
// Streams have to be handled manually as the tokens are part of the streamed chunks
26+
return;
27+
}
28+
29+
$rawResponse = $output->result->getRawResult()?->getObject();
30+
if (!$rawResponse instanceof ResponseInterface) {
31+
return;
32+
}
33+
34+
$metadata = $output->result->getMetadata();
35+
36+
$tokenUsage = new TokenUsage();
37+
38+
$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'];
46+
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;
53+
54+
$metadata->add('token_usage', $tokenUsage);
55+
}
56+
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ public function __construct(
2020
public ?int $promptTokens = null,
2121
public ?int $completionTokens = null,
2222
public ?int $thinkingTokens = null,
23+
public ?int $toolTokens = null,
2324
public ?int $cachedTokens = null,
2425
public ?int $remainingTokens = null,
2526
public ?int $remainingTokensMinute = null,
@@ -33,6 +34,7 @@ public function __construct(
3334
* prompt_tokens: ?int,
3435
* completion_tokens: ?int,
3536
* thinking_tokens: ?int,
37+
* tool_tokens: ?int,
3638
* cached_tokens: ?int,
3739
* remaining_tokens: ?int,
3840
* remaining_tokens_minute: ?int,
@@ -46,6 +48,7 @@ public function jsonSerialize(): array
4648
'prompt_tokens' => $this->promptTokens,
4749
'completion_tokens' => $this->completionTokens,
4850
'thinking_tokens' => $this->thinkingTokens,
51+
'tool_tokens' => $this->toolTokens,
4952
'cached_tokens' => $this->cachedTokens,
5053
'remaining_tokens' => $this->remainingTokens,
5154
'remaining_tokens_minute' => $this->remainingTokensMinute,

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

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
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\Tests\Bridge\Gemini;
13+
14+
use PHPUnit\Framework\Attributes\CoversClass;
15+
use PHPUnit\Framework\Attributes\Small;
16+
use PHPUnit\Framework\Attributes\UsesClass;
17+
use PHPUnit\Framework\TestCase;
18+
use Symfony\AI\Agent\Output;
19+
use Symfony\AI\Platform\Bridge\Gemini\TokenOutputProcessor;
20+
use Symfony\AI\Platform\Message\MessageBag;
21+
use Symfony\AI\Platform\Metadata\Metadata;
22+
use Symfony\AI\Platform\Metadata\TokenUsage;
23+
use Symfony\AI\Platform\Model;
24+
use Symfony\AI\Platform\Result\RawHttpResult;
25+
use Symfony\AI\Platform\Result\ResultInterface;
26+
use Symfony\AI\Platform\Result\StreamResult;
27+
use Symfony\AI\Platform\Result\TextResult;
28+
use Symfony\Contracts\HttpClient\ResponseInterface;
29+
30+
#[CoversClass(TokenOutputProcessor::class)]
31+
#[UsesClass(Output::class)]
32+
#[UsesClass(TextResult::class)]
33+
#[UsesClass(StreamResult::class)]
34+
#[UsesClass(Metadata::class)]
35+
#[UsesClass(TokenUsage::class)]
36+
#[Small]
37+
final class TokenOutputProcessorTest extends TestCase
38+
{
39+
public function testItHandlesStreamResponsesWithoutProcessing()
40+
{
41+
$processor = new TokenOutputProcessor();
42+
$streamResult = new StreamResult((static function () { yield 'test'; })());
43+
$output = $this->createOutput($streamResult);
44+
45+
$processor->processOutput($output);
46+
47+
$metadata = $output->result->getMetadata();
48+
$this->assertCount(0, $metadata);
49+
}
50+
51+
public function testItDoesNothingWithoutRawResponse()
52+
{
53+
$processor = new TokenOutputProcessor();
54+
$textResult = new TextResult('test');
55+
$output = $this->createOutput($textResult);
56+
57+
$processor->processOutput($output);
58+
59+
$metadata = $output->result->getMetadata();
60+
$this->assertCount(0, $metadata);
61+
}
62+
63+
public function testItAddsRemainingTokensToMetadata()
64+
{
65+
$processor = new TokenOutputProcessor();
66+
$textResult = new TextResult('test');
67+
68+
$textResult->setRawResult($this->createRawResult());
69+
70+
$output = $this->createOutput($textResult);
71+
72+
$processor->processOutput($output);
73+
74+
$metadata = $output->result->getMetadata();
75+
$tokenUsage = $metadata->get('token_usage');
76+
77+
$this->assertCount(1, $metadata);
78+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
79+
$this->assertNull($tokenUsage->remainingTokens);
80+
}
81+
82+
public function testItAddsUsageTokensToMetadata()
83+
{
84+
$processor = new TokenOutputProcessor();
85+
$textResult = new TextResult('test');
86+
87+
$rawResult = $this->createRawResult([
88+
'usageMetadata' => [
89+
'promptTokenCount' => 10,
90+
'candidatesTokenCount' => 20,
91+
'totalTokenCount' => 50,
92+
'thoughtsTokenCount' => 20,
93+
'cachedContentTokenCount' => 40,
94+
'toolUsePromptTokenCount' => 5,
95+
],
96+
]);
97+
98+
$textResult->setRawResult($rawResult);
99+
100+
$output = $this->createOutput($textResult);
101+
102+
$processor->processOutput($output);
103+
104+
$metadata = $output->result->getMetadata();
105+
$tokenUsage = $metadata->get('token_usage');
106+
107+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
108+
$this->assertSame(10, $tokenUsage->promptTokens);
109+
$this->assertSame(5, $tokenUsage->toolTokens);
110+
$this->assertSame(20, $tokenUsage->completionTokens);
111+
$this->assertNull($tokenUsage->remainingTokens);
112+
$this->assertSame(20, $tokenUsage->thinkingTokens);
113+
$this->assertSame(40, $tokenUsage->cachedTokens);
114+
$this->assertSame(50, $tokenUsage->totalTokens);
115+
}
116+
117+
public function testItHandlesMissingUsageFields()
118+
{
119+
$processor = new TokenOutputProcessor();
120+
$textResult = new TextResult('test');
121+
122+
$rawResult = $this->createRawResult([
123+
'usageMetadata' => [
124+
// Missing some fields
125+
'promptTokenCount' => 10,
126+
],
127+
]);
128+
129+
$textResult->setRawResult($rawResult);
130+
131+
$output = $this->createOutput($textResult);
132+
133+
$processor->processOutput($output);
134+
135+
$metadata = $output->result->getMetadata();
136+
$tokenUsage = $metadata->get('token_usage');
137+
138+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
139+
$this->assertSame(10, $tokenUsage->promptTokens);
140+
$this->assertNull($tokenUsage->remainingTokens);
141+
$this->assertNull($tokenUsage->completionTokens);
142+
$this->assertNull($tokenUsage->totalTokens);
143+
}
144+
145+
private function createRawResult(array $data = []): RawHttpResult
146+
{
147+
$rawResponse = $this->createStub(ResponseInterface::class);
148+
$rawResponse->method('toArray')->willReturn($data);
149+
150+
return new RawHttpResult($rawResponse);
151+
}
152+
153+
private function createOutput(ResultInterface $result): Output
154+
{
155+
return new Output(
156+
$this->createStub(Model::class),
157+
$result,
158+
$this->createStub(MessageBag::class),
159+
[],
160+
);
161+
}
162+
}

0 commit comments

Comments
 (0)