[Platform] Add TokenOutputProcessor for Mistral

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

@@ -0,0 +1,58 @@
+<?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\Mistral;
+
+use Symfony\AI\Agent\Output;
+use Symfony\AI\Agent\OutputProcessorInterface;
+use Symfony\AI\Platform\Response\StreamResponse;
+use Symfony\Contracts\HttpClient\ResponseInterface;
+
+/**
+ * @author Quentin Fahrner <fahrner.quentin@gmail.com>
+ */
+final class TokenOutputProcessor implements OutputProcessorInterface
+{
+    public function processOutput(Output $output): void
+    {
+        if ($output->response instanceof StreamResponse) {
+            // Streams have to be handled manually as the tokens are part of the streamed chunks
+            return;
+        }
+
+        $rawResponse = $output->response->getRawResponse()?->getRawObject();
+        if (!$rawResponse instanceof ResponseInterface) {
+            return;
+        }
+
+        $metadata = $output->response->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/platform/tests/Bridge/Mistral/TokenOutputProcessorTest.php

@@ -0,0 +1,161 @@
+<?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 Bridge\Mistral;
+
+use PHPUnit\Framework\Attributes\CoversClass;
+use PHPUnit\Framework\Attributes\Small;
+use PHPUnit\Framework\Attributes\Test;
+use PHPUnit\Framework\Attributes\UsesClass;
+use PHPUnit\Framework\TestCase;
+use Symfony\AI\Agent\Output;
+use Symfony\AI\Platform\Bridge\Mistral\TokenOutputProcessor;
+use Symfony\AI\Platform\Message\MessageBagInterface;
+use Symfony\AI\Platform\Model;
+use Symfony\AI\Platform\Response\Metadata\Metadata;
+use Symfony\AI\Platform\Response\RawHttpResponse;
+use Symfony\AI\Platform\Response\ResponseInterface;
+use Symfony\AI\Platform\Response\StreamResponse;
+use Symfony\AI\Platform\Response\TextResponse;
+use Symfony\Contracts\HttpClient\ResponseInterface as SymfonyHttpResponse;
+
+#[CoversClass(TokenOutputProcessor::class)]
+#[UsesClass(Output::class)]
+#[UsesClass(TextResponse::class)]
+#[UsesClass(StreamResponse::class)]
+#[UsesClass(Metadata::class)]
+#[Small]
+final class TokenOutputProcessorTest extends TestCase
+{
+    #[Test]
+    public function itHandlesStreamResponsesWithoutProcessing(): void
+    {
+        $processor = new TokenOutputProcessor();
+        $streamResponse = new StreamResponse((static function () { yield 'test'; })());
+        $output = $this->createOutput($streamResponse);
+
+        $processor->processOutput($output);
+
+        $metadata = $output->response->getMetadata();
+        self::assertCount(0, $metadata);
+    }
+
+    #[Test]
+    public function itDoesNothingWithoutRawResponse(): void
+    {
+        $processor = new TokenOutputProcessor();
+        $textResponse = new TextResponse('test');
+        $output = $this->createOutput($textResponse);
+
+        $processor->processOutput($output);
+
+        $metadata = $output->response->getMetadata();
+        self::assertCount(0, $metadata);
+    }
+
+    #[Test]
+    public function itAddsRemainingTokensToMetadata(): void
+    {
+        $processor = new TokenOutputProcessor();
+        $textResponse = new TextResponse('test');
+
+        $textResponse->setRawResponse($this->createRawResponse());
+
+        $output = $this->createOutput($textResponse);
+
+        $processor->processOutput($output);
+
+        $metadata = $output->response->getMetadata();
+        self::assertCount(2, $metadata);
+        self::assertSame(1000, $metadata->get('remaining_tokens_minute'));
+        self::assertSame(1000000, $metadata->get('remaining_tokens_month'));
+    }
+
+    #[Test]
+    public function itAddsUsageTokensToMetadata(): void
+    {
+        $processor = new TokenOutputProcessor();
+        $textResponse = new TextResponse('test');
+
+        $rawResponse = $this->createRawResponse([
+            'usage' => [
+                'prompt_tokens' => 10,
+                'completion_tokens' => 20,
+                'total_tokens' => 30,
+            ],
+        ]);
+
+        $textResponse->setRawResponse($rawResponse);
+
+        $output = $this->createOutput($textResponse);
+
+        $processor->processOutput($output);
+
+        $metadata = $output->response->getMetadata();
+        self::assertCount(5, $metadata);
+        self::assertSame(1000, $metadata->get('remaining_tokens_minute'));
+        self::assertSame(1000000, $metadata->get('remaining_tokens_month'));
+        self::assertSame(10, $metadata->get('prompt_tokens'));
+        self::assertSame(20, $metadata->get('completion_tokens'));
+        self::assertSame(30, $metadata->get('total_tokens'));
+    }
+
+    #[Test]
+    public function itHandlesMissingUsageFields(): void
+    {
+        $processor = new TokenOutputProcessor();
+        $textResponse = new TextResponse('test');
+
+        $rawResponse = $this->createRawResponse([
+            'usage' => [
+                // Missing some fields
+                'prompt_tokens' => 10,
+            ],
+        ]);
+
+        $textResponse->setRawResponse($rawResponse);
+
+        $output = $this->createOutput($textResponse);
+
+        $processor->processOutput($output);
+
+        $metadata = $output->response->getMetadata();
+        self::assertCount(5, $metadata);
+        self::assertSame(1000, $metadata->get('remaining_tokens_minute'));
+        self::assertSame(1000000, $metadata->get('remaining_tokens_month'));
+        self::assertSame(10, $metadata->get('prompt_tokens'));
+        self::assertNull($metadata->get('completion_tokens'));
+        self::assertNull($metadata->get('total_tokens'));
+    }
+
+    private function createRawResponse(array $data = []): RawHttpResponse
+    {
+        $rawResponse = self::createStub(SymfonyHttpResponse::class);
+        $rawResponse->method('getHeaders')->willReturn([
+            'x-ratelimit-limit-tokens-minute' => ['1000'],
+            'x-ratelimit-limit-tokens-month' => ['1000000'],
+        ]);
+
+        $rawResponse->method('toArray')->willReturn($data);
+
+        return new RawHttpResponse($rawResponse);
+    }
+
+    private function createOutput(ResponseInterface $response): Output
+    {
+        return new Output(
+            self::createStub(Model::class),
+            $response,
+            self::createStub(MessageBagInterface::class),
+            [],
+        );
+    }
+}