Skip to content

Commit 0cf0db1

Browse files
junaidbinfarooqjunaidbinfarooq
committed
feat(platform): Add support for Google vertex AI
- Adds support to calculate the token usage
1 parent 1b335cf commit 0cf0db1

File tree

3 files changed

+236
-0
lines changed

3 files changed

+236
-0
lines changed

‎examples/vertexai/token-metadata.php

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
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\VertexAi\Gemini\Model;
14+
use Symfony\AI\Platform\Bridge\VertexAi\PlatformFactory;
15+
use Symfony\AI\Platform\Message\Message;
16+
use Symfony\AI\Platform\Message\MessageBag;
17+
18+
require_once __DIR__.'/bootstrap.php';
19+
20+
$platform = PlatformFactory::create(env('GOOGLE_CLOUD_LOCATION'), env('GOOGLE_CLOUD_PROJECT'), adc_aware_http_client());
21+
$model = new Model(Model::GEMINI_2_0_FLASH_LITE);
22+
23+
$agent = new Agent($platform, $model, outputProcessors: [new Symfony\AI\Platform\Bridge\VertexAi\TokenOutputProcessor()], logger: logger());
24+
$messages = new MessageBag(
25+
Message::forSystem('You are an expert assistant in animal study.'),
26+
Message::ofUser('What does a cat usually eat?'),
27+
);
28+
$result = $agent->call($messages);
29+
30+
$metadata = $result->getMetadata();
31+
32+
echo 'Prompt Tokens: '.$metadata['prompt_tokens'].\PHP_EOL;
33+
echo 'Completion Tokens: '.$metadata['completion_tokens'].\PHP_EOL;
34+
echo 'Thinking Tokens: '.$metadata['thinking_tokens'].\PHP_EOL;
35+
echo 'Utilized Tokens: '.$metadata['total_tokens'].\PHP_EOL;

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

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
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\VertexAi;
13+
14+
use Symfony\AI\Agent\Output;
15+
use Symfony\AI\Agent\OutputProcessorInterface;
16+
use Symfony\AI\Platform\Result\StreamResult;
17+
use Symfony\Contracts\HttpClient\Exception\ClientExceptionInterface;
18+
use Symfony\Contracts\HttpClient\Exception\DecodingExceptionInterface;
19+
use Symfony\Contracts\HttpClient\Exception\RedirectionExceptionInterface;
20+
use Symfony\Contracts\HttpClient\Exception\ServerExceptionInterface;
21+
use Symfony\Contracts\HttpClient\Exception\TransportExceptionInterface;
22+
use Symfony\Contracts\HttpClient\ResponseInterface;
23+
24+
/**
25+
* @author Junaid Farooq <[email protected]>
26+
*/
27+
final class TokenOutputProcessor implements OutputProcessorInterface
28+
{
29+
/**
30+
* @throws TransportExceptionInterface
31+
* @throws ServerExceptionInterface
32+
* @throws RedirectionExceptionInterface
33+
* @throws DecodingExceptionInterface
34+
* @throws ClientExceptionInterface
35+
*/
36+
public function processOutput(Output $output): void
37+
{
38+
$metadata = $output->result->getMetadata();
39+
40+
if ($output->result instanceof StreamResult) {
41+
$lastChunk = null;
42+
43+
foreach ($output->result as $chunk) {
44+
// Store last event that contains usage metadata
45+
if (isset($chunk['usageMetadata'])) {
46+
$lastChunk = $chunk;
47+
}
48+
}
49+
50+
if ($lastChunk) {
51+
$this->extractUsageMetadata($lastChunk['usageMetadata'], $metadata);
52+
}
53+
54+
return;
55+
}
56+
57+
$rawResponse = $output->result->getRawResult()?->getObject();
58+
if (!$rawResponse instanceof ResponseInterface) {
59+
return;
60+
}
61+
62+
$content = $rawResponse->toArray(false);
63+
64+
if (!isset($content['usageMetadata'])) {
65+
return;
66+
}
67+
68+
$this->extractUsageMetadata($content['usageMetadata'], $metadata);
69+
}
70+
71+
private function extractUsageMetadata(array $usage, $metadata): void
72+
{
73+
$metadata->add('prompt_tokens', $usage['promptTokenCount'] ?? null);
74+
$metadata->add('completion_tokens', $usage['candidatesTokenCount'] ?? null);
75+
$metadata->add('thinking_tokens', $usage['thoughtsTokenCount'] ?? null);
76+
$metadata->add('total_tokens', $usage['totalTokenCount'] ?? null);
77+
}
78+
}

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

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
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\VertexAi;
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\VertexAi\TokenOutputProcessor;
20+
use Symfony\AI\Platform\Message\MessageBagInterface;
21+
use Symfony\AI\Platform\Model;
22+
use Symfony\AI\Platform\Result\Metadata\Metadata;
23+
use Symfony\AI\Platform\Result\RawHttpResult;
24+
use Symfony\AI\Platform\Result\ResultInterface;
25+
use Symfony\AI\Platform\Result\StreamResult;
26+
use Symfony\AI\Platform\Result\TextResult;
27+
use Symfony\Contracts\HttpClient\ResponseInterface;
28+
29+
#[CoversClass(TokenOutputProcessor::class)]
30+
#[UsesClass(Output::class)]
31+
#[UsesClass(TextResult::class)]
32+
#[UsesClass(StreamResult::class)]
33+
#[UsesClass(Metadata::class)]
34+
#[Small]
35+
final class TokenOutputProcessorTest extends TestCase
36+
{
37+
public function testItDoesNothingWithoutRawResponse()
38+
{
39+
$processor = new TokenOutputProcessor();
40+
$textResult = new TextResult('test');
41+
$output = $this->createOutput($textResult);
42+
43+
$processor->processOutput($output);
44+
45+
$this->assertCount(0, $output->result->getMetadata());
46+
}
47+
48+
public function testItAddsUsageTokensToMetadata()
49+
{
50+
// Arrange
51+
$textResult = new TextResult('test');
52+
53+
$rawResponse = $this->createRawResponse([
54+
'usageMetadata' => [
55+
'promptTokenCount' => 10,
56+
'candidatesTokenCount' => 20,
57+
'thoughtsTokenCount' => 20,
58+
'totalTokenCount' => 50,
59+
],
60+
]);
61+
62+
$textResult->setRawResult($rawResponse);
63+
$processor = new TokenOutputProcessor();
64+
$output = $this->createOutput($textResult);
65+
66+
// Act
67+
$processor->processOutput($output);
68+
69+
// Assert
70+
$metadata = $output->result->getMetadata();
71+
$this->assertCount(4, $metadata);
72+
$this->assertSame(10, $metadata->get('prompt_tokens'));
73+
$this->assertSame(20, $metadata->get('completion_tokens'));
74+
$this->assertSame(20, $metadata->get('thinking_tokens'));
75+
$this->assertSame(50, $metadata->get('total_tokens'));
76+
}
77+
78+
public function testItHandlesMissingUsageFields()
79+
{
80+
// Arrange
81+
$textResult = new TextResult('test');
82+
83+
$rawResponse = $this->createRawResponse([
84+
'usageMetadata' => [
85+
'promptTokenCount' => 10,
86+
],
87+
]);
88+
89+
$textResult->setRawResult($rawResponse);
90+
$processor = new TokenOutputProcessor();
91+
$output = $this->createOutput($textResult);
92+
93+
// Act
94+
$processor->processOutput($output);
95+
96+
// Assert
97+
$metadata = $output->result->getMetadata();
98+
$this->assertCount(4, $metadata);
99+
$this->assertSame(10, $metadata->get('prompt_tokens'));
100+
$this->assertNull($metadata->get('completion_tokens'));
101+
$this->assertNull($metadata->get('completion_tokens'));
102+
$this->assertNull($metadata->get('total_tokens'));
103+
}
104+
105+
private function createRawResponse(array $data = []): RawHttpResult
106+
{
107+
$rawResponse = $this->createStub(ResponseInterface::class);
108+
109+
$rawResponse->method('toArray')->willReturn($data);
110+
111+
return new RawHttpResult($rawResponse);
112+
}
113+
114+
private function createOutput(ResultInterface $result): Output
115+
{
116+
return new Output(
117+
$this->createStub(Model::class),
118+
$result,
119+
$this->createStub(MessageBagInterface::class),
120+
[],
121+
);
122+
}
123+
}

0 commit comments

Comments
 (0)