[Platform] Standardise token usage

examples/mistral/token-metadata.php

@@ -0,0 +1,48 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+use Symfony\AI\Agent\Agent;
+use Symfony\AI\Platform\Bridge\Mistral\Mistral;
+use Symfony\AI\Platform\Bridge\Mistral\PlatformFactory;
+use Symfony\AI\Platform\Bridge\Mistral\TokenUsageExtractor;
+use Symfony\AI\Platform\Message\Message;
+use Symfony\AI\Platform\Message\MessageBag;
+use Symfony\AI\Platform\Result\TokenUsage\TokenUsageOutputProcessor;
+
+require_once dirname(__DIR__).'/bootstrap.php';
+
+$platform = PlatformFactory::create(env('MISTRAL_API_KEY'), http_client());
+$model = new Mistral(Mistral::MISTRAL_SMALL, [
+    'temperature' => 0.5, // default options for the model
+]);
+$agent = new Agent(
+    $platform,
+    $model,
+    outputProcessors: [new TokenUsageOutputProcessor(new TokenUsageExtractor())],
+    logger: logger()
+);
+
+$messages = new MessageBag(
+    Message::forSystem('You are a pirate and you write funny.'),
+    Message::ofUser('What is the Symfony framework?'),
+);
+
+$result = $agent->call($messages, [
+    'max_tokens' => 500,
+]);
+
+if (null === $tokenUsage = $result->getTokenUsage()) {
+    throw new RuntimeException('Token usage is not available.');
+}
+
+echo 'Utilized Tokens: '.$tokenUsage->total.\PHP_EOL;
+echo '-- Prompt Tokens: '.$tokenUsage->prompt.\PHP_EOL;
+echo '-- Completion Tokens: '.$tokenUsage->completion.\PHP_EOL;

examples/openai/token-metadata.php

@@ -12,9 +12,10 @@
 use Symfony\AI\Agent\Agent;
 use Symfony\AI\Platform\Bridge\OpenAi\Gpt;
 use Symfony\AI\Platform\Bridge\OpenAi\PlatformFactory;
-use Symfony\AI\Platform\Bridge\OpenAi\TokenOutputProcessor;
+use Symfony\AI\Platform\Bridge\OpenAi\TokenUsageExtractor;
 use Symfony\AI\Platform\Message\Message;
 use Symfony\AI\Platform\Message\MessageBag;
+use Symfony\AI\Platform\Result\TokenUsage\TokenUsageOutputProcessor;
 
 require_once dirname(__DIR__).'/bootstrap.php';
 
@@ -23,18 +24,26 @@
     'temperature' => 0.5, // default options for the model
 ]);
 
-$agent = new Agent($platform, $model, outputProcessors: [new TokenOutputProcessor()], logger: logger());
+$agent = new Agent(
+    $platform,
+    $model,
+    outputProcessors: [new TokenUsageOutputProcessor(new TokenUsageExtractor())],
+    logger: logger()
+);
 $messages = new MessageBag(
     Message::forSystem('You are a pirate and you write funny.'),
     Message::ofUser('What is the Symfony framework?'),
 );
+
 $result = $agent->call($messages, [
     'max_tokens' => 500, // specific options just for this call
 ]);
 
-$metadata = $result->getMetadata();
+if (null === $tokenUsage = $result->getTokenUsage()) {
+    throw new RuntimeException('Token usage is not available.');
+}
 
-echo 'Utilized Tokens: '.$metadata['total_tokens'].\PHP_EOL;
-echo '-- Prompt Tokens: '.$metadata['prompt_tokens'].\PHP_EOL;
-echo '-- Completion Tokens: '.$metadata['completion_tokens'].\PHP_EOL;
-echo 'Remaining Tokens: '.$metadata['remaining_tokens'].\PHP_EOL;
+echo 'Utilized Tokens: '.$tokenUsage->total.\PHP_EOL;
+echo '-- Prompt Tokens: '.$tokenUsage->prompt.\PHP_EOL;
+echo '-- Completion Tokens: '.$tokenUsage->completion.\PHP_EOL;
+echo 'Remaining Tokens: '.$tokenUsage->remaining.\PHP_EOL;

src/ai-bundle/config/options.php

@@ -103,6 +103,7 @@
                             ->info('Include tool definitions at the end of the system prompt')
                             ->defaultFalse()
                         ->end()
+                        ->booleanNode('token_usage')->defaultFalse()->end()
                         ->arrayNode('tools')
                             ->addDefaultsIfNotSet()
                             ->treatFalseLike(['enabled' => false])

src/ai-bundle/src/AiBundle.php

@@ -21,6 +21,7 @@
 use Symfony\AI\Agent\Toolbox\Tool\Agent as AgentTool;
 use Symfony\AI\Agent\Toolbox\ToolFactory\ChainFactory;
 use Symfony\AI\Agent\Toolbox\ToolFactory\MemoryToolFactory;
+use Symfony\AI\AiBundle\DependencyInjection\Compiler\RegisterTokenUsageExtractorPass;
 use Symfony\AI\AiBundle\Exception\InvalidArgumentException;
 use Symfony\AI\AiBundle\Profiler\TraceablePlatform;
 use Symfony\AI\AiBundle\Profiler\TraceableToolbox;
@@ -38,6 +39,8 @@
 use Symfony\AI\Platform\ModelClientInterface;
 use Symfony\AI\Platform\Platform;
 use Symfony\AI\Platform\PlatformInterface;
+use Symfony\AI\Platform\Result\TokenUsage\TokenUsageExtractorInterface;
+use Symfony\AI\Platform\Result\TokenUsage\TokenUsageOutputProcessor;
 use Symfony\AI\Platform\ResultConverterInterface;
 use Symfony\AI\Store\Bridge\Azure\SearchStore as AzureSearchStore;
 use Symfony\AI\Store\Bridge\ChromaDb\Store as ChromaDbStore;
@@ -139,6 +142,8 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
             ->addTag('ai.platform.model_client');
         $builder->registerForAutoconfiguration(ResultConverterInterface::class)
             ->addTag('ai.platform.result_converter');
+        $builder->registerForAutoconfiguration(TokenUsageExtractorInterface::class)
+            ->addTag('ai.platform.token_usage_extractor');
 
         if (!ContainerBuilder::willBeAvailable('symfony/security-core', AuthorizationCheckerInterface::class, ['symfony/ai-bundle'])) {
             $builder->removeDefinition('ai.security.is_granted_attribute_listener');
@@ -154,6 +159,13 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
         }
     }
 
+    public function build(ContainerBuilder $container): void
+    {
+        parent::build($container);
+
+        $container->addCompilerPass(new RegisterTokenUsageExtractorPass());
+    }
+
     /**
      * @param array<string, mixed> $platform
      */
@@ -382,7 +394,7 @@ private function processAgentConfig(string $name, array $config, ContainerBuilde
                         $tool['service'] = \sprintf('ai.agent.%s', $tool['agent']);
                     }
                     $reference = new Reference($tool['service']);
-                    // We use the memory factory in case method, description and name are set
+                    // We use the memory factory in case, method, description and name are set
                     if (isset($tool['name'], $tool['description'])) {
                         if (isset($tool['agent'])) {
                             $agentWrapperDefinition = new Definition(AgentTool::class, [$reference]);
@@ -432,6 +444,24 @@ private function processAgentConfig(string $name, array $config, ContainerBuilde
             }
         }
 
+        // TOKEN USAGE
+        if ($config['token_usage'] ?? false) {
+            $platformServiceId = $config['platform'];
+            $platformName = str_replace('ai.platform.', '', $platformServiceId);
+            $extractorAlias = \sprintf('ai.platform.token_usage_extractor.%s', $platformName);
+
+            if (!$container->hasAlias($extractorAlias)) {
+                throw new InvalidArgumentException(\sprintf('Token usage is enabled for agent "%s", but no token usage extractor is registered for platform "%s".', $name, $platformName));
+            }
+
+            $processorId = \sprintf('ai.token_usage_output_processor.%s', $name);
+            $processorDefinition = (new Definition(TokenUsageOutputProcessor::class))
+                ->addArgument(new Reference($extractorAlias));
+
+            $container->setDefinition($processorId, $processorDefinition);
+            $outputProcessors[] = new Reference($processorId);
+        }
+
         // STRUCTURED OUTPUT
         if ($config['structured_output']) {
             $inputProcessors[] = new Reference('ai.agent.structured_output_processor');

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

@@ -0,0 +1,48 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\AiBundle\DependencyInjection\Compiler;
+
+use Symfony\AI\Platform\Exception\RuntimeException;
+use Symfony\AI\Platform\Result\TokenUsage\Attribute\AsTokenUsageExtractor;
+use Symfony\Component\DependencyInjection\Compiler\CompilerPassInterface;
+use Symfony\Component\DependencyInjection\ContainerBuilder;
+
+/**
+ * @author Junaid Farooq <[email protected]>
+ */
+class RegisterTokenUsageExtractorPass implements CompilerPassInterface
+{
+    public function process(ContainerBuilder $container): void
+    {
+        foreach ($container->findTaggedServiceIds('ai.platform.token_usage_extractor') as $serviceId => $tags) {
+            $serviceDefinition = $container->getDefinition($serviceId);
+            $serviceClass = $serviceDefinition->getClass();
+
+            if (!class_exists($serviceClass)) {
+                continue;
+            }
+
+            $reflectionClass = new \ReflectionClass($serviceClass);
+            $attributes = $reflectionClass->getAttributes(AsTokenUsageExtractor::class);
+
+            if (0 === \count($attributes)) {
+                throw new RuntimeException(\sprintf('Service "%s" is tagged as "ai.platform.token_usage_extractor" but does not have the "%s" attribute.', $serviceId, AsTokenUsageExtractor::class));
+            }
+
+            foreach ($attributes as $attribute) {
+                $platform = $attribute->newInstance()->platform;
+                $alias = \sprintf('ai.platform.token_usage_extractor.%s', $platform);
+                $container->setAlias($alias, $serviceId);
+            }
+        }
+    }
+}

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

@@ -0,0 +1,60 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\Platform\Bridge\Mistral;
+
+use Symfony\AI\Agent\Output;
+use Symfony\AI\Platform\Result\StreamResult;
+use Symfony\AI\Platform\Result\TokenUsage\Attribute\AsTokenUsageExtractor;
+use Symfony\AI\Platform\Result\TokenUsage\TokenUsage;
+use Symfony\AI\Platform\Result\TokenUsage\TokenUsageExtractorInterface;
+use Symfony\Contracts\HttpClient\ResponseInterface;
+
+/**
+ * @author Junaid Farooq <[email protected]>
+ */
+#[AsTokenUsageExtractor(platform: 'mistral')]
+class TokenUsageExtractor implements TokenUsageExtractorInterface
+{
+    public function extractTokenUsage(Output $output): ?TokenUsage
+    {
+        if ($output->result instanceof StreamResult) {
+            return null;
+        }
+
+        $rawResponse = $output->result->getRawResult()?->getObject();
+
+        if (!$rawResponse instanceof ResponseInterface) {
+            return null;
+        }
+
+        $remainingTokensMinute = $rawResponse->getHeaders(false)['x-ratelimit-limit-tokens-minute'][0] ?? null;
+        $remainingTokensMonth = $rawResponse->getHeaders(false)['x-ratelimit-limit-tokens-month'][0] ?? null;
+
+        $tokenUsage = new TokenUsage(
+            remainingTokensMinute: null !== $remainingTokensMinute ? (int) $remainingTokensMinute : null,
+            remainingTokensMonth: null !== $remainingTokensMonth ? (int) $remainingTokensMonth : null,
+        );
+
+        $data = $rawResponse->toArray(false);
+        $usage = $data['usage'] ?? null;
+
+        if (null === $usage) {
+            return $tokenUsage;
+        }
+
+        $tokenUsage->prompt = $usage['prompt_tokens'] ?? null;
+        $tokenUsage->completion = $usage['completion_tokens'] ?? null;
+        $tokenUsage->total = $usage['total_tokens'] ?? null;
+
+        return $tokenUsage;
+    }
+}