mirror of
https://github.com/symfony/ai-generic-platform.git
synced 2026-03-23 23:42:22 +01:00
Add support for token usage on generic bridge
This commit is contained in:
@@ -1,6 +1,11 @@
|
||||
CHANGELOG
|
||||
=========
|
||||
|
||||
0.4
|
||||
---
|
||||
|
||||
* Add support for token usage tracking
|
||||
|
||||
0.1
|
||||
---
|
||||
|
||||
|
||||
@@ -85,7 +85,7 @@ class ResultConverter implements ResultConverterInterface
|
||||
|
||||
public function getTokenUsageExtractor(): ?TokenUsageExtractorInterface
|
||||
{
|
||||
return null;
|
||||
return new TokenUsageExtractor();
|
||||
}
|
||||
|
||||
private function convertStream(RawResultInterface|RawHttpResult $result): \Generator
|
||||
|
||||
43
Completions/TokenUsageExtractor.php
Normal file
43
Completions/TokenUsageExtractor.php
Normal 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,
|
||||
);
|
||||
}
|
||||
}
|
||||
73
Tests/Completions/TokenUsageExtractorTest.php
Normal file
73
Tests/Completions/TokenUsageExtractorTest.php
Normal 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());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user