feat(image): stream text-to-image generation with SSE and fallback
Send stream=true for text-to-image requests and consume the SSE response, falling back to a non-streamed request when the stream errors in a way that warrants it. Raises the scanner line limit to handle large base64 image payloads. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -97,6 +97,9 @@ func TestImageClientSendsTransportDefaultSizeWhenUnset(t *testing.T) {
|
||||
if requestBody["size"] != "1024x1024" {
|
||||
t.Fatalf("expected transport default size, got %#v", requestBody["size"])
|
||||
}
|
||||
if requestBody["stream"] != true {
|
||||
t.Fatalf("expected image generations to use stream transport, got %#v", requestBody["stream"])
|
||||
}
|
||||
if _, ok := requestBody["n"]; ok {
|
||||
t.Fatalf("expected single image request to omit n for compatible transports, got %#v", requestBody["n"])
|
||||
}
|
||||
@@ -306,6 +309,112 @@ func TestImageClientParsesAlternateImageResponseShapes(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageClientParsesStreamingImageResponse(t *testing.T) {
|
||||
var requestBody map[string]any
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = w.Write([]byte("event: image_generation.completed\n"))
|
||||
_, _ = w.Write([]byte(`data: {"type":"image_generation.completed","result":"Z2VuZXJhdGVkIGJ5dGVz"}` + "\n\n"))
|
||||
_, _ = w.Write([]byte("data: [DONE]\n\n"))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
|
||||
image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if requestBody["stream"] != true {
|
||||
t.Fatalf("expected stream request body, got %#v", requestBody)
|
||||
}
|
||||
if image.URL != "data:image/png;base64,Z2VuZXJhdGVkIGJ5dGVz" {
|
||||
t.Fatalf("expected streamed base64 image, got %#v", image)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageClientFallsBackWhenStreamIsUnsupported(t *testing.T) {
|
||||
calls := 0
|
||||
streamValues := make([]any, 0, 2)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls++
|
||||
var requestBody map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
streamValues = append(streamValues, requestBody["stream"])
|
||||
if calls == 1 {
|
||||
http.Error(w, `{"error":{"message":"unknown parameter: stream"}}`, http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
|
||||
image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if image.URL != "https://example.com/image.png" {
|
||||
t.Fatalf("expected fallback image, got %#v", image)
|
||||
}
|
||||
if calls != 2 {
|
||||
t.Fatalf("expected stream then json fallback, got %d calls", calls)
|
||||
}
|
||||
if len(streamValues) != 2 || streamValues[0] != true || streamValues[1] != nil {
|
||||
t.Fatalf("expected first request stream=true and fallback without stream, got %#v", streamValues)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageClientFallsBackAfterIdempotentStreamFailures(t *testing.T) {
|
||||
calls := 0
|
||||
streamValues := make([]any, 0, 4)
|
||||
keys := make([]string, 0, 4)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls++
|
||||
keys = append(keys, r.Header.Get("Idempotency-Key"))
|
||||
var requestBody map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
streamValues = append(streamValues, requestBody["stream"])
|
||||
if calls <= 3 {
|
||||
http.Error(w, `{"error":{"message":"temporary stream failure"}}`, http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/recovered.png"}]}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
|
||||
image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{
|
||||
Model: "gpt-image-2",
|
||||
IdempotencyKey: "same-generation",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if image.URL != "https://example.com/recovered.png" {
|
||||
t.Fatalf("expected fallback image, got %#v", image)
|
||||
}
|
||||
if calls != 4 {
|
||||
t.Fatalf("expected three stream attempts then json fallback, got %d calls", calls)
|
||||
}
|
||||
if len(streamValues) != 4 || streamValues[0] != true || streamValues[1] != true || streamValues[2] != true || streamValues[3] != nil {
|
||||
t.Fatalf("expected stream retries then json fallback, got %#v", streamValues)
|
||||
}
|
||||
for _, key := range keys {
|
||||
if key != "same-generation" {
|
||||
t.Fatalf("expected idempotency key to be preserved, got %#v", keys)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageClientRetriesTransientFailures(t *testing.T) {
|
||||
calls := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
Reference in New Issue
Block a user