Skip to content

Commit ca9e8d4

Browse files
GuikingoneGuikingone
committed
refactor(platform): add TokenUsage for Ollama
1 parent 8e03d1e commit ca9e8d4

File tree

8 files changed

+211
-6
lines changed

8 files changed

+211
-6
lines changed

‎examples/ollama/chat-llama-with-token-usage.php‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
<?php
2+
3+
/*
4+
* This file is part of the Symfony package.
5+
*
6+
* (c) Fabien Potencier <[email protected]>
7+
*
8+
* For the full copyright and license information, please view the LICENSE
9+
* file that was distributed with this source code.
10+
*/
11+
12+
use Symfony\AI\Platform\Bridge\Ollama\PlatformFactory;
13+
use Symfony\AI\Platform\Message\Message;
14+
use Symfony\AI\Platform\Message\MessageBag;
15+
16+
require_once dirname(__DIR__).'/bootstrap.php';
17+
18+
$platform = PlatformFactory::create(env('OLLAMA_HOST_URL'), http_client());
19+
20+
try {
21+
$result = $platform->invoke(env('OLLAMA_LLM'), new MessageBag(
22+
Message::forSystem('You are a helpful assistant.'),
23+
Message::ofUser('Tina has one brother and one sister. How many sisters do Tina\'s siblings have?'),
24+
));
25+
26+
echo $result->asText().\PHP_EOL;
27+
28+
print_token_usage($result->getMetadata());
29+
} catch (InvalidArgumentException $e) {
30+
echo $e->getMessage()."\nMaybe use a different model?\n";
31+
}

‎examples/ollama/stream-with-token-usage.php‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
<?php
2+
3+
/*
4+
* This file is part of the Symfony package.
5+
*
6+
* (c) Fabien Potencier <[email protected]>
7+
*
8+
* For the full copyright and license information, please view the LICENSE
9+
* file that was distributed with this source code.
10+
*/
11+
12+
use Symfony\AI\Platform\Bridge\Ollama\PlatformFactory;
13+
use Symfony\AI\Platform\Message\Message;
14+
use Symfony\AI\Platform\Message\MessageBag;
15+
16+
require_once dirname(__DIR__).'/bootstrap.php';
17+
18+
$platform = PlatformFactory::create(env('OLLAMA_HOST_URL'), http_client());
19+
20+
$messages = new MessageBag(
21+
Message::forSystem('You are a helpful assistant.'),
22+
Message::ofUser('Tina has one brother and one sister. How many sisters do Tina\'s siblings have?'),
23+
);
24+
25+
$result = $platform->invoke(env('OLLAMA_LLM'), $messages, ['stream' => true]);
26+
27+
print_stream($result);
28+
print_token_usage($result->getMetadata());

‎src/platform/src/Bridge/Ollama/OllamaClient.php‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,8 @@ public function supports(Model $model): bool
3838
public function request(Model $model, array|string $payload, array $options = []): RawHttpResult
3939
{
4040
return match (true) {
41-
\in_array(Capability::INPUT_MESSAGES, $model->getCapabilities(), true) => $this->doCompletionRequest($payload, $options),
42-
\in_array(Capability::EMBEDDINGS, $model->getCapabilities(), true) => $this->doEmbeddingsRequest($model, $payload, $options),
41+
$model->supports(Capability::INPUT_MESSAGES) => $this->doCompletionRequest($payload, $options),
42+
$model->supports(Capability::EMBEDDINGS) => $this->doEmbeddingsRequest($model, $payload, $options),
4343
default => throw new InvalidArgumentException(\sprintf('Unsupported model "%s": "%s".', $model::class, $model->getName())),
4444
};
4545
}

‎src/platform/src/Bridge/Ollama/OllamaResultConverter.php‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,9 +47,9 @@ public function convert(RawResultInterface $result, array $options = []): Result
4747
: $this->doConvertCompletion($data);
4848
}
4949

50-
public function getTokenUsageExtractor(): ?TokenUsageExtractorInterface
50+
public function getTokenUsageExtractor(): TokenUsageExtractorInterface
5151
{
52-
return null;
52+
return new TokenUsageExtractor();
5353
}
5454

5555
/**

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
<?php
2+
3+
/*
4+
* This file is part of the Symfony package.
5+
*
6+
* (c) Fabien Potencier <[email protected]>
7+
*
8+
* For the full copyright and license information, please view the LICENSE
9+
* file that was distributed with this source code.
10+
*/
11+
12+
namespace Symfony\AI\Platform\Bridge\Ollama;
13+
14+
use Symfony\AI\Platform\Result\RawResultInterface;
15+
use Symfony\AI\Platform\TokenUsage\TokenUsage;
16+
use Symfony\AI\Platform\TokenUsage\TokenUsageExtractorInterface;
17+
use Symfony\AI\Platform\TokenUsage\TokenUsageInterface;
18+
use Symfony\Contracts\HttpClient\ResponseInterface;
19+
20+
/**
21+
* @author Guillaume Loulier <[email protected]>
22+
*/
23+
final class TokenUsageExtractor implements TokenUsageExtractorInterface
24+
{
25+
public function extract(RawResultInterface $rawResult, array $options = []): ?TokenUsageInterface
26+
{
27+
$response = $rawResult->getObject();
28+
if (!$response instanceof ResponseInterface) {
29+
return null;
30+
}
31+
32+
if ($options['stream'] ?? false) {
33+
foreach ($rawResult->getDataStream() as $chunk) {
34+
if ($chunk['done']) {
35+
return new TokenUsage(
36+
$chunk['prompt_eval_count'],
37+
$chunk['eval_count']
38+
);
39+
}
40+
}
41+
42+
return null;
43+
}
44+
45+
$payload = $response->toArray();
46+
47+
return new TokenUsage(
48+
$payload['prompt_eval_count'],
49+
$payload['eval_count']
50+
);
51+
}
52+
}

‎src/platform/tests/Bridge/Ollama/OllamaClientTest.php‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ final class OllamaClientTest extends TestCase
2828
{
2929
public function testSupportsModel()
3030
{
31-
$client = new OllamaClient(new MockHttpClient(), 'http://localhost:1234');
31+
$client = new OllamaClient(new MockHttpClient(), 'http://127.0.0.1:1234');
3232

3333
$this->assertTrue($client->supports(new Ollama('llama3.2')));
3434
$this->assertFalse($client->supports(new Model('any-model')));
@@ -97,6 +97,8 @@ public function testStreamingIsSupported()
9797
'created_at' => '2025-08-23T10:00:00Z',
9898
'message' => ['role' => 'assistant', 'content' => 'Hello world'],
9999
'done' => true,
100+
'prompt_eval_count' => 10,
101+
'eval_count' => 10,
100102
])."\n\n", [
101103
'response_headers' => [
102104
'content-type' => 'text/event-stream',

‎src/platform/tests/Bridge/Ollama/TokenUsageExtractorTest.php‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
<?php
2+
3+
/*
4+
* This file is part of the Symfony package.
5+
*
6+
* (c) Fabien Potencier <[email protected]>
7+
*
8+
* For the full copyright and license information, please view the LICENSE
9+
* file that was distributed with this source code.
10+
*/
11+
12+
namespace Symfony\AI\Platform\Tests\Bridge\Ollama;
13+
14+
use PHPUnit\Framework\MockObject\MockObject;
15+
use PHPUnit\Framework\TestCase;
16+
use Symfony\AI\Platform\Bridge\Ollama\TokenUsageExtractor;
17+
use Symfony\AI\Platform\Result\InMemoryRawResult;
18+
use Symfony\AI\Platform\TokenUsage\TokenUsage;
19+
use Symfony\Contracts\HttpClient\ResponseInterface;
20+
21+
final class TokenUsageExtractorTest extends TestCase
22+
{
23+
public function testItHandlesStreamResponsesWithoutProcessing()
24+
{
25+
$extractor = new TokenUsageExtractor();
26+
27+
$this->assertNull($extractor->extract(new InMemoryRawResult(), ['stream' => true]));
28+
}
29+
30+
public function testItDoesNothingWithoutUsageData()
31+
{
32+
$extractor = new TokenUsageExtractor();
33+
34+
$this->assertNull($extractor->extract(new InMemoryRawResult(['some' => 'data'])));
35+
}
36+
37+
public function testItExtractsTokenUsage()
38+
{
39+
$extractor = new TokenUsageExtractor();
40+
$result = new InMemoryRawResult([], object: $this->createResponseObject());
41+
42+
$tokenUsage = $extractor->extract($result);
43+
44+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
45+
$this->assertSame(10, $tokenUsage->getPromptTokens());
46+
$this->assertSame(10, $tokenUsage->getCompletionTokens());
47+
}
48+
49+
public function testItExtractsTokenUsageFromStreamResult()
50+
{
51+
$extractor = new TokenUsageExtractor();
52+
53+
$result = new InMemoryRawResult([], [
54+
[
55+
'model' => 'foo',
56+
'response' => 'First chunk',
57+
'done' => false,
58+
],
59+
[
60+
'model' => 'foo',
61+
'response' => 'Hello World!',
62+
'done' => true,
63+
'prompt_eval_count' => 10,
64+
'eval_count' => 10,
65+
],
66+
], object: $this->createResponseObject());
67+
68+
$tokenUsage = $extractor->extract($result, ['stream' => true]);
69+
70+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
71+
$this->assertSame(10, $tokenUsage->getPromptTokens());
72+
$this->assertSame(10, $tokenUsage->getCompletionTokens());
73+
}
74+
75+
private function createResponseObject(): ResponseInterface|MockObject
76+
{
77+
$response = $this->createStub(ResponseInterface::class);
78+
$response->method('toArray')->willReturn([
79+
'model' => 'foo',
80+
'response' => 'Hello World!',
81+
'done' => true,
82+
'prompt_eval_count' => 10,
83+
'eval_count' => 10,
84+
]);
85+
86+
return $response;
87+
}
88+
}

‎src/platform/tests/CachedPlatformTest.php‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ final class CachedPlatformTest extends TestCase
2626
{
2727
public function testPlatformCanReturnCachedResultWhenCalledTwice()
2828
{
29-
$httpResponse = $this->createStub(SymfonyHttpResponse::class);
29+
$httpResponse = $this->createMock(SymfonyHttpResponse::class);
3030
$rawHttpResult = new RawHttpResult($httpResponse);
3131

3232
$resultConverter = self::createMock(ResultConverterInterface::class);
@@ -47,12 +47,16 @@ public function testPlatformCanReturnCachedResultWhenCalledTwice()
4747
'prompt_cache_key' => 'symfony',
4848
]);
4949

50+
$this->assertTrue($deferredResult->getMetadata()->has('cached_at'));
51+
5052
$this->assertSame('test content', $deferredResult->getResult()->getContent());
5153

5254
$secondDeferredResult = $cachedPlatform->invoke('foo', 'bar', [
5355
'prompt_cache_key' => 'symfony',
5456
]);
5557

5658
$this->assertSame('test content', $secondDeferredResult->getResult()->getContent());
59+
$this->assertTrue($secondDeferredResult->getMetadata()->has('cached_at'));
60+
$this->assertSame($deferredResult->getMetadata()->get('cached_at'), $secondDeferredResult->getMetadata()->get('cached_at'));
5761
}
5862
}

0 commit comments

Comments
 (0)