Skip to content
Open
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
6 changes: 4 additions & 2 deletions src/Providers/Gemini/Handlers/Stream.php
Original file line number Diff line number Diff line change
Expand Up @@ -241,7 +241,8 @@ protected function processStream(Response $response, Request $request, int $dept
$this->state->markStepFinished();
yield new StepFinishEvent(
id: EventID::generate(),
timestamp: time()
timestamp: time(),
usage: $this->state->usage(),
);

yield $this->emitStreamEndEvent();
Expand Down Expand Up @@ -341,7 +342,8 @@ protected function handleToolCalls(
$this->state->markStepFinished();
yield new StepFinishEvent(
id: EventID::generate(),
timestamp: time()
timestamp: time(),
usage: $this->state->usage(),
);

$request->addMessage(new AssistantMessage($this->state->currentText(), $mappedToolCalls));
Expand Down
16 changes: 16 additions & 0 deletions src/Streaming/Events/StepFinishEvent.php
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,18 @@
namespace Prism\Prism\Streaming\Events;

use Prism\Prism\Enums\StreamEventType;
use Prism\Prism\ValueObjects\Usage;

readonly class StepFinishEvent extends StreamEvent
{
public function __construct(
string $id,
int $timestamp,
public ?Usage $usage = null, // Token usage information
) {
parent::__construct($id, $timestamp);
}

public function type(): StreamEventType
{
return StreamEventType::StepFinish;
Expand All @@ -18,6 +27,13 @@ public function toArray(): array
return [
'id' => $this->id,
'timestamp' => $this->timestamp,
'usage' => $this->usage instanceof Usage ? [
'prompt_tokens' => $this->usage->promptTokens,
'completion_tokens' => $this->usage->completionTokens,
'cache_write_input_tokens' => $this->usage->cacheWriteInputTokens,
'cache_read_input_tokens' => $this->usage->cacheReadInputTokens,
'thought_tokens' => $this->usage->thoughtTokens,
] : null,
];
}
}
3 changes: 2 additions & 1 deletion src/Testing/PrismFake.php
Original file line number Diff line number Diff line change
Expand Up @@ -336,7 +336,8 @@ protected function streamEventsFromTextResponse(TextResponse $response, TextRequ

yield new StepFinishEvent(
id: EventID::generate(),
timestamp: time()
timestamp: time(),
usage: $response->usage
);

yield new StreamEndEvent(
Expand Down
7 changes: 6 additions & 1 deletion tests/Providers/Gemini/GeminiStreamTest.php
Original file line number Diff line number Diff line change
Expand Up @@ -310,9 +310,14 @@
expect($stepStartEvents)->toHaveCount(1);

// Check for StepFinishEvent before StreamEndEvent
$stepFinishEvents = array_filter($events, fn (StreamEvent $e): bool => $e instanceof StepFinishEvent);
$stepFinishEvents = array_values(array_filter($events, fn (StreamEvent $e): bool => $e instanceof StepFinishEvent));
expect($stepFinishEvents)->toHaveCount(1);

// Verify StepFinishEvent contains usage data
expect($stepFinishEvents[0]->usage)->not->toBeNull();
expect($stepFinishEvents[0]->usage->promptTokens)->toBe(21);
expect($stepFinishEvents[0]->usage->completionTokens)->toBe(47);

// Verify order: StreamStart -> StepStart -> ... -> StepFinish -> StreamEnd
$eventTypes = array_map(get_class(...), $events);
$streamStartIndex = array_search(StreamStartEvent::class, $eventTypes);
Expand Down
82 changes: 82 additions & 0 deletions tests/Streaming/PrismStreamIntegrationTest.php
Original file line number Diff line number Diff line change
Expand Up @@ -1011,4 +1011,86 @@
expect($array)->toHaveKey('timestamp');
}
});

it('step finish event has default zero usage when fake response has no explicit usage', function (): void {
Prism::fake([
TextResponseFake::make()->withText('Test'),
]);

$events = iterator_to_array(
Prism::text()
->using('openai', 'gpt-4')
->withPrompt('Test')
->asStream()
);

$stepFinishEvents = array_values(array_filter(
$events,
fn (StreamEvent $e): bool => $e instanceof StepFinishEvent
));

expect($stepFinishEvents)->not->toBeEmpty();
expect($stepFinishEvents[0]->usage)->not->toBeNull();
expect($stepFinishEvents[0]->usage->promptTokens)->toBe(0);
expect($stepFinishEvents[0]->usage->completionTokens)->toBe(0);
});

it('step finish event contains usage when fake response has usage', function (): void {
Prism::fake([
TextResponseFake::make()
->withText('Test')
->withUsage(new Usage(100, 50)),
]);

$events = iterator_to_array(
Prism::text()
->using('openai', 'gpt-4')
->withPrompt('Test')
->asStream()
);

$stepFinishEvents = array_values(array_filter(
$events,
fn (StreamEvent $e): bool => $e instanceof StepFinishEvent
));

expect($stepFinishEvents)->not->toBeEmpty();
expect($stepFinishEvents[0]->usage)->not->toBeNull();
expect($stepFinishEvents[0]->usage->promptTokens)->toBe(100);
expect($stepFinishEvents[0]->usage->completionTokens)->toBe(50);
});

it('step finish event toArray includes usage when set', function (): void {
Prism::fake([
TextResponseFake::make()
->withText('Test')
->withUsage(new Usage(20, 10)),
]);

$events = iterator_to_array(
Prism::text()
->using('openai', 'gpt-4')
->withPrompt('Test')
->asStream()
);

$stepFinishEvents = array_values(array_filter(
$events,
fn (StreamEvent $e): bool => $e instanceof StepFinishEvent
));

expect($stepFinishEvents)->not->toBeEmpty();
$array = $stepFinishEvents[0]->toArray();
expect($array)->toHaveKey('usage');
expect($array['usage'])->not->toBeNull();
expect($array['usage']['prompt_tokens'])->toBe(20);
expect($array['usage']['completion_tokens'])->toBe(10);
});

it('step finish event toArray has null usage when event is constructed without usage', function (): void {
$event = new StepFinishEvent(id: 'test-id', timestamp: 1234567890);
$array = $event->toArray();
expect($array)->toHaveKey('usage');
expect($array['usage'])->toBeNull();
});
});