Skip to content

Commit 7a34c26

Browse files
junaidbinfarooqjunaidbinfarooq
committed
feat(platform): Standardize token usage
- Makes token usage explicit - Adds support to automatically register token usage extractors so that the token usage information is available in the result when a model is invoked - Makes the token usage auto-registration configurable
1 parent 1336ce7 commit 7a34c26

18 files changed

+498
-197
lines changed

‎examples/mistral/token-metadata.php

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
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\Mistral\Mistral;
14+
use Symfony\AI\Platform\Bridge\Mistral\PlatformFactory;
15+
use Symfony\AI\Platform\Bridge\Mistral\TokenUsageExtractor;
16+
use Symfony\AI\Platform\Message\Message;
17+
use Symfony\AI\Platform\Message\MessageBag;
18+
use Symfony\AI\Platform\Result\TokenUsage\TokenUsageOutputProcessor;
19+
20+
require_once dirname(__DIR__).'/bootstrap.php';
21+
22+
$platform = PlatformFactory::create(env('MISTRAL_API_KEY'), http_client());
23+
$model = new Mistral(Mistral::MISTRAL_SMALL, [
24+
'temperature' => 0.5, // default options for the model
25+
]);
26+
$agent = new Agent(
27+
$platform,
28+
$model,
29+
outputProcessors: [new TokenUsageOutputProcessor(new TokenUsageExtractor())],
30+
logger: logger()
31+
);
32+
33+
$messages = new MessageBag(
34+
Message::forSystem('You are a pirate and you write funny.'),
35+
Message::ofUser('What is the Symfony framework?'),
36+
);
37+
38+
$result = $agent->call($messages, [
39+
'max_tokens' => 500,
40+
]);
41+
42+
if (null === $tokenUsage = $result->getTokenUsage()) {
43+
throw new RuntimeException('Token usage is not available.');
44+
}
45+
46+
echo 'Utilized Tokens: '.$tokenUsage->total.\PHP_EOL;
47+
echo '-- Prompt Tokens: '.$tokenUsage->prompt.\PHP_EOL;
48+
echo '-- Completion Tokens: '.$tokenUsage->completion.\PHP_EOL;

‎examples/openai/token-metadata.php

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,10 @@
1212
use Symfony\AI\Agent\Agent;
1313
use Symfony\AI\Platform\Bridge\OpenAi\Gpt;
1414
use Symfony\AI\Platform\Bridge\OpenAi\PlatformFactory;
15-
use Symfony\AI\Platform\Bridge\OpenAi\TokenOutputProcessor;
15+
use Symfony\AI\Platform\Bridge\OpenAi\TokenUsageExtractor;
1616
use Symfony\AI\Platform\Message\Message;
1717
use Symfony\AI\Platform\Message\MessageBag;
18+
use Symfony\AI\Platform\Result\TokenUsage\TokenUsageOutputProcessor;
1819

1920
require_once dirname(__DIR__).'/bootstrap.php';
2021

@@ -23,18 +24,26 @@
2324
'temperature' => 0.5, // default options for the model
2425
]);
2526

26-
$agent = new Agent($platform, $model, outputProcessors: [new TokenOutputProcessor()], logger: logger());
27+
$agent = new Agent(
28+
$platform,
29+
$model,
30+
outputProcessors: [new TokenUsageOutputProcessor(new TokenUsageExtractor())],
31+
logger: logger()
32+
);
2733
$messages = new MessageBag(
2834
Message::forSystem('You are a pirate and you write funny.'),
2935
Message::ofUser('What is the Symfony framework?'),
3036
);
37+
3138
$result = $agent->call($messages, [
3239
'max_tokens' => 500, // specific options just for this call
3340
]);
3441

35-
$metadata = $result->getMetadata();
42+
if (null === $tokenUsage = $result->getTokenUsage()) {
43+
throw new RuntimeException('Token usage is not available.');
44+
}
3645

37-
echo 'Utilized Tokens: '.$metadata['total_tokens'].\PHP_EOL;
38-
echo '-- Prompt Tokens: '.$metadata['prompt_tokens'].\PHP_EOL;
39-
echo '-- Completion Tokens: '.$metadata['completion_tokens'].\PHP_EOL;
40-
echo 'Remaining Tokens: '.$metadata['remaining_tokens'].\PHP_EOL;
46+
echo 'Utilized Tokens: '.$tokenUsage->total.\PHP_EOL;
47+
echo '-- Prompt Tokens: '.$tokenUsage->prompt.\PHP_EOL;
48+
echo '-- Completion Tokens: '.$tokenUsage->completion.\PHP_EOL;
49+
echo 'Remaining Tokens: '.$tokenUsage->remaining.\PHP_EOL;

‎src/ai-bundle/config/options.php

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,7 @@
103103
->info('Include tool definitions at the end of the system prompt')
104104
->defaultFalse()
105105
->end()
106+
->booleanNode('token_usage')->defaultFalse()->end()
106107
->arrayNode('tools')
107108
->addDefaultsIfNotSet()
108109
->treatFalseLike(['enabled' => false])

‎src/ai-bundle/src/AiBundle.php

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
use Symfony\AI\Agent\Toolbox\Tool\Agent as AgentTool;
2222
use Symfony\AI\Agent\Toolbox\ToolFactory\ChainFactory;
2323
use Symfony\AI\Agent\Toolbox\ToolFactory\MemoryToolFactory;
24+
use Symfony\AI\AiBundle\DependencyInjection\Compiler\RegisterTokenUsageExtractorPass;
2425
use Symfony\AI\AiBundle\Exception\InvalidArgumentException;
2526
use Symfony\AI\AiBundle\Profiler\TraceablePlatform;
2627
use Symfony\AI\AiBundle\Profiler\TraceableToolbox;
@@ -38,6 +39,8 @@
3839
use Symfony\AI\Platform\ModelClientInterface;
3940
use Symfony\AI\Platform\Platform;
4041
use Symfony\AI\Platform\PlatformInterface;
42+
use Symfony\AI\Platform\Result\TokenUsage\TokenUsageExtractorInterface;
43+
use Symfony\AI\Platform\Result\TokenUsage\TokenUsageOutputProcessor;
4144
use Symfony\AI\Platform\ResultConverterInterface;
4245
use Symfony\AI\Store\Bridge\Azure\SearchStore as AzureSearchStore;
4346
use Symfony\AI\Store\Bridge\ChromaDb\Store as ChromaDbStore;
@@ -67,7 +70,6 @@
6770
use Symfony\Component\HttpKernel\Bundle\AbstractBundle;
6871
use Symfony\Component\Security\Core\Authorization\AuthorizationCheckerInterface;
6972
use Symfony\Contracts\HttpClient\HttpClientInterface;
70-
7173
use function Symfony\Component\String\u;
7274

7375
/**
@@ -141,6 +143,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
141143
->addTag('ai.platform.model_client');
142144
$builder->registerForAutoconfiguration(ResultConverterInterface::class)
143145
->addTag('ai.platform.result_converter');
146+
$builder->registerForAutoconfiguration(TokenUsageExtractorInterface::class)
147+
->addTag('ai.platform.token_usage_extractor');
144148

145149
if (!ContainerBuilder::willBeAvailable('symfony/security-core', AuthorizationCheckerInterface::class, ['symfony/ai-bundle'])) {
146150
$builder->removeDefinition('ai.security.is_granted_attribute_listener');
@@ -156,6 +160,13 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
156160
}
157161
}
158162

163+
public function build(ContainerBuilder $container): void
164+
{
165+
parent::build($container);
166+
167+
$container->addCompilerPass(new RegisterTokenUsageExtractorPass());
168+
}
169+
159170
/**
160171
* @param array<string, mixed> $platform
161172
*/
@@ -384,7 +395,7 @@ private function processAgentConfig(string $name, array $config, ContainerBuilde
384395
$tool['service'] = \sprintf('ai.agent.%s', $tool['agent']);
385396
}
386397
$reference = new Reference($tool['service']);
387-
// We use the memory factory in case method, description and name are set
398+
// We use the memory factory in case, method, description and name are set
388399
if (isset($tool['name'], $tool['description'])) {
389400
if (isset($tool['agent'])) {
390401
$agentWrapperDefinition = new Definition(AgentTool::class, [$reference]);
@@ -434,6 +445,24 @@ private function processAgentConfig(string $name, array $config, ContainerBuilde
434445
}
435446
}
436447

448+
// TOKEN USAGE
449+
if ($config['token_usage'] ?? false) {
450+
$platformServiceId = $config['platform'];
451+
$platformName = str_replace('ai.platform.', '', $platformServiceId);
452+
$extractorAlias = \sprintf('ai.platform.token_usage_extractor.%s', $platformName);
453+
454+
if (!$container->hasAlias($extractorAlias)) {
455+
throw new InvalidArgumentException(\sprintf('Token usage is enabled for agent "%s", but no token usage extractor is registered for platform "%s".', $name, $platformName));
456+
}
457+
458+
$processorId = \sprintf('ai.token_usage_output_processor.%s', $name);
459+
$processorDefinition = (new Definition(TokenUsageOutputProcessor::class))
460+
->addArgument(new Reference($extractorAlias));
461+
462+
$container->setDefinition($processorId, $processorDefinition);
463+
$outputProcessors[] = new Reference($processorId);
464+
}
465+
437466
// STRUCTURED OUTPUT
438467
if ($config['structured_output']) {
439468
$inputProcessors[] = new Reference('ai.agent.structured_output_processor');

‎src/ai-bundle/src/DependencyInjection/Compiler/RegisterTokenUsageExtractorPass.php

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
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\AiBundle\DependencyInjection\Compiler;
13+
14+
use Symfony\AI\Platform\Exception\RuntimeException;
15+
use Symfony\AI\Platform\Result\TokenUsage\Attribute\AsTokenUsageExtractor;
16+
use Symfony\Component\DependencyInjection\Compiler\CompilerPassInterface;
17+
use Symfony\Component\DependencyInjection\ContainerBuilder;
18+
19+
class RegisterTokenUsageExtractorPass implements CompilerPassInterface
20+
{
21+
public function process(ContainerBuilder $container): void
22+
{
23+
foreach ($container->findTaggedServiceIds('ai.platform.token_usage_extractor') as $serviceId => $tags) {
24+
$serviceDefinition = $container->getDefinition($serviceId);
25+
$serviceClass = $serviceDefinition->getClass();
26+
27+
if (!class_exists($serviceClass)) {
28+
continue;
29+
}
30+
31+
$reflectionClass = new \ReflectionClass($serviceClass);
32+
$attributes = $reflectionClass->getAttributes(AsTokenUsageExtractor::class);
33+
34+
if (0 === \count($attributes)) {
35+
throw new RuntimeException(\sprintf('Service "%s" is tagged as "ai.platform.token_usage_extractor" but does not have the "%s" attribute.', $serviceId, AsTokenUsageExtractor::class));
36+
}
37+
38+
foreach ($attributes as $attribute) {
39+
$platform = $attribute->newInstance()->platform;
40+
$alias = \sprintf('ai.platform.token_usage_extractor.%s', $platform);
41+
$container->setAlias($alias, $serviceId);
42+
}
43+
}
44+
}
45+
}

‎src/platform/src/Bridge/Mistral/TokenOutputProcessor.php

Lines changed: 0 additions & 58 deletions
This file was deleted.

‎src/platform/src/Bridge/Mistral/TokenUsageExtractor.php

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
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\Mistral;
13+
14+
use Symfony\AI\Agent\Output;
15+
use Symfony\AI\Platform\Result\StreamResult;
16+
use Symfony\AI\Platform\Result\TokenUsage\Attribute\AsTokenUsageExtractor;
17+
use Symfony\AI\Platform\Result\TokenUsage\TokenUsage;
18+
use Symfony\AI\Platform\Result\TokenUsage\TokenUsageExtractorInterface;
19+
use Symfony\Contracts\HttpClient\ResponseInterface;
20+
21+
/**
22+
* @author Junaid Farooq <[email protected]>
23+
*/
24+
#[AsTokenUsageExtractor(platform: 'mistral')]
25+
class TokenUsageExtractor implements TokenUsageExtractorInterface
26+
{
27+
public function extractTokenUsage(Output $output): ?TokenUsage
28+
{
29+
if ($output->result instanceof StreamResult) {
30+
return null;
31+
}
32+
33+
$rawResponse = $output->result->getRawResult()?->getObject();
34+
35+
if (!$rawResponse instanceof ResponseInterface) {
36+
return null;
37+
}
38+
39+
$remainingTokensMinute = $rawResponse->getHeaders(false)['x-ratelimit-limit-tokens-minute'][0] ?? null;
40+
$remainingTokensMonth = $rawResponse->getHeaders(false)['x-ratelimit-limit-tokens-month'][0] ?? null;
41+
42+
$tokenUsage = new TokenUsage(
43+
remainingTokensMinute: null !== $remainingTokensMinute ? (int) $remainingTokensMinute : null,
44+
remainingTokensMonth: null !== $remainingTokensMonth ? (int) $remainingTokensMonth : null,
45+
);
46+
47+
$data = $rawResponse->toArray(false);
48+
$usage = $data['usage'] ?? null;
49+
50+
if (null === $usage) {
51+
return $tokenUsage;
52+
}
53+
54+
$tokenUsage->prompt = $usage['prompt_tokens'] ?? null;
55+
$tokenUsage->completion = $usage['completion_tokens'] ?? null;
56+
$tokenUsage->total = $usage['total_tokens'] ?? null;
57+
58+
return $tokenUsage;
59+
}
60+
}

0 commit comments

Comments
 (0)