104 lines
3.1 KiB
Go
104 lines
3.1 KiB
Go
package codex
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
)
|
|
|
|
const CanvasImageBytesLimit = 20 << 20
|
|
const canvasArtifactResponseLimit = 32 << 20
|
|
|
|
type CanvasImage struct {
|
|
MIME string `json:"mime_type"`
|
|
Data []byte `json:"data"`
|
|
}
|
|
type CanvasOutput struct {
|
|
Kind string `json:"kind"`
|
|
Text string `json:"text,omitempty"`
|
|
Image *CanvasImage `json:"image,omitempty"`
|
|
}
|
|
|
|
// Encoded image bytes exist only on the private Unix IPC channel. Public Canvas
|
|
// jobs expose stable asset IDs/URLs and never serialize these bytes to browsers.
|
|
func (c *CanvasClient) GenerateOutput(ctx context.Context, input CanvasTask) (CanvasOutput, error) {
|
|
if c == nil {
|
|
return CanvasOutput{}, errors.New("codex_unavailable")
|
|
}
|
|
if len(input.Prompt) > CanvasInputLimit || len(input.Images) > 4 {
|
|
return CanvasOutput{}, errors.New("context_too_large")
|
|
}
|
|
total := 0
|
|
for _, image := range input.Images {
|
|
total += len(image.Data)
|
|
}
|
|
if total > CanvasImageBytesLimit {
|
|
return CanvasOutput{}, errors.New("context_too_large")
|
|
}
|
|
raw, err := json.Marshal(input)
|
|
if err != nil {
|
|
return CanvasOutput{}, err
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, "POST", "http://codex/canvas/exec", bytes.NewReader(raw))
|
|
if err != nil {
|
|
return CanvasOutput{}, err
|
|
}
|
|
response, err := c.client.Do(request)
|
|
if err != nil {
|
|
return CanvasOutput{}, err
|
|
}
|
|
defer response.Body.Close()
|
|
raw, err = io.ReadAll(io.LimitReader(response.Body, canvasArtifactResponseLimit+1))
|
|
if err != nil || len(raw) > canvasArtifactResponseLimit {
|
|
return CanvasOutput{}, errors.New("generation_failed")
|
|
}
|
|
var result struct {
|
|
Result string `json:"result"`
|
|
Output *CanvasOutput `json:"output"`
|
|
Error string `json:"error"`
|
|
}
|
|
if json.Unmarshal(raw, &result) != nil {
|
|
return CanvasOutput{}, errors.New("generation_failed")
|
|
}
|
|
if response.StatusCode != 200 {
|
|
switch result.Error {
|
|
case "codex_timeout":
|
|
return CanvasOutput{}, errors.New("generation_timeout")
|
|
case "codex_busy":
|
|
if response.StatusCode == 429 {
|
|
return CanvasOutput{}, errors.New("codex_busy")
|
|
}
|
|
return CanvasOutput{}, errors.New("generation_failed")
|
|
case "context_too_large":
|
|
return CanvasOutput{}, errors.New("context_too_large")
|
|
case "model_unavailable":
|
|
return CanvasOutput{}, errors.New("model_unavailable")
|
|
default:
|
|
return CanvasOutput{}, errors.New("generation_failed")
|
|
}
|
|
}
|
|
if result.Output == nil {
|
|
if input.OutputModality == "image" {
|
|
return CanvasOutput{}, errors.New("generation_failed")
|
|
}
|
|
return CanvasOutput{Kind: "text", Text: result.Result}, nil
|
|
}
|
|
output := *result.Output
|
|
switch output.Kind {
|
|
case "text":
|
|
if len(output.Text) > maxOutput || output.Image != nil || input.OutputModality == "image" {
|
|
return CanvasOutput{}, errors.New("generation_failed")
|
|
}
|
|
case "image":
|
|
if output.Image == nil || len(output.Image.Data) == 0 || len(output.Image.Data) > CanvasImageBytesLimit || input.OutputModality != "image" {
|
|
return CanvasOutput{}, errors.New("generation_failed")
|
|
}
|
|
default:
|
|
return CanvasOutput{}, errors.New("generation_failed")
|
|
}
|
|
return output, nil
|
|
}
|