[Platform] Add token usage extraction for embeddings responses

Amoifr · chr-hertel · commit 9e206f8ef6c9 · 2026-04-03T23:09:13.000+02:00
Add TokenUsageExtractor for embeddings in OpenAi, Generic, Scaleway,
DockerModelRunner and VertexAi bridges.

The extractors parse token usage information from embedding API responses
and return a TokenUsage object with prompt and total token counts.

VertexAi uses a different response format with per-embedding statistics,
so the extractor sums up token_count from all predictions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,7 @@ CHANGELOG
 0.7
 ---
 
+ * Add token usage extraction for embeddings responses
  * [BC BREAK] OpenAI-compatible completion streams now yield `TextDelta`, `ThinkingDelta`, `ThinkingComplete`, `ToolCallStart`, `ToolInputDelta`, `ToolCallComplete`, and streamed `TokenUsage` deltas
 
 0.4
diff --git a/Embeddings/ResultConverter.php b/Embeddings/ResultConverter.php
@@ -68,8 +68,8 @@ public function convert(RawResultInterface $result, array $options = []): Vector
         );
     }
 
-    public function getTokenUsageExtractor(): ?TokenUsageExtractorInterface
+    public function getTokenUsageExtractor(): TokenUsageExtractorInterface
     {
-        return null;
+        return new TokenUsageExtractor();
     }
 }
diff --git a/Embeddings/TokenUsageExtractor.php b/Embeddings/TokenUsageExtractor.php
@@ -0,0 +1,37 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <fabien@symfony.com>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\Platform\Bridge\Generic\Embeddings;
+
+use Symfony\AI\Platform\Result\RawResultInterface;
+use Symfony\AI\Platform\TokenUsage\TokenUsage;
+use Symfony\AI\Platform\TokenUsage\TokenUsageExtractorInterface;
+use Symfony\AI\Platform\TokenUsage\TokenUsageInterface;
+
+/**
+ * @author Pascal Cescon <pascal.cescon@gmail.com>
+ */
+final class TokenUsageExtractor implements TokenUsageExtractorInterface
+{
+    public function extract(RawResultInterface $rawResult, array $options = []): ?TokenUsageInterface
+    {
+        $content = $rawResult->getData();
+
+        if (!\array_key_exists('usage', $content)) {
+            return null;
+        }
+
+        return new TokenUsage(
+            promptTokens: $content['usage']['prompt_tokens'] ?? null,
+            totalTokens: $content['usage']['total_tokens'] ?? null,
+        );
+    }
+}
diff --git a/Tests/Embeddings/TokenUsageExtractorTest.php b/Tests/Embeddings/TokenUsageExtractorTest.php
@@ -0,0 +1,61 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <fabien@symfony.com>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\Platform\Bridge\Generic\Tests\Embeddings;
+
+use PHPUnit\Framework\TestCase;
+use Symfony\AI\Platform\Bridge\Generic\Embeddings\TokenUsageExtractor;
+use Symfony\AI\Platform\Result\InMemoryRawResult;
+use Symfony\AI\Platform\TokenUsage\TokenUsage;
+
+final class TokenUsageExtractorTest extends TestCase
+{
+    public function testItReturnsNullWithoutUsageData()
+    {
+        $extractor = new TokenUsageExtractor();
+
+        $this->assertNull($extractor->extract(new InMemoryRawResult(['some' => 'data'])));
+    }
+
+    public function testItExtractsTokenUsage()
+    {
+        $extractor = new TokenUsageExtractor();
+        $result = new InMemoryRawResult([
+            'usage' => [
+                'prompt_tokens' => 20,
+                'total_tokens' => 20,
+            ],
+        ]);
+
+        $tokenUsage = $extractor->extract($result);
+
+        $this->assertInstanceOf(TokenUsage::class, $tokenUsage);
+        $this->assertSame(20, $tokenUsage->getPromptTokens());
+        $this->assertNull($tokenUsage->getCompletionTokens());
+        $this->assertSame(20, $tokenUsage->getTotalTokens());
+    }
+
+    public function testItHandlesMissingUsageFields()
+    {
+        $extractor = new TokenUsageExtractor();
+        $result = new InMemoryRawResult([
+            'usage' => [
+                'prompt_tokens' => 5,
+            ],
+        ]);
+
+        $tokenUsage = $extractor->extract($result);
+
+        $this->assertInstanceOf(TokenUsage::class, $tokenUsage);
+        $this->assertSame(5, $tokenUsage->getPromptTokens());
+        $this->assertNull($tokenUsage->getTotalTokens());
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -68,8 +68,8 @@ public function convert(RawResultInterface $result, array $options = []): Vector`
`68`	`68`	`);`
`69`	`69`	`}`
`70`	`70`
`71`		`- public function getTokenUsageExtractor(): ?TokenUsageExtractorInterface`
	`71`	`+ public function getTokenUsageExtractor(): TokenUsageExtractorInterface`
`72`	`72`	`{`
`73`		`- return null;`
	`73`	`+ return new TokenUsageExtractor();`
`74`	`74`	`}`
`75`	`75`	`}`