symfony/archived-ai-ollama-platform
1
0
Fork 0
mirror of https://github.com/symfony/ai-ollama-platform.git synced 2026-03-23 23:22:07 +01:00
Code Issues Packages Projects Releases Wiki Activity

Disable Ollama token extraction for streams

Browse Source
This commit is contained in:
Christopher Hertel
2025-12-23 11:33:11 +01:00
parent cf536a33f3
commit 0c4cfde3de
2 changed files with 3 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
Tests/TokenUsageExtractorTest.php
View File

@@ -54,25 +54,8 @@ final class TokenUsageExtractorTest extends TestCase
{
$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(new InMemoryRawResult(), ['stream' => true]);
$tokenUsage = $extractor->extract($result, ['stream' => true]);
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
$this->assertSame(10, $tokenUsage->getPromptTokens());
$this->assertSame(10, $tokenUsage->getCompletionTokens());
$this->assertNull($tokenUsage);
}
}

10
TokenUsageExtractor.php
View File

@@ -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;
}