Skip to content

Commit 605bf5e

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

File tree

7 files changed

+256
-10
lines changed

7 files changed

+256
-10
lines changed

‎examples/composer.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
"async-aws/bedrock-runtime": "^1.1",
99
"codewithkyrian/transformers": "^0.6.1",
1010
"doctrine/dbal": "^3.3|^4.0",
11+
"google/auth": "^1.47",
1112
"mrmysql/youtube-transcript": "^0.0.5",
1213
"php-http/discovery": "^1.20",
1314
"probots-io/pinecone-php": "^1.1",

‎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/ai-bundle/doc/index.rst

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ Configuration
3434
class: 'Symfony\AI\Platform\Bridge\OpenAi\Gpt'
3535
name: !php/const Symfony\AI\Platform\Bridge\OpenAi\Gpt::GPT_4O_MINI
3636
37-
**Advanced Example with Anthropic, Azure, Gemini, VertexAi and multiple agents**
37+
**Advanced Example with Anthropic, Azure, Gemini, Vertex AI and multiple agents**
3838

3939
.. code-block:: yaml
4040
@@ -52,7 +52,7 @@ Configuration
5252
api_version: '%env(AZURE_GPT_VERSION)%'
5353
gemini:
5454
api_key: '%env(GEMINI_API_KEY)%'
55-
vertexAi:
55+
vertexai:
5656
location: '%env(GOOGLE_CLOUD_LOCATION)%'
5757
project_id: '%env(GOOGLE_CLOUD_PROJECT)%'
5858
ollama:

‎src/platform/doc/index.rst

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -73,15 +73,15 @@ usually defined by the specific models and their documentation.
7373
* `OpenAI's GPT`_ with `OpenAI`_ and `Azure`_ as Platform
7474
* `Anthropic's Claude`_ with `Anthropic`_ and `AWS Bedrock`_ as Platform
7575
* `Meta's Llama`_ with `Azure`_, `Ollama`_, `Replicate`_ and `AWS Bedrock`_ as Platform
76-
* `Gemini`_ with `Google`_, `VertexAi`_ and `OpenRouter`_ as Platform
77-
* `VertexAi Gen AI`_ with `VertexAi`_ as Platform
76+
* `Gemini`_ with `Google`_, `Vertex AI`_ and `OpenRouter`_ as Platform
77+
* `Vertex AI Gen AI`_ with `Vertex AI`_ as Platform
7878
* `DeepSeek's R1`_ with `OpenRouter`_ as Platform
7979
* `Amazon's Nova`_ with `AWS Bedrock`_ as Platform
8080
* `Mistral's Mistral`_ with `Mistral`_ as Platform
8181
* `Albert API`_ models with `Albert`_ as Platform (French government's sovereign AI gateway)
8282
* **Embeddings Models**
8383
* `Gemini Text Embeddings`_ with `Google`_
84-
* `VertexAi Text Embeddings`_ with `VertexAi`_
84+
* `Vertex AI Text Embeddings`_ with `Vertex AI`_
8585
* `OpenAI's Text Embeddings`_ with `OpenAI`_ and `Azure`_ as Platform
8686
* `Voyage's Embeddings`_ with `Voyage`_ as Platform
8787
* `Mistral Embed`_ with `Mistral`_ as Platform
@@ -266,7 +266,7 @@ Server Tools
266266
Some platforms provide built-in server-side tools for enhanced capabilities without custom implementations:
267267

268268
1. **[Gemini](gemini-server-tools.rst)** - URL Context, Google Search, Code Execution
269-
2. **[VertexAi](vertexai-server-tools.rst)** - URL Context, Google Search, Code Execution
269+
2. **[Vertex AI](vertexai-server-tools.rst)** - URL Context, Google Search, Code Execution
270270

271271
Parallel Platform Calls
272272
-----------------------
@@ -366,7 +366,7 @@ This allows fast and isolated testing of AI-powered features without relying on
366366
.. _`Ollama`: https://ollama.com/
367367
.. _`Replicate`: https://replicate.com/
368368
.. _`Gemini`: https://gemini.google.com/
369-
.. _`VertexAi`: https://cloud.google.com/vertex-ai/generative-ai/docs
369+
.. _`Vertex AI`: https://cloud.google.com/vertex-ai/generative-ai/docs
370370
.. _`Google`: https://ai.google.dev/
371371
.. _`OpenRouter`: https://www.openrouter.ai/
372372
.. _`DeepSeek's R1`: https://www.deepseek.com/
@@ -376,8 +376,8 @@ This allows fast and isolated testing of AI-powered features without relying on
376376
.. _`Albert`: https://alliance.numerique.gouv.fr/produit/albert/
377377
.. _`Mistral`: https://www.mistral.ai/
378378
.. _`Gemini Text Embeddings`: https://ai.google.dev/gemini-api/docs/embeddings
379-
.. _`VertexAi Gen AI`: https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference
380-
.. _`VertexAi Text Embeddings`: https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api
379+
.. _`Vertex AI Gen AI`: https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference
380+
.. _`Vertex AI Text Embeddings`: https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api
381381
.. _`OpenAI's Text Embeddings`: https://platform.openai.com/docs/guides/embeddings/embedding-models
382382
.. _`Voyage's Embeddings`: https://docs.voyageai.com/docs/embeddings
383383
.. _`Voyage`: https://www.voyageai.com/

‎src/platform/src/Bridge/VertexAi/Contract/AssistantMessageNormalizer.php

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@ final class AssistantMessageNormalizer extends ModelContractNormalizer
3131
* name: string,
3232
* args?: array<int|string, mixed>
3333
* }
34-
* }
3534
* }
35+
* }
3636
*/
3737
public function normalize(mixed $data, ?string $format = null, array $context = []): array
3838
{

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