Revert "feat(platform): Standardize token usage"

This reverts commit 1f1e6be.

Author: junaidbinfarooq

examples/mistral/token-metadata.php

@@ -12,10 +12,9 @@
 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\TokenUsageExtractor;
+use Symfony\AI\Platform\Bridge\OpenAi\TokenOutputProcessor;
 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';
 
@@ -24,26 +23,18 @@
     'temperature' => 0.5, // default options for the model
 ]);
 
-$agent = new Agent(
-    $platform,
-    $model,
-    outputProcessors: [new TokenUsageOutputProcessor(new TokenUsageExtractor())],
-    logger: logger()
-);
+$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 Symfony framework?'),
 );
-
 $result = $agent->call($messages, [
     'max_tokens' => 500, // specific options just for this call
 ]);
 
-if (null === $tokenUsage = $result->getTokenUsage()) {
-    throw new RuntimeException('Token usage is not available.');
-}
+$metadata = $result->getMetadata();
 
-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;
+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;

examples/openai/token-metadata.php

@@ -103,7 +103,6 @@
                             ->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/config/options.php

@@ -21,7 +21,6 @@
 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;
@@ -39,8 +38,6 @@
 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;
@@ -142,8 +139,6 @@ 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');
@@ -159,13 +154,6 @@ 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
      */
@@ -394,7 +382,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]);
@@ -444,24 +432,6 @@ 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/AiBundle.php

@@ -0,0 +1,58 @@
+<?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\Agent\OutputProcessorInterface;
+use Symfony\AI\Platform\Result\StreamResult;
+use Symfony\Contracts\HttpClient\ResponseInterface;
+
+/**
+ * @author Quentin Fahrner <[email protected]>
+ */
+final class TokenOutputProcessor implements OutputProcessorInterface
+{
+    public function processOutput(Output $output): void
+    {
+        if ($output->result instanceof StreamResult) {
+            // Streams have to be handled manually as the tokens are part of the streamed chunks
+            return;
+        }
+
+        $rawResponse = $output->result->getRawResult()?->getObject();
+        if (!$rawResponse instanceof ResponseInterface) {
+            return;
+        }
+
+        $metadata = $output->result->getMetadata();
+
+        $metadata->add(
+            'remaining_tokens_minute',
+            (int) $rawResponse->getHeaders(false)['x-ratelimit-limit-tokens-minute'][0],
+        );
+
+        $metadata->add(
+            'remaining_tokens_month',
+            (int) $rawResponse->getHeaders(false)['x-ratelimit-limit-tokens-month'][0],
+        );
+
+        $content = $rawResponse->toArray(false);
+
+        if (!\array_key_exists('usage', $content)) {
+            return;
+        }
+
+        $metadata->add('prompt_tokens', $content['usage']['prompt_tokens'] ?? null);
+        $metadata->add('completion_tokens', $content['usage']['completion_tokens'] ?? null);
+        $metadata->add('total_tokens', $content['usage']['total_tokens'] ?? null);
+    }
+}

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

@@ -0,0 +1,53 @@
+<?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\OpenAi;
+
+use Symfony\AI\Agent\Output;
+use Symfony\AI\Agent\OutputProcessorInterface;
+use Symfony\AI\Platform\Result\StreamResult;
+use Symfony\Contracts\HttpClient\ResponseInterface;
+
+/**
+ * @author Denis Zunke <[email protected]>
+ */
+final class TokenOutputProcessor implements OutputProcessorInterface
+{
+    public function processOutput(Output $output): void
+    {
+        if ($output->result instanceof StreamResult) {
+            // Streams have to be handled manually as the tokens are part of the streamed chunks
+            return;
+        }
+
+        $rawResponse = $output->result->getRawResult()?->getObject();
+        if (!$rawResponse instanceof ResponseInterface) {
+            return;
+        }
+
+        $metadata = $output->result->getMetadata();
+
+        $metadata->add(
+            'remaining_tokens',
+            (int) $rawResponse->getHeaders(false)['x-ratelimit-remaining-tokens'][0],
+        );
+
+        $content = $rawResponse->toArray(false);
+
+        if (!\array_key_exists('usage', $content)) {
+            return;
+        }
+
+        $metadata->add('prompt_tokens', $content['usage']['prompt_tokens'] ?? null);
+        $metadata->add('completion_tokens', $content['usage']['completion_tokens'] ?? null);
+        $metadata->add('total_tokens', $content['usage']['total_tokens'] ?? null);
+    }
+}