Chat stream usage: update annotations and add test

Author: gehrisandro

src/Responses/Chat/CreateStreamedResponse.php

@@ -9,12 +9,12 @@
 use OpenAI\Testing\Responses\Concerns\FakeableForStreamedResponse;
 
 /**
- * @implements ResponseContract<array{id: string, object: string, created: int, model: string, choices: array<int, array{index: int, delta: array{role?: string, content?: string}|array{role?: string, content: null, function_call: array{name?: string, arguments?: string}}, finish_reason: string|null}>}>
+ * @implements ResponseContract<array{id: string, object: string, created: int, model: string, choices: array<int, array{index: int, delta: array{role?: string, content?: string}|array{role?: string, content: null, function_call: array{name?: string, arguments?: string}}, finish_reason: string|null}>, usage?: array{prompt_tokens: int, completion_tokens: int|null, total_tokens: int}}>
  */
 final class CreateStreamedResponse implements ResponseContract
 {
     /**
-     * @use ArrayAccessible<array{id: string, object: string, created: int, model: string, choices: array<int, array{index: int, delta: array{role?: string, content?: string}|array{role?: string, content: null, function_call: array{name?: string, arguments?: string}}, finish_reason: string|null}>}>
+     * @use ArrayAccessible<array{id: string, object: string, created: int, model: string, choices: array<int, array{index: int, delta: array{role?: string, content?: string}|array{role?: string, content: null, function_call: array{name?: string, arguments?: string}}, finish_reason: string|null}>, usage?: array{prompt_tokens: int, completion_tokens: int|null, total_tokens: int}}>
      */
     use ArrayAccessible;
 
@@ -36,7 +36,7 @@ private function __construct(
     /**
      * Acts as static factory, and returns a new Response instance.
      *
-     * @param  array{id: string, object: string, created: int, model: string, choices: array<int, array{index: int, delta: array{role?: string, content?: string}, finish_reason: string|null}>}  $attributes
+     * @param  array{id: string, object: string, created: int, model: string, choices: array<int, array{index: int, delta: array{role?: string, content?: string}, finish_reason: string|null}>, usage?: array{prompt_tokens: int, completion_tokens: int|null, total_tokens: int}}  $attributes
      */
     public static function from(array $attributes): self
     {
@@ -59,7 +59,7 @@ public static function from(array $attributes): self
      */
     public function toArray(): array
     {
-        return [
+        $data = [
             'id' => $this->id,
             'object' => $this->object,
             'created' => $this->created,
@@ -68,7 +68,12 @@ public function toArray(): array
                 static fn (CreateStreamedResponseChoice $result): array => $result->toArray(),
                 $this->choices,
             ),
-            'usage' => $this->usage?->toArray() ?? null,
         ];
+
+        if ($this->usage instanceof \OpenAI\Responses\Chat\CreateResponseUsage) {
+            $data['usage'] = $this->usage->toArray();
+        }
+
+        return $data;
     }
 }

tests/Fixtures/Chat.php

@@ -206,6 +206,22 @@ function chatCompletionStreamContentChunk(): array
     ];
 }
 
+function chatCompletionStreamUsageChunk(): array
+{
+    return [
+        'id' => 'chatcmpl-6wdIE4DsUtqf1srdMTsfkJp0VWZgz',
+        'object' => 'chat.completion.chunk',
+        'created' => 1679432086,
+        'model' => 'gpt-4-0314',
+        'choices' => [],
+        'usage' => [
+            'prompt_tokens' => 9,
+            'completion_tokens' => 12,
+            'total_tokens' => 21,
+        ],
+    ];
+}
+
 function chatCompletionStreamFunctionCallChunk(): array
 {
     return [

tests/Responses/Chat/CreateStreamedResponse.php

@@ -1,5 +1,6 @@
 <?php
 
+use OpenAI\Responses\Chat\CreateResponseUsage;
 use OpenAI\Responses\Chat\CreateStreamedResponse;
 use OpenAI\Responses\Chat\CreateStreamedResponseChoice;
 
@@ -16,6 +17,22 @@
         ->choices->each->toBeInstanceOf(CreateStreamedResponseChoice::class);
 });
 
+test('from usage chunk', function () {
+    $completion = CreateStreamedResponse::from(chatCompletionStreamUsageChunk());
+
+    expect($completion)
+        ->toBeInstanceOf(CreateStreamedResponse::class)
+        ->id->toBe('chatcmpl-6wdIE4DsUtqf1srdMTsfkJp0VWZgz')
+        ->object->toBe('chat.completion.chunk')
+        ->created->toBe(1679432086)
+        ->model->toBe('gpt-4-0314')
+        ->choices->toBeArray()->toHaveCount(0)
+        ->usage->toBeInstanceOf(CreateResponseUsage::class)
+        ->usage->promptTokens->toBe(9)
+        ->usage->completionTokens->toBe(12)
+        ->usage->totalTokens->toBe(21);
+});
+
 test('as array accessible', function () {
     $completion = CreateStreamedResponse::from(chatCompletionStreamFirstChunk());