Skip to content

Commit b08d7e7

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

File tree

3 files changed

+245
-0
lines changed

3 files changed

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

‎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)