Skip to content

Commit 15c2f84

Browse files
ilvalerioneilvalerione
authored
Merge pull request #542 from CodeWrap/fix/openai-image-streaming
fix: OpenAIImage stream TypeError + error / truncation handling
2 parents ad7a4c1 + 2a875a2 commit 15c2f84

File tree

2 files changed

+289
-9
lines changed

2 files changed

+289
-9
lines changed

‎src/Providers/OpenAI/Image/OpenAIImage.php‎

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,10 @@
2222
use NeuronAI\Providers\MessageMapperInterface;
2323
use NeuronAI\Providers\SSEParser;
2424
use NeuronAI\Providers\ToolMapperInterface;
25+
use NeuronAI\UniqueIdGenerator;
2526

2627
use function end;
28+
use function is_string;
2729

2830
class OpenAIImage implements AIProviderInterface
2931
{
@@ -142,25 +144,47 @@ public function stream(Message ...$messages): Generator
142144
)
143145
);
144146

145-
$content = '';
147+
$messageId = UniqueIdGenerator::generateId('msg_');
148+
$content = null;
146149
$usage = new Usage(0, 0);
147150

148151
while (! $stream->eof()) {
149152
if (!$line = SSEParser::parseNextSSEEvent($stream)) {
150153
continue;
151154
}
152155

153-
// Image APIs stream entire partially generated images, not base64 chunks.
154-
// The last content streamed is the final image.
155-
if ($line['type'] === 'image_generation.partial_image') {
156-
$content = $line['b64_json'];
157-
yield new ImageChunk($line['partial_image_index'], $line['b64_json']);
156+
$type = $line['type'] ?? null;
157+
158+
if ($type === 'error') {
159+
throw new ProviderException(
160+
$line['error']['message'] ?? 'Image generation failed.'
161+
);
158162
}
159163

160-
if (isset($line['usage'])) {
161-
$usage->inputTokens = $line['usage']['input_tokens'] ?? 0;
162-
$usage->outputTokens = $line['usage']['output_tokens'] ?? 0;
164+
if ($type === 'image_generation.partial_image') {
165+
$b64 = $line['b64_json'] ?? null;
166+
if (!is_string($b64) || $b64 === '') {
167+
throw new ProviderException('Received a partial image event without b64_json payload.');
168+
}
169+
yield new ImageChunk($messageId, $b64);
163170
}
171+
172+
if ($type === 'image_generation.completed') {
173+
$b64 = $line['b64_json'] ?? null;
174+
if (!is_string($b64) || $b64 === '') {
175+
throw new ProviderException('Received a completed image event without b64_json payload.');
176+
}
177+
$content = $b64;
178+
179+
if (isset($line['usage'])) {
180+
$usage->inputTokens = $line['usage']['input_tokens'] ?? 0;
181+
$usage->outputTokens = $line['usage']['output_tokens'] ?? 0;
182+
}
183+
}
184+
}
185+
186+
if ($content === null) {
187+
throw new ProviderException('Image generation stream ended before a completed image was received.');
164188
}
165189

166190
$result = new AssistantMessage(

‎tests/Providers/OpenAIImageTest.php‎

Lines changed: 256 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,256 @@
1+
<?php
2+
3+
declare(strict_types=1);
4+
5+
namespace NeuronAI\Tests\Providers;
6+
7+
use GuzzleHttp\Handler\MockHandler;
8+
use GuzzleHttp\HandlerStack;
9+
use GuzzleHttp\Middleware;
10+
use GuzzleHttp\Psr7\Response;
11+
use NeuronAI\Chat\Enums\SourceType;
12+
use NeuronAI\Chat\Messages\AssistantMessage;
13+
use NeuronAI\Chat\Messages\ContentBlocks\ImageContent;
14+
use NeuronAI\Chat\Messages\Stream\Chunks\ImageChunk;
15+
use NeuronAI\Chat\Messages\UserMessage;
16+
use NeuronAI\Exceptions\ProviderException;
17+
use NeuronAI\HttpClient\GuzzleHttpClient;
18+
use NeuronAI\Providers\OpenAI\Image\OpenAIImage;
19+
use PHPUnit\Framework\TestCase;
20+
21+
use function json_decode;
22+
23+
class OpenAIImageTest extends TestCase
24+
{
25+
public function test_chat_returns_image_message(): void
26+
{
27+
$body = '{"data":[{"b64_json":"FINAL_BASE64"}],"usage":{"input_tokens":10,"output_tokens":20,"total_tokens":30}}';
28+
29+
$mockHandler = new MockHandler([
30+
new Response(status: 200, body: $body),
31+
]);
32+
$stack = HandlerStack::create($mockHandler);
33+
34+
$provider = new OpenAIImage(
35+
key: 'test-key',
36+
model: 'gpt-image-1',
37+
output_format: 'webp',
38+
httpClient: new GuzzleHttpClient(handler: $stack),
39+
);
40+
41+
$message = $provider->chat(new UserMessage('A cat'));
42+
43+
$this->assertInstanceOf(AssistantMessage::class, $message);
44+
$blocks = $message->getContentBlocks();
45+
$this->assertInstanceOf(ImageContent::class, $blocks[0]);
46+
$this->assertSame('FINAL_BASE64', $blocks[0]->content);
47+
}
48+
49+
public function test_stream_yields_partials_and_uses_completed_event_for_final_image(): void
50+
{
51+
// Mock OpenAI Images streaming response:
52+
// - partial_image_index is an integer per the API spec.
53+
// - `completed` is the authoritative final image; no partial_image_index.
54+
$streamBody = 'data: {"type":"image_generation.partial_image","partial_image_index":0,"b64_json":"PARTIAL_ONE"}'."\n\n";
55+
$streamBody .= 'data: {"type":"image_generation.partial_image","partial_image_index":1,"b64_json":"PARTIAL_TWO"}'."\n\n";
56+
$streamBody .= 'data: {"type":"image_generation.completed","b64_json":"FINAL","usage":{"input_tokens":15,"output_tokens":100,"total_tokens":115}}'."\n\n";
57+
$streamBody .= "data: [DONE]\n\n";
58+
59+
$mockHandler = new MockHandler([
60+
new Response(status: 200, body: $streamBody),
61+
]);
62+
$stack = HandlerStack::create($mockHandler);
63+
64+
$provider = new OpenAIImage(
65+
key: 'test-key',
66+
model: 'gpt-image-1',
67+
output_format: 'webp',
68+
parameters: ['partial_images' => 2],
69+
httpClient: new GuzzleHttpClient(handler: $stack),
70+
);
71+
72+
$generator = $provider->stream(new UserMessage('A cat'));
73+
74+
$chunks = [];
75+
foreach ($generator as $chunk) {
76+
$chunks[] = $chunk;
77+
}
78+
79+
$message = $generator->getReturn();
80+
81+
// Only partial events yield chunks; the completed event does not.
82+
$this->assertCount(2, $chunks);
83+
foreach ($chunks as $chunk) {
84+
$this->assertInstanceOf(ImageChunk::class, $chunk);
85+
// messageId is a non-empty string (StreamChunk contract).
86+
$this->assertIsString($chunk->messageId);
87+
$this->assertNotEmpty($chunk->messageId);
88+
}
89+
90+
// All chunks in one stream share the same messageId.
91+
$this->assertSame($chunks[0]->messageId, $chunks[1]->messageId);
92+
93+
// messageId uses the library-wide `msg_` prefix for consistency with
94+
// OpenAI text and audio providers.
95+
$this->assertStringStartsWith('msg_', $chunks[0]->messageId);
96+
97+
// Chunk contents come directly from the API payload.
98+
$this->assertSame('PARTIAL_ONE', $chunks[0]->content);
99+
$this->assertSame('PARTIAL_TWO', $chunks[1]->content);
100+
101+
// The final assistant message uses the completed image, not the last partial.
102+
$this->assertInstanceOf(AssistantMessage::class, $message);
103+
$image = $message->getContentBlocks()[0];
104+
$this->assertInstanceOf(ImageContent::class, $image);
105+
$this->assertSame('FINAL', $image->content);
106+
$this->assertSame(SourceType::BASE64, $image->sourceType);
107+
$this->assertSame('image/webp', $image->mediaType);
108+
109+
// Usage captured from the completed event only.
110+
$this->assertSame(15, $message->getUsage()->inputTokens);
111+
$this->assertSame(100, $message->getUsage()->outputTokens);
112+
}
113+
114+
public function test_stream_supports_partial_images_zero_returning_only_completed(): void
115+
{
116+
$streamBody = 'data: {"type":"image_generation.completed","b64_json":"ONLY","usage":{"input_tokens":5,"output_tokens":42}}'."\n\n";
117+
$streamBody .= "data: [DONE]\n\n";
118+
119+
$mockHandler = new MockHandler([
120+
new Response(status: 200, body: $streamBody),
121+
]);
122+
$stack = HandlerStack::create($mockHandler);
123+
124+
$provider = new OpenAIImage(
125+
key: 'test-key',
126+
model: 'gpt-image-1',
127+
httpClient: new GuzzleHttpClient(handler: $stack),
128+
);
129+
130+
$generator = $provider->stream(new UserMessage('A dog'));
131+
132+
$chunks = [];
133+
foreach ($generator as $chunk) {
134+
$chunks[] = $chunk;
135+
}
136+
137+
$message = $generator->getReturn();
138+
139+
$this->assertCount(0, $chunks);
140+
$image = $message->getContentBlocks()[0];
141+
$this->assertInstanceOf(ImageContent::class, $image);
142+
$this->assertSame('ONLY', $image->content);
143+
$this->assertSame(42, $message->getUsage()->outputTokens);
144+
}
145+
146+
public function test_stream_sends_stream_flag_and_parameters_in_request_body(): void
147+
{
148+
$sentRequests = [];
149+
$history = Middleware::history($sentRequests);
150+
151+
$streamBody = 'data: {"type":"image_generation.completed","b64_json":"ONLY"}'."\n\n";
152+
$streamBody .= "data: [DONE]\n\n";
153+
154+
$mockHandler = new MockHandler([
155+
new Response(status: 200, body: $streamBody),
156+
]);
157+
$stack = HandlerStack::create($mockHandler);
158+
$stack->push($history);
159+
160+
$provider = new OpenAIImage(
161+
key: 'test-key',
162+
model: 'gpt-image-1',
163+
output_format: 'png',
164+
parameters: ['partial_images' => 3, 'size' => '1024x1024'],
165+
httpClient: new GuzzleHttpClient(handler: $stack),
166+
);
167+
168+
foreach ($provider->stream(new UserMessage('A fox')) as $chunk) {
169+
// drain generator
170+
}
171+
172+
$this->assertCount(1, $sentRequests);
173+
$body = json_decode((string) $sentRequests[0]['request']->getBody(), true);
174+
$this->assertTrue($body['stream']);
175+
$this->assertSame(3, $body['partial_images']);
176+
$this->assertSame('1024x1024', $body['size']);
177+
$this->assertSame('A fox', $body['prompt']);
178+
}
179+
180+
public function test_stream_throws_provider_exception_on_error_event(): void
181+
{
182+
// Real error payload shape as sent by the OpenAI Images streaming API.
183+
$streamBody = 'data: {"type":"error","error":{"type":"image_generation_server_error","code":"image_generation_failed","message":"Image generation failed","param":null}}'."\n\n";
184+
$streamBody .= "data: [DONE]\n\n";
185+
186+
$mockHandler = new MockHandler([
187+
new Response(status: 200, body: $streamBody),
188+
]);
189+
$stack = HandlerStack::create($mockHandler);
190+
191+
$provider = new OpenAIImage(
192+
key: 'test-key',
193+
model: 'gpt-image-1',
194+
httpClient: new GuzzleHttpClient(handler: $stack),
195+
);
196+
197+
$this->expectException(ProviderException::class);
198+
$this->expectExceptionMessage('Image generation failed');
199+
200+
foreach ($provider->stream(new UserMessage('A fox')) as $chunk) {
201+
// drain generator — should throw before yielding
202+
}
203+
}
204+
205+
public function test_stream_throws_when_stream_closes_before_completed_event(): void
206+
{
207+
// Stream drops after partials — no completed event. The old code would
208+
// silently return the last partial as the "final" image.
209+
$streamBody = 'data: {"type":"image_generation.partial_image","partial_image_index":0,"b64_json":"PARTIAL_ONE"}'."\n\n";
210+
$streamBody .= "data: [DONE]\n\n";
211+
212+
$mockHandler = new MockHandler([
213+
new Response(status: 200, body: $streamBody),
214+
]);
215+
$stack = HandlerStack::create($mockHandler);
216+
217+
$provider = new OpenAIImage(
218+
key: 'test-key',
219+
model: 'gpt-image-1',
220+
httpClient: new GuzzleHttpClient(handler: $stack),
221+
);
222+
223+
$this->expectException(ProviderException::class);
224+
$this->expectExceptionMessage('stream ended before a completed image');
225+
226+
$generator = $provider->stream(new UserMessage('A fox'));
227+
foreach ($generator as $chunk) {
228+
// drain
229+
}
230+
$generator->getReturn(); // triggers the post-loop assertion
231+
}
232+
233+
public function test_stream_throws_when_completed_event_has_no_b64_json(): void
234+
{
235+
$streamBody = 'data: {"type":"image_generation.completed"}'."\n\n";
236+
$streamBody .= "data: [DONE]\n\n";
237+
238+
$mockHandler = new MockHandler([
239+
new Response(status: 200, body: $streamBody),
240+
]);
241+
$stack = HandlerStack::create($mockHandler);
242+
243+
$provider = new OpenAIImage(
244+
key: 'test-key',
245+
model: 'gpt-image-1',
246+
httpClient: new GuzzleHttpClient(handler: $stack),
247+
);
248+
249+
$this->expectException(ProviderException::class);
250+
$this->expectExceptionMessage('completed image event without b64_json');
251+
252+
foreach ($provider->stream(new UserMessage('A fox')) as $chunk) {
253+
// drain
254+
}
255+
}
256+
}

0 commit comments

Comments
 (0)
⚡