Skip to content

Commit 3901999

Browse files
nictou-weglotnictou-weglot
committed
[Platform][Anthropic] Add TokenOutputProcessor
1 parent 4134b97 commit 3901999

File tree

2 files changed

+222
-0
lines changed

2 files changed

+222
-0
lines changed

‎src/Bridge/Anthropic/TokenOutputProcessor.php‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
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\Anthropic;
13+
14+
use Symfony\AI\Agent\Output;
15+
use Symfony\AI\Agent\OutputProcessorInterface;
16+
use Symfony\AI\Platform\Metadata\TokenUsage;
17+
use Symfony\AI\Platform\Result\StreamResult;
18+
use Symfony\Contracts\HttpClient\ResponseInterface;
19+
20+
final class TokenOutputProcessor implements OutputProcessorInterface
21+
{
22+
public function processOutput(Output $output): void
23+
{
24+
if ($output->result instanceof StreamResult) {
25+
// Streams have to be handled manually as the tokens are part of the streamed chunks
26+
return;
27+
}
28+
29+
$rawResponse = $output->result->getRawResult()?->getObject();
30+
if (!$rawResponse instanceof ResponseInterface) {
31+
return;
32+
}
33+
34+
$metadata = $output->result->getMetadata();
35+
36+
$tokenUsage = new TokenUsage();
37+
38+
$content = $rawResponse->toArray(false);
39+
if (!\array_key_exists('usage', $content)) {
40+
$metadata->add('token_usage', $tokenUsage);
41+
42+
return;
43+
}
44+
45+
$usage = $content['usage'];
46+
47+
$tokenUsage->promptTokens = $usage['input_tokens'] ?? null;
48+
$tokenUsage->completionTokens = $usage['output_tokens'] ?? null;
49+
$tokenUsage->toolTokens = $usage['server_tool_use']['web_search_requests'] ?? null;
50+
51+
$cachedTokens = null;
52+
if (\array_key_exists('cache_creation_input_tokens', $usage) || \array_key_exists('cache_read_input_tokens', $usage)) {
53+
$cachedTokens = ($usage['cache_creation_input_tokens'] ?? 0) + ($usage['cache_read_input_tokens'] ?? 0);
54+
}
55+
$tokenUsage->cachedTokens = $cachedTokens;
56+
57+
$metadata->add('token_usage', $tokenUsage);
58+
}
59+
}

‎tests/Bridge/Anthropic/TokenOutputProcessorTest.php‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
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 Bridge\Anthropic;
13+
14+
use PHPUnit\Framework\Attributes\CoversClass;
15+
use PHPUnit\Framework\Attributes\Small;
16+
use PHPUnit\Framework\Attributes\UsesClass;
17+
use PHPUnit\Framework\TestCase;
18+
use Symfony\AI\Agent\Output;
19+
use Symfony\AI\Platform\Bridge\Anthropic\TokenOutputProcessor;
20+
use Symfony\AI\Platform\Message\MessageBag;
21+
use Symfony\AI\Platform\Metadata\Metadata;
22+
use Symfony\AI\Platform\Metadata\TokenUsage;
23+
use Symfony\AI\Platform\Model;
24+
use Symfony\AI\Platform\Result\RawHttpResult;
25+
use Symfony\AI\Platform\Result\ResultInterface;
26+
use Symfony\AI\Platform\Result\StreamResult;
27+
use Symfony\AI\Platform\Result\TextResult;
28+
use Symfony\Contracts\HttpClient\ResponseInterface;
29+
30+
#[CoversClass(TokenOutputProcessor::class)]
31+
#[UsesClass(Output::class)]
32+
#[UsesClass(TextResult::class)]
33+
#[UsesClass(StreamResult::class)]
34+
#[UsesClass(Metadata::class)]
35+
#[UsesClass(TokenUsage::class)]
36+
#[Small]
37+
final class TokenOutputProcessorTest extends TestCase
38+
{
39+
public function testItHandlesStreamResponsesWithoutProcessing()
40+
{
41+
$processor = new TokenOutputProcessor();
42+
$streamResult = new StreamResult((static function () { yield 'test'; })());
43+
$output = $this->createOutput($streamResult);
44+
45+
$processor->processOutput($output);
46+
47+
$metadata = $output->result->getMetadata();
48+
$this->assertCount(0, $metadata);
49+
}
50+
51+
public function testItDoesNothingWithoutRawResponse()
52+
{
53+
$processor = new TokenOutputProcessor();
54+
$textResult = new TextResult('test');
55+
$output = $this->createOutput($textResult);
56+
57+
$processor->processOutput($output);
58+
59+
$metadata = $output->result->getMetadata();
60+
$this->assertCount(0, $metadata);
61+
}
62+
63+
public function testItAddsRemainingTokensToMetadata()
64+
{
65+
$processor = new TokenOutputProcessor();
66+
$textResult = new TextResult('test');
67+
68+
$textResult->setRawResult($this->createRawResult());
69+
70+
$output = $this->createOutput($textResult);
71+
72+
$processor->processOutput($output);
73+
74+
$metadata = $output->result->getMetadata();
75+
$tokenUsage = $metadata->get('token_usage');
76+
77+
$this->assertCount(1, $metadata);
78+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
79+
$this->assertNull($tokenUsage->remainingTokens);
80+
}
81+
82+
public function testItAddsUsageTokensToMetadata()
83+
{
84+
$processor = new TokenOutputProcessor();
85+
$textResult = new TextResult('test');
86+
87+
$rawResult = $this->createRawResult([
88+
'usage' => [
89+
'input_tokens' => 10,
90+
'output_tokens' => 20,
91+
'server_tool_use' => [
92+
'web_search_requests' => 30,
93+
],
94+
'cache_creation_input_tokens' => 40,
95+
'cache_read_input_tokens' => 50,
96+
],
97+
]);
98+
99+
$textResult->setRawResult($rawResult);
100+
101+
$output = $this->createOutput($textResult);
102+
103+
$processor->processOutput($output);
104+
105+
$metadata = $output->result->getMetadata();
106+
$tokenUsage = $metadata->get('token_usage');
107+
108+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
109+
$this->assertSame(10, $tokenUsage->promptTokens);
110+
$this->assertSame(30, $tokenUsage->toolTokens);
111+
$this->assertSame(20, $tokenUsage->completionTokens);
112+
$this->assertNull($tokenUsage->remainingTokens);
113+
$this->assertNull($tokenUsage->thinkingTokens);
114+
$this->assertSame(90, $tokenUsage->cachedTokens);
115+
$this->assertNull($tokenUsage->totalTokens);
116+
}
117+
118+
public function testItHandlesMissingUsageFields()
119+
{
120+
$processor = new TokenOutputProcessor();
121+
$textResult = new TextResult('test');
122+
123+
$rawResult = $this->createRawResult([
124+
'usage' => [
125+
// Missing some fields
126+
'input_tokens' => 10,
127+
],
128+
]);
129+
130+
$textResult->setRawResult($rawResult);
131+
132+
$output = $this->createOutput($textResult);
133+
134+
$processor->processOutput($output);
135+
136+
$metadata = $output->result->getMetadata();
137+
$tokenUsage = $metadata->get('token_usage');
138+
139+
$this->assertInstanceOf(TokenUsage::class, $tokenUsage);
140+
$this->assertSame(10, $tokenUsage->promptTokens);
141+
$this->assertNull($tokenUsage->remainingTokens);
142+
$this->assertNull($tokenUsage->completionTokens);
143+
$this->assertNull($tokenUsage->totalTokens);
144+
}
145+
146+
private function createRawResult(array $data = []): RawHttpResult
147+
{
148+
$rawResponse = $this->createStub(ResponseInterface::class);
149+
$rawResponse->method('toArray')->willReturn($data);
150+
151+
return new RawHttpResult($rawResponse);
152+
}
153+
154+
private function createOutput(ResultInterface $result): Output
155+
{
156+
return new Output(
157+
$this->createStub(Model::class),
158+
$result,
159+
new MessageBag(),
160+
[],
161+
);
162+
}
163+
}

0 commit comments

Comments
 (0)