Skip to content

[Platform] Standardise token usage #311

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions examples/openai/token-metadata.php
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
use Symfony\AI\Platform\Bridge\OpenAi\TokenOutputProcessor;
use Symfony\AI\Platform\Message\Message;
use Symfony\AI\Platform\Message\MessageBag;
use Symfony\AI\Platform\Result\Metadata\TokenUsage;

require_once dirname(__DIR__).'/bootstrap.php';

Expand All @@ -33,8 +34,11 @@
]);

$metadata = $result->getMetadata();
$tokenUsage = $metadata->get('token_usage');

echo 'Utilized Tokens: '.$metadata['total_tokens'].\PHP_EOL;
echo '-- Prompt Tokens: '.$metadata['prompt_tokens'].\PHP_EOL;
echo '-- Completion Tokens: '.$metadata['completion_tokens'].\PHP_EOL;
echo 'Remaining Tokens: '.$metadata['remaining_tokens'].\PHP_EOL;
assert($tokenUsage instanceof TokenUsage);

echo 'Utilized Tokens: '.$tokenUsage->totalTokens.\PHP_EOL;
echo '-- Prompt Tokens: '.$tokenUsage->promptTokens.\PHP_EOL;
echo '-- Completion Tokens: '.$tokenUsage->completionTokens.\PHP_EOL;
echo 'Remaining Tokens: '.$tokenUsage->remainingTokens.\PHP_EOL;
27 changes: 16 additions & 11 deletions src/platform/src/Bridge/Mistral/TokenOutputProcessor.php
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

use Symfony\AI\Agent\Output;
use Symfony\AI\Agent\OutputProcessorInterface;
use Symfony\AI\Platform\Result\Metadata\TokenUsage;
use Symfony\AI\Platform\Result\StreamResult;
use Symfony\Contracts\HttpClient\ResponseInterface;

Expand All @@ -34,25 +35,29 @@ public function processOutput(Output $output): void
}

$metadata = $output->result->getMetadata();
$headers = $rawResponse->getHeaders(false);

$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],
$remainingTokensMinute = $headers['x-ratelimit-limit-tokens-minute'][0] ?? null;
$remainingTokensMonth = $headers['x-ratelimit-limit-tokens-month'][0] ?? null;
$tokenUsage = new TokenUsage(
remainingTokensMinute: null !== $remainingTokensMinute ? (int) $remainingTokensMinute : null,
remainingTokensMonth: null !== $remainingTokensMonth ? (int) $remainingTokensMonth : null,
);

$content = $rawResponse->toArray(false);

if (!\array_key_exists('usage', $content)) {
$metadata->add('token_usage', $tokenUsage);

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);
$usage = $content['usage'];

$tokenUsage->promptTokens = $usage['prompt_tokens'] ?? null;
$tokenUsage->completionTokens = $usage['completion_tokens'] ?? null;
$tokenUsage->totalTokens = $usage['total_tokens'] ?? null;

$metadata->add('token_usage', $tokenUsage);
}
}
21 changes: 15 additions & 6 deletions src/platform/src/Bridge/OpenAi/TokenOutputProcessor.php
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

use Symfony\AI\Agent\Output;
use Symfony\AI\Agent\OutputProcessorInterface;
use Symfony\AI\Platform\Result\Metadata\TokenUsage;
use Symfony\AI\Platform\Result\StreamResult;
use Symfony\Contracts\HttpClient\ResponseInterface;

Expand All @@ -35,19 +36,27 @@ public function processOutput(Output $output): void

$metadata = $output->result->getMetadata();

$metadata->add(
'remaining_tokens',
(int) $rawResponse->getHeaders(false)['x-ratelimit-remaining-tokens'][0],
$remainingTokens = $rawResponse->getHeaders(false)['x-ratelimit-remaining-tokens'][0] ?? null;
$tokenUsage = new TokenUsage(
remainingTokens: null !== $remainingTokens ? (int) $remainingTokens : null,
);

$content = $rawResponse->toArray(false);

if (!\array_key_exists('usage', $content)) {
$metadata->add('token_usage', $tokenUsage);

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);
$usage = $content['usage'];

$tokenUsage->promptTokens = $usage['prompt_tokens'] ?? null;
$tokenUsage->completionTokens = $usage['completion_tokens'] ?? null;
$tokenUsage->thinkingTokens = $usage['completion_tokens_details']['reasoning_tokens'] ?? null;
$tokenUsage->cachedTokens = $usage['prompt_tokens_details']['cached_tokens'] ?? null;
$tokenUsage->totalTokens = $usage['total_tokens'] ?? null;

$metadata->add('token_usage', $tokenUsage);
}
}
56 changes: 56 additions & 0 deletions src/platform/src/Result/Metadata/TokenUsage.php
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
<?php

/*
* This file is part of the Symfony package.
*
* (c) Fabien Potencier <[email protected]>
*
* For the full copyright and license information, please view the LICENSE
* file that was distributed with this source code.
*/

namespace Symfony\AI\Platform\Result\Metadata;

/**
* @author Junaid Farooq <[email protected]>
*/
final class TokenUsage implements \JsonSerializable
{
public function __construct(
public ?int $promptTokens = null,
public ?int $completionTokens = null,
public ?int $thinkingTokens = null,
public ?int $cachedTokens = null,
public ?int $remainingTokens = null,
public ?int $remainingTokensMinute = null,
public ?int $remainingTokensMonth = null,
public ?int $totalTokens = null,
) {
}

/**
* @return array{
* prompt_tokens: ?int,
* completion_tokens: ?int,
* thinking_tokens: ?int,
* cached_tokens: ?int,
* remaining_tokens: ?int,
* remaining_tokens_minute: ?int,
* remaining_tokens_month: ?int,
* total_tokens: ?int,
* }
*/
public function jsonSerialize(): array
{
return [
'prompt_tokens' => $this->promptTokens,
'completion_tokens' => $this->completionTokens,
'thinking_tokens' => $this->thinkingTokens,
'cached_tokens' => $this->cachedTokens,
'remaining_tokens' => $this->remainingTokens,
'remaining_tokens_minute' => $this->remainingTokensMinute,
'remaining_tokens_month' => $this->remainingTokensMonth,
'total_tokens' => $this->totalTokens,
];
}
}
38 changes: 23 additions & 15 deletions src/platform/tests/Bridge/Mistral/TokenOutputProcessorTest.php
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
use Symfony\AI\Platform\Message\MessageBagInterface;
use Symfony\AI\Platform\Model;
use Symfony\AI\Platform\Result\Metadata\Metadata;
use Symfony\AI\Platform\Result\Metadata\TokenUsage;
use Symfony\AI\Platform\Result\RawHttpResult;
use Symfony\AI\Platform\Result\ResultInterface;
use Symfony\AI\Platform\Result\StreamResult;
Expand Down Expand Up @@ -70,9 +71,12 @@ public function testItAddsRemainingTokensToMetadata()
$processor->processOutput($output);

$metadata = $output->result->getMetadata();
$this->assertCount(2, $metadata);
$this->assertSame(1000, $metadata->get('remaining_tokens_minute'));
$this->assertSame(1000000, $metadata->get('remaining_tokens_month'));
$tokenUsage = $metadata->get('token_usage');

$this->assertCount(1, $metadata);
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
$this->assertSame(1000, $tokenUsage->remainingTokensMinute);
$this->assertSame(1000000, $tokenUsage->remainingTokensMonth);
}

public function testItAddsUsageTokensToMetadata()
Expand All @@ -95,12 +99,14 @@ public function testItAddsUsageTokensToMetadata()
$processor->processOutput($output);

$metadata = $output->result->getMetadata();
$this->assertCount(5, $metadata);
$this->assertSame(1000, $metadata->get('remaining_tokens_minute'));
$this->assertSame(1000000, $metadata->get('remaining_tokens_month'));
$this->assertSame(10, $metadata->get('prompt_tokens'));
$this->assertSame(20, $metadata->get('completion_tokens'));
$this->assertSame(30, $metadata->get('total_tokens'));
$tokenUsage = $metadata->get('token_usage');

$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
$this->assertSame(1000, $tokenUsage->remainingTokensMinute);
$this->assertSame(1000000, $tokenUsage->remainingTokensMonth);
$this->assertSame(10, $tokenUsage->promptTokens);
$this->assertSame(20, $tokenUsage->completionTokens);
$this->assertSame(30, $tokenUsage->totalTokens);
}

public function testItHandlesMissingUsageFields()
Expand All @@ -122,12 +128,14 @@ public function testItHandlesMissingUsageFields()
$processor->processOutput($output);

$metadata = $output->result->getMetadata();
$this->assertCount(5, $metadata);
$this->assertSame(1000, $metadata->get('remaining_tokens_minute'));
$this->assertSame(1000000, $metadata->get('remaining_tokens_month'));
$this->assertSame(10, $metadata->get('prompt_tokens'));
$this->assertNull($metadata->get('completion_tokens'));
$this->assertNull($metadata->get('total_tokens'));
$tokenUsage = $metadata->get('token_usage');

$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
$this->assertSame(1000, $tokenUsage->remainingTokensMinute);
$this->assertSame(1000000, $tokenUsage->remainingTokensMonth);
$this->assertSame(10, $tokenUsage->promptTokens);
$this->assertNull($tokenUsage->completionTokens);
$this->assertNull($tokenUsage->totalTokens);
}

private function createRawResponse(array $data = []): RawHttpResult
Expand Down
41 changes: 29 additions & 12 deletions src/platform/tests/Bridge/OpenAi/TokenOutputProcessorTest.php
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
use Symfony\AI\Platform\Message\MessageBagInterface;
use Symfony\AI\Platform\Model;
use Symfony\AI\Platform\Result\Metadata\Metadata;
use Symfony\AI\Platform\Result\Metadata\TokenUsage;
use Symfony\AI\Platform\Result\RawHttpResult;
use Symfony\AI\Platform\Result\ResultInterface;
use Symfony\AI\Platform\Result\StreamResult;
Expand All @@ -31,6 +32,7 @@
#[UsesClass(TextResult::class)]
#[UsesClass(StreamResult::class)]
#[UsesClass(Metadata::class)]
#[UsesClass(TokenUsage::class)]
#[Small]
final class TokenOutputProcessorTest extends TestCase
{
Expand Down Expand Up @@ -70,8 +72,11 @@ public function testItAddsRemainingTokensToMetadata()
$processor->processOutput($output);

$metadata = $output->result->getMetadata();
$tokenUsage = $metadata->get('token_usage');
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What about having a dedicated setTokenUsage/getTokenUsage on the metadata object?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure about this.
Metadata object already has a lot of methods, and each of them is generic in nature as far as I can see, and doesn't concern any specific model characteristic.


$this->assertCount(1, $metadata);
$this->assertSame(1000, $metadata->get('remaining_tokens'));
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
$this->assertSame(1000, $tokenUsage->remainingTokens);
}

public function testItAddsUsageTokensToMetadata()
Expand All @@ -83,7 +88,13 @@ public function testItAddsUsageTokensToMetadata()
'usage' => [
'prompt_tokens' => 10,
'completion_tokens' => 20,
'total_tokens' => 30,
'total_tokens' => 50,
'completion_tokens_details' => [
'reasoning_tokens' => 20,
],
'prompt_tokens_details' => [
'cached_tokens' => 40,
],
],
]);

Expand All @@ -94,11 +105,15 @@ public function testItAddsUsageTokensToMetadata()
$processor->processOutput($output);

$metadata = $output->result->getMetadata();
$this->assertCount(4, $metadata);
$this->assertSame(1000, $metadata->get('remaining_tokens'));
$this->assertSame(10, $metadata->get('prompt_tokens'));
$this->assertSame(20, $metadata->get('completion_tokens'));
$this->assertSame(30, $metadata->get('total_tokens'));
$tokenUsage = $metadata->get('token_usage');

$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
$this->assertSame(10, $tokenUsage->promptTokens);
$this->assertSame(20, $tokenUsage->completionTokens);
$this->assertSame(1000, $tokenUsage->remainingTokens);
$this->assertSame(20, $tokenUsage->thinkingTokens);
$this->assertSame(40, $tokenUsage->cachedTokens);
$this->assertSame(50, $tokenUsage->totalTokens);
}

public function testItHandlesMissingUsageFields()
Expand All @@ -120,11 +135,13 @@ public function testItHandlesMissingUsageFields()
$processor->processOutput($output);

$metadata = $output->result->getMetadata();
$this->assertCount(4, $metadata);
$this->assertSame(1000, $metadata->get('remaining_tokens'));
$this->assertSame(10, $metadata->get('prompt_tokens'));
$this->assertNull($metadata->get('completion_tokens'));
$this->assertNull($metadata->get('total_tokens'));
$tokenUsage = $metadata->get('token_usage');

$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
$this->assertSame(10, $tokenUsage->promptTokens);
$this->assertSame(1000, $tokenUsage->remainingTokens);
$this->assertNull($tokenUsage->completionTokens);
$this->assertNull($tokenUsage->totalTokens);
}

private function createRawResult(array $data = []): RawHttpResult
Expand Down