bug #1267 [Platform][Ollama][VertexAi] Disable token extraction for streams (chr-hertel)

This PR was merged into the main branch.

Discussion
----------

[Platform][Ollama][VertexAi] Disable token extraction for streams

| Q             | A
| ------------- | ---
| Bug fix?      | yes
| New feature?  | no
| Docs?         | no <!-- required for new features -->
| Issues        | Fix #... <!-- prefix each issue number with "Fix #", no need to create an issue if none exist, explain below instead -->
| License       | MIT

For now we need to revert this, since it consumes the stream iterator already and output streaming is broken than.

Commits
-------

42dd99d Disable Ollama token extraction for streams

Author: chr-hertel

src/platform/src/Bridge/Ollama/Tests/TokenUsageExtractorTest.php

@@ -54,25 +54,8 @@ public function testItExtractsTokenUsageFromStreamResult()
     {
         $extractor = new TokenUsageExtractor();
 
-        $result = new InMemoryRawResult([], [
-            [
-                'model' => 'foo',
-                'response' => 'First chunk',
-                'done' => false,
-            ],
-            [
-                'model' => 'foo',
-                'response' => 'Hello World!',
-                'done' => true,
-                'prompt_eval_count' => 10,
-                'eval_count' => 10,
-            ],
-        ]);
-
-        $tokenUsage = $extractor->extract($result, ['stream' => true]);
+        $tokenUsage = $extractor->extract(new InMemoryRawResult(), ['stream' => true]);
 
-        $this->assertInstanceOf(TokenUsage::class, $tokenUsage);
-        $this->assertSame(10, $tokenUsage->getPromptTokens());
-        $this->assertSame(10, $tokenUsage->getCompletionTokens());
+        $this->assertNull($tokenUsage);
     }
 }

src/platform/src/Bridge/Ollama/TokenUsageExtractor.php

@@ -24,15 +24,7 @@ final class TokenUsageExtractor implements TokenUsageExtractorInterface
     public function extract(RawResultInterface $rawResult, array $options = []): ?TokenUsageInterface
     {
         if ($options['stream'] ?? false) {
-            foreach ($rawResult->getDataStream() as $chunk) {
-                if ($chunk['done']) {
-                    return new TokenUsage(
-                        $chunk['prompt_eval_count'],
-                        $chunk['eval_count']
-                    );
-                }
-            }
-
+            // Streams have to be handled manually as the tokens are part of the streamed chunks
             return null;
         }
 

src/platform/src/Bridge/VertexAi/Gemini/TokenUsageExtractor.php

@@ -24,19 +24,7 @@ final class TokenUsageExtractor implements TokenUsageExtractorInterface
     public function extract(RawResultInterface $rawResult, array $options = []): ?TokenUsageInterface
     {
         if ($options['stream'] ?? false) {
-            $lastChunk = null;
-
-            foreach ($rawResult->getDataStream() as $chunk) {
-                // Store last event that contains usage metadata
-                if (isset($chunk['usageMetadata'])) {
-                    $lastChunk = $chunk;
-                }
-            }
-
-            if ($lastChunk) {
-                return $this->extractUsageMetadata($lastChunk['usageMetadata']);
-            }
-
+            // Streams have to be handled manually as the tokens are part of the streamed chunks
             return null;
         }
 

src/platform/src/Bridge/VertexAi/Tests/TokenUsageExtractorTest.php

@@ -68,21 +68,9 @@ public function testItHandlesMissingUsageFields()
     public function testItHandlesStreamResults()
     {
         $extractor = new TokenUsageExtractor();
-        $result = new InMemoryRawResult(dataStream: [
-            ['content' => 'chunk1'],
-            ['content' => 'chunk2', 'usageMetadata' => [
-                'promptTokenCount' => 15,
-                'candidatesTokenCount' => 25,
-                'totalTokenCount' => 40,
-            ]],
-        ]);
 
-        $tokenUsage = $extractor->extract($result, ['stream' => true]);
+        $tokenUsage = $extractor->extract(new InMemoryRawResult(), ['stream' => true]);
 
-        $this->assertInstanceOf(TokenUsage::class, $tokenUsage);
-        $this->assertSame(15, $tokenUsage->getPromptTokens());
-        $this->assertSame(25, $tokenUsage->getCompletionTokens());
-        $this->assertNull($tokenUsage->getThinkingTokens());
-        $this->assertSame(40, $tokenUsage->getTotalTokens());
+        $this->assertNull($tokenUsage);
     }
 }