refactor(platform): add TokenUsage for Ollama

OllamaClient.php

@@ -38,8 +38,8 @@ final class OllamaClient implements ModelClientInterface
     public function request(Model $model, array|string $payload, array $options = []): RawHttpResult
     {
         return match (true) {
-            \in_array(Capability::INPUT_MESSAGES, $model->getCapabilities(), true) => $this->doCompletionRequest($payload, $options),
-            \in_array(Capability::EMBEDDINGS, $model->getCapabilities(), true) => $this->doEmbeddingsRequest($model, $payload, $options),
+            $model->supports(Capability::INPUT_MESSAGES) => $this->doCompletionRequest($payload, $options),
+            $model->supports(Capability::EMBEDDINGS) => $this->doEmbeddingsRequest($model, $payload, $options),
             default => throw new InvalidArgumentException(\sprintf('Unsupported model "%s": "%s".', $model::class, $model->getName())),
         };
     }

OllamaResultConverter.php

@@ -47,9 +47,9 @@ final class OllamaResultConverter implements ResultConverterInterface
             : $this->doConvertCompletion($data);
     }
 
-    public function getTokenUsageExtractor(): ?TokenUsageExtractorInterface
+    public function getTokenUsageExtractor(): TokenUsageExtractorInterface
     {
-        return null;
+        return new TokenUsageExtractor();
     }
 
     /**

TokenUsageExtractor.php

@@ -0,0 +1,52 @@
+<?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\Ollama;
+
+use Symfony\AI\Platform\Result\RawResultInterface;
+use Symfony\AI\Platform\TokenUsage\TokenUsage;
+use Symfony\AI\Platform\TokenUsage\TokenUsageExtractorInterface;
+use Symfony\AI\Platform\TokenUsage\TokenUsageInterface;
+use Symfony\Contracts\HttpClient\ResponseInterface;
+
+/**
+ * @author Guillaume Loulier <personal@guillaumeloulier.fr>
+ */
+final class TokenUsageExtractor implements TokenUsageExtractorInterface
+{
+    public function extract(RawResultInterface $rawResult, array $options = []): ?TokenUsageInterface
+    {
+        $response = $rawResult->getObject();
+        if (!$response instanceof ResponseInterface) {
+            return null;
+        }
+
+        if ($options['stream'] ?? false) {
+            foreach ($rawResult->getDataStream() as $chunk) {
+                if ($chunk['done']) {
+                    return new TokenUsage(
+                        $chunk['prompt_eval_count'],
+                        $chunk['eval_count']
+                    );
+                }
+            }
+
+            return null;
+        }
+
+        $payload = $response->toArray();
+
+        return new TokenUsage(
+            $payload['prompt_eval_count'],
+            $payload['eval_count']
+        );
+    }
+}