feature #329 [AI Bundle] Activate token usage output processors (junaidbinfarooq)

This PR was merged into the main branch.

Discussion
----------

[AI Bundle] Activate token usage output processors

| Q             | A
| ------------- | ---
| Bug fix?      | no
| New feature?  | yes
| Docs?         | no
| Issues        | Fix #208 Continues #311
| License       | MIT

# Changes proposed:
- Adds support to automatically configure the token usage processors so that they are available to end users (devs) when needed
- Makes the availability of token usage configurable and enabled by default

Commits
-------

b426948 feat(ai-bundle): Auto-configure token usage processors - Adds support to automatically configure the token usage processors so that they are available to end users (devs) when needed - Makes the availability of token usage configurable and disabled by default

Author: chr-hertel

examples/bootstrap.php

@@ -11,6 +11,8 @@
 
 use Psr\Log\LoggerAwareInterface;
 use Psr\Log\LoggerInterface;
+use Symfony\AI\Platform\Metadata\Metadata;
+use Symfony\AI\Platform\Metadata\TokenUsage;
 use Symfony\Component\Console\Logger\ConsoleLogger;
 use Symfony\Component\Console\Output\ConsoleOutput;
 use Symfony\Component\Dotenv\Dotenv;
@@ -52,3 +54,19 @@ function logger(): LoggerInterface
 
     return new ConsoleLogger(new ConsoleOutput($verbosity));
 }
+
+function print_token_usage(Metadata $metadata): void
+{
+    $tokenUsage = $metadata->get('token_usage');
+
+    assert($tokenUsage instanceof TokenUsage);
+
+    echo 'Prompt tokens: '.$tokenUsage->promptTokens.\PHP_EOL;
+    echo 'Completion tokens: '.$tokenUsage->completionTokens.\PHP_EOL;
+    echo 'Thinking tokens: '.$tokenUsage->thinkingTokens.\PHP_EOL;
+    echo 'Cached tokens: '.$tokenUsage->cachedTokens.\PHP_EOL;
+    echo 'Remaining tokens minute: '.$tokenUsage->remainingTokensMinute.\PHP_EOL;
+    echo 'Remaining tokens month: '.$tokenUsage->remainingTokensMonth.\PHP_EOL;
+    echo 'Remaining tokens: '.$tokenUsage->remainingTokens.\PHP_EOL;
+    echo 'Utilized tokens: '.$tokenUsage->totalTokens.\PHP_EOL;
+}

examples/mistral/token-metadata.php

@@ -0,0 +1,36 @@
+<?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\TokenOutputProcessor;
+use Symfony\AI\Platform\Message\Message;
+use Symfony\AI\Platform\Message\MessageBag;
+
+require_once dirname(__DIR__).'/bootstrap.php';
+
+$platform = PlatformFactory::create(env('MISTRAL_API_KEY'), http_client());
+$model = new Mistral();
+
+$agent = new Agent($platform, $model, outputProcessors: [new TokenOutputProcessor()], logger: logger());
+
+$messages = new MessageBag(
+    Message::forSystem('You are a pirate and you write funny.'),
+    Message::ofUser('What is the best French cuisine?'),
+);
+
+$result = $agent->call($messages, [
+    'temperature' => 0.3,
+    'max_tokens' => 500,
+]);
+
+print_token_usage($result->getMetadata());

examples/openai/token-metadata.php

@@ -15,7 +15,6 @@
 use Symfony\AI\Platform\Bridge\OpenAi\TokenOutputProcessor;
 use Symfony\AI\Platform\Message\Message;
 use Symfony\AI\Platform\Message\MessageBag;
-use Symfony\AI\Platform\Result\Metadata\TokenUsage;
 
 require_once dirname(__DIR__).'/bootstrap.php';
 
@@ -33,12 +32,4 @@
     'max_tokens' => 500, // specific options just for this call
 ]);
 
-$metadata = $result->getMetadata();
-$tokenUsage = $metadata->get('token_usage');
-
-assert($tokenUsage instanceof TokenUsage);
-
-echo 'Utilized Tokens: '.$tokenUsage->totalTokens.\PHP_EOL;
-echo '-- Prompt Tokens: '.$tokenUsage->promptTokens.\PHP_EOL;
-echo '-- Completion Tokens: '.$tokenUsage->completionTokens.\PHP_EOL;
-echo 'Remaining Tokens: '.$tokenUsage->remainingTokens.\PHP_EOL;
+print_token_usage($result->getMetadata());

src/ai-bundle/CHANGELOG.md

@@ -25,4 +25,9 @@ CHANGELOG
    - Factory services for creating platforms
  * Add bundle configuration with semantic validation
  * Add support for fault-tolerant tool execution
- * Add structured output configuration support
+ * Add structured output configuration support
+ * Add token usage tracking:
+   - `track_token_usage` option for agents to monitor AI model consumption
+   - Automatic registration of token output processors for Mistral, OpenAI and Vertex AI
+   - Token usage metadata in agent results including prompt, completion, total, cached, and thinking tokens
+   - Rate limit information tracking for supported platforms

src/ai-bundle/config/options.php

@@ -103,6 +103,10 @@
                             ->info('Service name of platform')
                             ->defaultValue(PlatformInterface::class)
                         ->end()
+                        ->booleanNode('track_token_usage')
+                            ->info('Enable tracking of token usage for the agent')
+                            ->defaultTrue()
+                        ->end()
                         ->arrayNode('model')
                             ->children()
                                 ->scalarNode('class')->isRequired()->end()

src/ai-bundle/config/services.php

@@ -26,9 +26,12 @@
 use Symfony\AI\AiBundle\Security\EventListener\IsGrantedToolAttributeListener;
 use Symfony\AI\Platform\Bridge\Anthropic\Contract\AnthropicContract;
 use Symfony\AI\Platform\Bridge\Gemini\Contract\GeminiContract;
+use Symfony\AI\Platform\Bridge\Mistral\TokenOutputProcessor as MistralTokenOutputProcessor;
 use Symfony\AI\Platform\Bridge\Ollama\Contract\OllamaContract;
 use Symfony\AI\Platform\Bridge\OpenAi\Contract\OpenAiContract;
+use Symfony\AI\Platform\Bridge\OpenAi\TokenOutputProcessor as OpenAiTokenOutputProcessor;
 use Symfony\AI\Platform\Bridge\VertexAi\Contract\GeminiContract as VertexAiGeminiContract;
+use Symfony\AI\Platform\Bridge\VertexAi\TokenOutputProcessor as VertexAiTokenOutputProcessor;
 use Symfony\AI\Platform\Contract;
 use Symfony\AI\Platform\Contract\JsonSchema\DescriptionParser;
 use Symfony\AI\Platform\Contract\JsonSchema\Factory as SchemaFactory;
@@ -131,5 +134,10 @@
                 service('.inner'),
             ])
             ->tag('ai.traceable_toolbox')
+
+        // token usage processors
+        ->set('ai.platform.token_usage_processor.mistral', MistralTokenOutputProcessor::class)
+        ->set('ai.platform.token_usage_processor.openai', OpenAiTokenOutputProcessor::class)
+        ->set('ai.platform.token_usage_processor.vertexai', VertexAiTokenOutputProcessor::class)
     ;
 };

src/ai-bundle/doc/index.rst

@@ -65,6 +65,7 @@ Configuration
             rag:
                 platform: 'ai.platform.azure.gpt_deployment'
                 structured_output: false # Disables support for "output_structure" option, default is true
+                track_token_usage: true # Enable tracking of token usage for the agent, default is true
                 model:
                     class: 'Symfony\AI\Platform\Bridge\OpenAi\Gpt'
                     name: !php/const Symfony\AI\Platform\Bridge\OpenAi\Gpt::GPT_4O_MINI
@@ -230,6 +231,66 @@ The attribute ``IsGrantedTool`` can be added on class- or method-level - even mu
 times. If multiple attributes apply to one tool call, a logical AND is used and all access
 decisions have to grant access.
 
+Token Usage Tracking
+--------------------
+
+Token usage tracking is a feature provided by some of the Platform's bridges, for monitoring and analyzing the
+consumption of tokens by your agents. This feature is particularly useful for understanding costs and performance.
+
+When enabled, the agent will automatically track token usage information and add it
+to the result metadata. The tracked information includes:
+
+* **Prompt tokens**: Number of tokens used in the input/prompt
+* **Completion tokens**: Number of tokens generated in the response
+* **Total tokens**: Total number of tokens used (prompt + completion)
+* **Remaining tokens**: Number of remaining tokens in rate limits (when available)
+* **Cached tokens**: Number of cached tokens used (when available)
+* **Thinking tokens**: Number of reasoning tokens used (for models that support reasoning)
+
+The token usage information can be accessed from the result metadata::
+
+    use Symfony\AI\Agent\AgentInterface;
+    use Symfony\AI\Platform\Message\Message;
+    use Symfony\AI\Platform\Message\MessageBag;
+    use Symfony\AI\Platform\Result\Metadata\TokenUsage\TokenUsage;
+
+    final readonly class MyService
+    {
+        public function __construct(
+            private AgentInterface $agent,
+        ) {
+        }
+
+        public function getTokenUsage(string $message): ?TokenUsage
+        {
+            $messages = new MessageBag(Message::ofUser($message));
+            $result = $this->agent->call($messages);
+
+            return $result->getMetadata()->get('token_usage');
+        }
+    }
+
+**Supported Platforms**
+
+Token usage tracking is currently supported, and by default enabled, for the following platforms:
+
+* **OpenAI**: Tracks all token types including cached and thinking tokens
+* **Mistral**: Tracks basic token usage and rate limit information
+
+**Disable Tracking**
+
+To disable token usage tracking for an agent, set the ``track_token_usage`` option to ``false``:
+
+.. code-block:: yaml
+
+    ai:
+        agent:
+            my_agent:
+                track_token_usage: false
+                model:
+                    class: 'Symfony\AI\Platform\Bridge\OpenAi\Gpt'
+                    name: !php/const Symfony\AI\Platform\Bridge\OpenAi\Gpt::GPT_4O_MINI
+
 Profiler
 --------
 

src/ai-bundle/src/AiBundle.php

@@ -514,6 +514,27 @@ private function processAgentConfig(string $name, array $config, ContainerBuilde
             $outputProcessors[] = new Reference('ai.agent.structured_output_processor');
         }
 
+        // TOKEN USAGE TRACKING
+        if ($config['track_token_usage'] ?? true) {
+            $platformServiceId = $config['platform'];
+
+            if ($container->hasAlias($platformServiceId)) {
+                $platformServiceId = (string) $container->getAlias($platformServiceId);
+            }
+
+            if (str_starts_with($platformServiceId, 'ai.platform.')) {
+                $platform = u($platformServiceId)->after('ai.platform.')->toString();
+
+                if (str_contains($platform, 'azure')) {
+                    $platform = 'azure';
+                }
+
+                if ($container->hasDefinition('ai.platform.token_usage_processor.'.$platform)) {
+                    $outputProcessors[] = new Reference('ai.platform.token_usage_processor.'.$platform);
+                }
+            }
+        }
+
         // SYSTEM PROMPT
         if (\is_string($config['system_prompt'])) {
             $systemPromptInputProcessorDefinition = new Definition(SystemPromptInputProcessor::class, [

src/ai-bundle/tests/DependencyInjection/AiBundleTest.php

@@ -344,6 +344,7 @@ private function getFullConfig(): array
                             ],
                         ],
                         'structured_output' => false,
+                        'track_token_usage' => true,
                         'system_prompt' => 'You are a helpful assistant.',
                         'include_tools' => true,
                         'tools' => [

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

@@ -13,7 +13,7 @@
 
 use Symfony\AI\Agent\Output;
 use Symfony\AI\Agent\OutputProcessorInterface;
-use Symfony\AI\Platform\Result\Metadata\TokenUsage;
+use Symfony\AI\Platform\Metadata\TokenUsage;
 use Symfony\AI\Platform\Result\StreamResult;
 use Symfony\Contracts\HttpClient\ResponseInterface;