Merge prism-php#1000 — fix(gemini): add Usage to StepFinishEvent

Author: wishborn

src/Providers/Gemini/Handlers/Stream.php

@@ -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();
@@ -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));

src/Streaming/Events/StepFinishEvent.php

@@ -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;
@@ -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,
         ];
     }
 }

src/Testing/PrismFake.php

@@ -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(

tests/Providers/Gemini/GeminiStreamTest.php

@@ -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);

tests/Streaming/PrismStreamIntegrationTest.php

@@ -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();
+    });
 });