mirror of
https://github.com/symfony/ai.git
synced 2026-03-23 23:42:18 +01:00
[Platform][Mistral] Add token usage extraction for embeddings
Add TokenUsageExtractor for Mistral embeddings to extract prompt_tokens, total_tokens, and rate limit headers from embedding responses. Wire it into the embeddings ResultConverter.
This commit is contained in:
committed by
Christopher Hertel
parent
f127af70c2
commit
81caade29f
@@ -23,3 +23,4 @@ $result = $platform->invoke('mistral-embed', <<<TEXT
|
||||
TEXT);
|
||||
|
||||
print_vectors($result);
|
||||
print_token_usage($result->getMetadata()->get('token_usage'));
|
||||
|
||||
@@ -1,6 +1,11 @@
|
||||
CHANGELOG
|
||||
=========
|
||||
|
||||
0.7
|
||||
---
|
||||
|
||||
* Add token usage extraction for embeddings
|
||||
|
||||
0.1
|
||||
---
|
||||
|
||||
|
||||
@@ -53,8 +53,8 @@ final class ResultConverter implements ResultConverterInterface
|
||||
);
|
||||
}
|
||||
|
||||
public function getTokenUsageExtractor(): ?TokenUsageExtractorInterface
|
||||
public function getTokenUsageExtractor(): TokenUsageExtractorInterface
|
||||
{
|
||||
return null;
|
||||
return new TokenUsageExtractor();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,46 @@
|
||||
<?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\Mistral\Embeddings;
|
||||
|
||||
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 Johannes Wachter <johannes@sulu.io>
|
||||
*/
|
||||
final class TokenUsageExtractor implements TokenUsageExtractorInterface
|
||||
{
|
||||
public function extract(RawResultInterface $rawResult, array $options = []): ?TokenUsageInterface
|
||||
{
|
||||
$rawResponse = $rawResult->getObject();
|
||||
if (!$rawResponse instanceof ResponseInterface) {
|
||||
return null;
|
||||
}
|
||||
|
||||
$headers = $rawResponse->getHeaders(false);
|
||||
|
||||
$remainingTokensMinute = $headers['x-ratelimit-limit-tokens-minute'][0] ?? null;
|
||||
$remainingTokensMonth = $headers['x-ratelimit-limit-tokens-month'][0] ?? null;
|
||||
|
||||
$content = $rawResult->getData();
|
||||
|
||||
return new TokenUsage(
|
||||
promptTokens: $content['usage']['prompt_tokens'] ?? null,
|
||||
remainingTokensMinute: null !== $remainingTokensMinute ? (int) $remainingTokensMinute : null,
|
||||
remainingTokensMonth: null !== $remainingTokensMonth ? (int) $remainingTokensMonth : null,
|
||||
totalTokens: $content['usage']['total_tokens'] ?? null,
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
<?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\Tests\Bridge\Mistral\Embeddings;
|
||||
|
||||
use PHPUnit\Framework\MockObject\MockObject;
|
||||
use PHPUnit\Framework\TestCase;
|
||||
use Symfony\AI\Platform\Bridge\Mistral\Embeddings\TokenUsageExtractor;
|
||||
use Symfony\AI\Platform\Result\InMemoryRawResult;
|
||||
use Symfony\AI\Platform\TokenUsage\TokenUsage;
|
||||
use Symfony\Contracts\HttpClient\ResponseInterface;
|
||||
|
||||
/**
|
||||
* @author Johannes Wachter <johannes@sulu.io>
|
||||
*/
|
||||
final class TokenUsageExtractorTest extends TestCase
|
||||
{
|
||||
public function testItDoesNothingWithoutUsageData()
|
||||
{
|
||||
$extractor = new TokenUsageExtractor();
|
||||
|
||||
$this->assertNull($extractor->extract(new InMemoryRawResult(['some' => 'data'])));
|
||||
}
|
||||
|
||||
public function testItExtractsTokenUsage()
|
||||
{
|
||||
$extractor = new TokenUsageExtractor();
|
||||
$result = new InMemoryRawResult([
|
||||
'usage' => [
|
||||
'prompt_tokens' => 10,
|
||||
'total_tokens' => 10,
|
||||
],
|
||||
], object: $this->createResponseObject());
|
||||
|
||||
$tokenUsage = $extractor->extract($result);
|
||||
|
||||
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
|
||||
$this->assertSame(1000, $tokenUsage->getRemainingTokensMinute());
|
||||
$this->assertSame(1000000, $tokenUsage->getRemainingTokensMonth());
|
||||
$this->assertSame(10, $tokenUsage->getPromptTokens());
|
||||
$this->assertNull($tokenUsage->getCompletionTokens());
|
||||
$this->assertSame(10, $tokenUsage->getTotalTokens());
|
||||
}
|
||||
|
||||
public function testItHandlesMissingUsageFields()
|
||||
{
|
||||
$extractor = new TokenUsageExtractor();
|
||||
$result = new InMemoryRawResult([
|
||||
'usage' => [
|
||||
'prompt_tokens' => 10,
|
||||
],
|
||||
], object: $this->createResponseObject());
|
||||
|
||||
$tokenUsage = $extractor->extract($result);
|
||||
|
||||
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
|
||||
$this->assertSame(1000, $tokenUsage->getRemainingTokensMinute());
|
||||
$this->assertSame(1000000, $tokenUsage->getRemainingTokensMonth());
|
||||
$this->assertSame(10, $tokenUsage->getPromptTokens());
|
||||
$this->assertNull($tokenUsage->getCompletionTokens());
|
||||
$this->assertNull($tokenUsage->getTotalTokens());
|
||||
}
|
||||
|
||||
private function createResponseObject(): ResponseInterface|MockObject
|
||||
{
|
||||
$response = $this->createStub(ResponseInterface::class);
|
||||
$response->method('getHeaders')->willReturn([
|
||||
'x-ratelimit-limit-tokens-minute' => ['1000'],
|
||||
'x-ratelimit-limit-tokens-month' => ['1000000'],
|
||||
]);
|
||||
|
||||
return $response;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user