Add support for token usage on generic bridge

This commit is contained in:
Christopher Hertel
2026-01-29 16:01:30 +01:00
parent 271fa3d5fd
commit fc7ef91d05
4 changed files with 122 additions and 1 deletions

View File

@@ -1,6 +1,11 @@
CHANGELOG
=========
0.4
---
* Add support for token usage tracking
0.1
---

View File

@@ -85,7 +85,7 @@ class ResultConverter implements ResultConverterInterface
public function getTokenUsageExtractor(): ?TokenUsageExtractorInterface
{
return null;
return new TokenUsageExtractor();
}
private function convertStream(RawResultInterface|RawHttpResult $result): \Generator

View File

@@ -0,0 +1,43 @@
<?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\Completions;
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 Christopher Hertel <mail@christopher-hertel.de>
*/
final class TokenUsageExtractor implements TokenUsageExtractorInterface
{
public function extract(RawResultInterface $rawResult, array $options = []): ?TokenUsageInterface
{
if ($options['stream'] ?? false) {
return null;
}
$content = $rawResult->getData();
if (!\array_key_exists('usage', $content)) {
return null;
}
return new TokenUsage(
promptTokens: $content['usage']['prompt_tokens'] ?? null,
completionTokens: $content['usage']['completion_tokens'] ?? null,
cachedTokens: $content['usage']['num_cached_tokens'] ?? null,
totalTokens: $content['usage']['total_tokens'] ?? null,
);
}
}

View File

@@ -0,0 +1,73 @@
<?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\Completions;
use PHPUnit\Framework\TestCase;
use Symfony\AI\Platform\Bridge\Generic\Completions\TokenUsageExtractor;
use Symfony\AI\Platform\Result\InMemoryRawResult;
use Symfony\AI\Platform\TokenUsage\TokenUsage;
final class TokenUsageExtractorTest extends TestCase
{
public function testItHandlesStreamResponsesWithoutProcessing()
{
$extractor = new TokenUsageExtractor();
$this->assertNull($extractor->extract(new InMemoryRawResult(), ['stream' => true]));
}
public function testItDoesNothingWithoutUsageData()
{
$extractor = new TokenUsageExtractor();
$this->assertNull($extractor->extract(new InMemoryRawResult(['some' => 'data'])));
}
public function testItExtractsTokenUsage()
{
$extractor = new TokenUsageExtractor();
$result = new InMemoryRawResult([
'usage' => [
'prompt_tokens' => 10,
'completion_tokens' => 20,
'total_tokens' => 30,
'num_cached_tokens' => 5,
],
]);
$tokenUsage = $extractor->extract($result);
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
$this->assertSame(10, $tokenUsage->getPromptTokens());
$this->assertSame(20, $tokenUsage->getCompletionTokens());
$this->assertSame(5, $tokenUsage->getCachedTokens());
$this->assertSame(30, $tokenUsage->getTotalTokens());
}
public function testItHandlesMissingUsageFields()
{
$extractor = new TokenUsageExtractor();
$result = new InMemoryRawResult([
'usage' => [
'prompt_tokens' => 10,
],
]);
$tokenUsage = $extractor->extract($result);
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
$this->assertSame(10, $tokenUsage->getPromptTokens());
$this->assertNull($tokenUsage->getCompletionTokens());
$this->assertNull($tokenUsage->getCachedTokens());
$this->assertNull($tokenUsage->getTotalTokens());
}
}