Initial open-source release (MIT): image2api AI gateway
Full Go backend + Vue 3 frontend, OpenAI-compatible API, multi-provider account pools, billing/admin, Docker one-command deploy with auto HTTPS. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Claude Opus 4.8
@@ -0,0 +1,396 @@
|
||||
// Package leonardo implements the Leonardo.ai (app.leonardo.ai) provider client.
|
||||
// Unlike chatgpt/runway (whose JWT IS the stored credential), Leonardo's durable
|
||||
// credential is the browser COOKIE (better-auth session): the bearer access token
|
||||
// it mints lives only ~1h. So every call here takes the cookie and derives a
|
||||
// fresh JWT on the fly via /api/auth/get-session — there is no long-lived token to
|
||||
// store or a separate refresh profile to maintain. tls-client gives a Chrome
|
||||
// JA3/JA4 fingerprint so the requests aren't flagged.
|
||||
package leonardo
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
http "github.com/bogdanfinn/fhttp"
|
||||
tlsclient "github.com/bogdanfinn/tls-client"
|
||||
"github.com/bogdanfinn/tls-client/profiles"
|
||||
)
|
||||
|
||||
const (
|
||||
appBase = "https://app.leonardo.ai"
|
||||
graphqlURL = "https://api.leonardo.ai/v1/graphql"
|
||||
schemaVersion = "1.187.0"
|
||||
userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrAuth = errors.New("leonardo auth failed")
|
||||
ErrQuotaExhausted = errors.New("leonardo quota exhausted")
|
||||
ErrTemporaryUpstream = errors.New("leonardo upstream temporary error")
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
proxy string
|
||||
// sessions caches the short-lived access token per cookie so we don't hit
|
||||
// /api/auth/get-session on every call — Leonardo rate-limits that endpoint
|
||||
// (429) hard, so re-using the ~1h JWT is essential.
|
||||
mu sync.Mutex
|
||||
sessions map[string]*Session
|
||||
}
|
||||
|
||||
func NewClient(proxy string) *Client {
|
||||
return &Client{proxy: strings.TrimSpace(proxy), sessions: map[string]*Session{}}
|
||||
}
|
||||
|
||||
func (c *Client) SetProxy(proxy string) {
|
||||
c.proxy = strings.TrimSpace(proxy)
|
||||
}
|
||||
|
||||
// IsLeonardoCookie reports whether a pasted credential is a Leonardo cookie: it
|
||||
// carries the better-auth session cookie name. This is what disambiguates it from
|
||||
// an Adobe cookie at import time.
|
||||
func IsLeonardoCookie(value string) bool {
|
||||
return strings.Contains(value, "__Secure-better-auth.session_token") ||
|
||||
strings.Contains(value, "better-auth.session_data")
|
||||
}
|
||||
|
||||
// Session is the result of /api/auth/get-session: the short-lived bearer plus the
|
||||
// ids the GraphQL API needs (cognitoSub for the quota query, userId for the feed
|
||||
// and the CDN image path) and the human-facing account fields.
|
||||
type Session struct {
|
||||
AccessToken string
|
||||
CognitoSub string
|
||||
UserID string
|
||||
Email string
|
||||
Name string
|
||||
ExpiresAt int64
|
||||
}
|
||||
|
||||
// GetSession exchanges the cookie for a fresh access token + account ids. A 401/403
|
||||
// (or a response with no access token) means the cookie/session is dead → ErrAuth.
|
||||
func (c *Client) GetSession(ctx context.Context, cookie string) (*Session, error) {
|
||||
cookie = strings.TrimSpace(cookie)
|
||||
if cookie == "" {
|
||||
return nil, ErrAuth
|
||||
}
|
||||
// Re-use a cached, still-valid access token (keep a 60s safety margin) instead
|
||||
// of hitting the heavily rate-limited get-session endpoint again.
|
||||
c.mu.Lock()
|
||||
if cs, ok := c.sessions[cookie]; ok && cs.ExpiresAt-60 > time.Now().Unix() {
|
||||
c.mu.Unlock()
|
||||
return cs, nil
|
||||
}
|
||||
c.mu.Unlock()
|
||||
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodGet, appBase+"/api/auth/get-session", nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"*/*"},
|
||||
"accept-language": {"en-US,en;q=0.9"},
|
||||
"cookie": {cookie},
|
||||
"origin": {appBase},
|
||||
"referer": {appBase + "/"},
|
||||
"user-agent": {userAgent},
|
||||
"sec-fetch-dest": {"empty"},
|
||||
"sec-fetch-mode": {"cors"},
|
||||
"sec-fetch-site": {"same-origin"},
|
||||
http.HeaderOrderKey: {
|
||||
"accept", "accept-language", "cookie", "origin", "referer",
|
||||
"user-agent", "sec-fetch-dest", "sec-fetch-mode", "sec-fetch-site",
|
||||
},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode == 401 || resp.StatusCode == 403 {
|
||||
return nil, ErrAuth
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return nil, fmt.Errorf("%w: get-session http %d: %s", ErrTemporaryUpstream, resp.StatusCode, clip(body, 160))
|
||||
}
|
||||
var raw struct {
|
||||
Session struct {
|
||||
AccessToken string `json:"accessToken"`
|
||||
CognitoSub string `json:"cognitoSub"`
|
||||
UserID string `json:"userId"`
|
||||
HasuraUserID string `json:"hasuraUserId"`
|
||||
TokenExpiry int64 `json:"accessTokenExpiry"`
|
||||
} `json:"session"`
|
||||
User struct {
|
||||
ID string `json:"id"`
|
||||
Email string `json:"email"`
|
||||
Name string `json:"name"`
|
||||
} `json:"user"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &raw); err != nil {
|
||||
return nil, fmt.Errorf("%w: get-session non-json", ErrTemporaryUpstream)
|
||||
}
|
||||
if strings.TrimSpace(raw.Session.AccessToken) == "" {
|
||||
// No bearer despite 200 → the cookie no longer authenticates.
|
||||
return nil, ErrAuth
|
||||
}
|
||||
uid := raw.Session.UserID
|
||||
if uid == "" {
|
||||
uid = raw.Session.HasuraUserID
|
||||
}
|
||||
if uid == "" {
|
||||
uid = raw.User.ID
|
||||
}
|
||||
sess := &Session{
|
||||
AccessToken: raw.Session.AccessToken,
|
||||
CognitoSub: raw.Session.CognitoSub,
|
||||
UserID: uid,
|
||||
Email: strings.TrimSpace(raw.User.Email),
|
||||
Name: strings.TrimSpace(raw.User.Name),
|
||||
ExpiresAt: raw.Session.TokenExpiry,
|
||||
}
|
||||
if sess.ExpiresAt > time.Now().Unix() {
|
||||
c.mu.Lock()
|
||||
c.sessions[cookie] = sess
|
||||
c.mu.Unlock()
|
||||
}
|
||||
return sess, nil
|
||||
}
|
||||
|
||||
const qGetTokens = `query GetUserTokensFromSub($sub: String) {
|
||||
user_details(where: {cognitoId: {_eq: $sub}}) {
|
||||
id
|
||||
plan
|
||||
subscriptionTokens
|
||||
paidTokens
|
||||
rolloverTokens
|
||||
tokenRenewalDate
|
||||
__typename
|
||||
}
|
||||
}`
|
||||
|
||||
// FetchCreditsBalance derives a JWT from the cookie then reads the account's image
|
||||
// token balance. Returns a normalized map mirroring the other providers so the
|
||||
// TokenService quota plumbing is uniform. remaining = subscription+paid+rollover
|
||||
// (the spendable image tokens); available_until carries the daily renewal time so
|
||||
// the maintenance sweep can auto-recover a 限额 account.
|
||||
func (c *Client) FetchCreditsBalance(ctx context.Context, cookie string) (map[string]any, error) {
|
||||
sess, err := c.GetSession(ctx, cookie)
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrAuth) {
|
||||
return nil, ErrAuth
|
||||
}
|
||||
return unknownBalance(err.Error()), nil
|
||||
}
|
||||
if sess.CognitoSub == "" {
|
||||
return unknownBalance("no cognitoSub"), nil
|
||||
}
|
||||
|
||||
payload, _ := json.Marshal(map[string]any{
|
||||
"operationName": "GetUserTokensFromSub",
|
||||
"variables": map[string]any{"sub": sess.CognitoSub},
|
||||
"query": qGetTokens,
|
||||
})
|
||||
body, status, err := c.graphql(ctx, sess.AccessToken, payload)
|
||||
if err != nil {
|
||||
return unknownBalance("network: " + err.Error()), nil
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return nil, ErrAuth
|
||||
}
|
||||
if status != 200 {
|
||||
return unknownBalance(fmt.Sprintf("http %d: %s", status, clip(body, 160))), nil
|
||||
}
|
||||
var result struct {
|
||||
Data struct {
|
||||
UserDetails []struct {
|
||||
Plan string `json:"plan"`
|
||||
SubscriptionTokens int `json:"subscriptionTokens"`
|
||||
PaidTokens int `json:"paidTokens"`
|
||||
RolloverTokens int `json:"rolloverTokens"`
|
||||
TokenRenewalDate string `json:"tokenRenewalDate"`
|
||||
} `json:"user_details"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &result); err != nil {
|
||||
return unknownBalance("non-json"), nil
|
||||
}
|
||||
if len(result.Data.UserDetails) == 0 {
|
||||
return unknownBalance("no user_details"), nil
|
||||
}
|
||||
ud := result.Data.UserDetails[0]
|
||||
remaining := ud.SubscriptionTokens + ud.PaidTokens + ud.RolloverTokens
|
||||
return map[string]any{
|
||||
"remaining": remaining,
|
||||
"used": nil,
|
||||
"total": nil,
|
||||
"unknown": false,
|
||||
"error": nil,
|
||||
"plan": ud.Plan,
|
||||
"available_until": strings.TrimSpace(ud.TokenRenewalDate),
|
||||
"email": emptyStringNil(sess.Email),
|
||||
"display_name": emptyStringNil(sess.Name),
|
||||
"user_id": emptyStringNil(sess.UserID),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// graphql POSTs a GraphQL body to the Leonardo API with the bearer + schema header,
|
||||
// returning the raw response body and status.
|
||||
func (c *Client) graphql(ctx context.Context, accessToken string, payload []byte) ([]byte, int, error) {
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, graphqlURL, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"content-type": {"application/json"},
|
||||
"accept": {"*/*"},
|
||||
"accept-language": {"en-US,en;q=0.9"},
|
||||
"origin": {appBase},
|
||||
"referer": {appBase + "/"},
|
||||
"user-agent": {userAgent},
|
||||
"authorization": {"Bearer " + accessToken},
|
||||
"x-leo-schema-version": {schemaVersion},
|
||||
"sec-fetch-dest": {"empty"},
|
||||
"sec-fetch-mode": {"cors"},
|
||||
"sec-fetch-site": {"same-site"},
|
||||
http.HeaderOrderKey: {
|
||||
"content-type", "accept", "accept-language", "origin", "referer",
|
||||
"user-agent", "authorization", "x-leo-schema-version",
|
||||
"sec-fetch-dest", "sec-fetch-mode", "sec-fetch-site",
|
||||
},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, resp.StatusCode, err
|
||||
}
|
||||
return body, resp.StatusCode, nil
|
||||
}
|
||||
|
||||
func unknownBalance(reason string) map[string]any {
|
||||
return map[string]any{
|
||||
"remaining": nil,
|
||||
"used": nil,
|
||||
"total": nil,
|
||||
"unknown": true,
|
||||
"error": reason,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) newTLSClient() (tlsclient.HttpClient, error) {
|
||||
// Match the fingerprint proven to work against Leonardo's Cloudflare edge:
|
||||
// Chrome_120, fixed extension order. A randomized JA3 (Chrome_133 +
|
||||
// WithRandomTLSExtensionOrder) gets flagged and 429'd at get-session.
|
||||
options := []tlsclient.HttpClientOption{
|
||||
tlsclient.WithTimeoutSeconds(60),
|
||||
tlsclient.WithClientProfile(profiles.Chrome_120),
|
||||
}
|
||||
if c.proxy != "" {
|
||||
options = append(options, tlsclient.WithProxyUrl(c.proxy))
|
||||
}
|
||||
return tlsclient.NewHttpClient(tlsclient.NewNoopLogger(), options...)
|
||||
}
|
||||
|
||||
// downloadImage fetches a generated image (cdn.leonardo.ai) and returns the bytes.
|
||||
func (c *Client) downloadImage(ctx context.Context, imageURL string) ([]byte, error) {
|
||||
if _, err := url.Parse(imageURL); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodGet, imageURL, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"accept": {"image/avif,image/webp,image/png,image/*,*/*;q=0.8"},
|
||||
"user-agent": {userAgent},
|
||||
"referer": {appBase + "/"},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
return nil, fmt.Errorf("%w: image download http %d", ErrTemporaryUpstream, resp.StatusCode)
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
func stringValue(v any) string {
|
||||
switch x := v.(type) {
|
||||
case string:
|
||||
return x
|
||||
case nil:
|
||||
return ""
|
||||
default:
|
||||
b, _ := json.Marshal(x)
|
||||
return strings.TrimSpace(string(b))
|
||||
}
|
||||
}
|
||||
|
||||
func intValue(v any) int {
|
||||
switch x := v.(type) {
|
||||
case int:
|
||||
return x
|
||||
case int64:
|
||||
return int(x)
|
||||
case float64:
|
||||
return int(x)
|
||||
case json.Number:
|
||||
n, _ := x.Int64()
|
||||
return int(n)
|
||||
case string:
|
||||
n, _ := strconv.Atoi(strings.TrimSpace(x))
|
||||
return n
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func emptyStringNil(v string) any {
|
||||
v = strings.TrimSpace(v)
|
||||
if v == "" {
|
||||
return nil
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func clip(b []byte, n int) string {
|
||||
s := strings.TrimSpace(string(b))
|
||||
if len(s) > n {
|
||||
return s[:n]
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,326 @@
|
||||
package leonardo
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"mime/multipart"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
http "github.com/bogdanfinn/fhttp"
|
||||
)
|
||||
|
||||
// defaultStyleID is the "Dynamic" style applied when the caller doesn't specify
|
||||
// one — Leonardo's Generate mutation expects a style_ids entry.
|
||||
const defaultStyleID = "111dc692-d470-4eec-b791-3475abac4c46"
|
||||
|
||||
const mGenerate = `mutation Generate($request: CreateGenerationRequest!) {
|
||||
generate(request: $request) {
|
||||
apiCreditCost
|
||||
generationId
|
||||
__typename
|
||||
}
|
||||
}`
|
||||
|
||||
// qGenerationImages polls one generation's status AND its produced images in a
|
||||
// single round-trip (where: id _in [genId]).
|
||||
const qGenerationImages = `query GenerationImages($where: generations_bool_exp = {}) {
|
||||
generations(where: $where) {
|
||||
id
|
||||
status
|
||||
generated_images {
|
||||
id
|
||||
url
|
||||
__typename
|
||||
}
|
||||
__typename
|
||||
}
|
||||
}`
|
||||
|
||||
const mUploadImage = `mutation UploadImage($uploadImageInput: UploadImageInput!) {
|
||||
uploadImage(arg1: $uploadImageInput) {
|
||||
uploadId
|
||||
url
|
||||
fields
|
||||
__typename
|
||||
}
|
||||
}`
|
||||
|
||||
// uploadInitImage uploads a reference (init) image for image-to-image: it asks
|
||||
// Leonardo for a presigned S3 POST, uploads the bytes, and returns the upload id
|
||||
// to reference in the Generate request's image_reference guidance.
|
||||
func (c *Client) uploadInitImage(ctx context.Context, accessToken string, img []byte) (string, error) {
|
||||
payload, _ := json.Marshal(map[string]any{
|
||||
"operationName": "UploadImage",
|
||||
"query": mUploadImage,
|
||||
"variables": map[string]any{"uploadImageInput": map[string]any{"uploadType": "INIT", "extension": "png"}},
|
||||
})
|
||||
body, status, err := c.graphql(ctx, accessToken, payload)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: upload-init: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return "", ErrAuth
|
||||
}
|
||||
if status != 200 {
|
||||
return "", fmt.Errorf("%w: upload-init http %d: %s", ErrTemporaryUpstream, status, clip(body, 160))
|
||||
}
|
||||
if e := graphqlError(body); e != nil {
|
||||
return "", e
|
||||
}
|
||||
var ur struct {
|
||||
Data struct {
|
||||
UploadImage struct {
|
||||
UploadID string `json:"uploadId"`
|
||||
URL string `json:"url"`
|
||||
Fields string `json:"fields"`
|
||||
} `json:"uploadImage"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &ur); err != nil {
|
||||
return "", fmt.Errorf("%w: upload-init non-json", ErrTemporaryUpstream)
|
||||
}
|
||||
up := ur.Data.UploadImage
|
||||
if up.UploadID == "" || up.URL == "" {
|
||||
return "", fmt.Errorf("%w: no upload url", ErrTemporaryUpstream)
|
||||
}
|
||||
var fields map[string]string
|
||||
if err := json.Unmarshal([]byte(up.Fields), &fields); err != nil {
|
||||
return "", fmt.Errorf("%w: bad upload fields", ErrTemporaryUpstream)
|
||||
}
|
||||
|
||||
// Presigned S3 POST: all policy fields first, the file part LAST.
|
||||
var buf bytes.Buffer
|
||||
w := multipart.NewWriter(&buf)
|
||||
for k, v := range fields {
|
||||
_ = w.WriteField(k, v)
|
||||
}
|
||||
fw, err := w.CreateFormFile("file", "image.png")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, err := fw.Write(img); err != nil {
|
||||
return "", err
|
||||
}
|
||||
_ = w.Close()
|
||||
|
||||
client, err := c.newTLSClient()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, up.URL, &buf)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header = http.Header{
|
||||
"content-type": {w.FormDataContentType()},
|
||||
"user-agent": {userAgent},
|
||||
"origin": {appBase},
|
||||
"referer": {appBase + "/"},
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: s3 upload: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != 204 && resp.StatusCode != 200 && resp.StatusCode != 201 {
|
||||
return "", fmt.Errorf("%w: s3 upload http %d", ErrTemporaryUpstream, resp.StatusCode)
|
||||
}
|
||||
return up.UploadID, nil
|
||||
}
|
||||
|
||||
// GenerateImage runs the full Leonardo image pipeline against one account cookie:
|
||||
// mint a JWT, (for image-to-image) upload each reference image, submit the
|
||||
// Generate mutation, poll until COMPLETE, then download the first produced image.
|
||||
// Returns the image bytes, an info map, and a classified error.
|
||||
func (c *Client) GenerateImage(ctx context.Context, cookie, model, prompt string, width, height int, styleIDs []string, refImages [][]byte) ([]byte, map[string]any, error) {
|
||||
sess, err := c.GetSession(ctx, cookie)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if len(styleIDs) == 0 {
|
||||
styleIDs = []string{defaultStyleID}
|
||||
}
|
||||
if strings.TrimSpace(model) == "" {
|
||||
model = "seedream-4.5"
|
||||
}
|
||||
|
||||
// Image-to-image: upload each reference and collect its guidance entry.
|
||||
var imageRefs []map[string]any
|
||||
for _, img := range refImages {
|
||||
if len(img) == 0 {
|
||||
continue
|
||||
}
|
||||
uploadID, upErr := c.uploadInitImage(ctx, sess.AccessToken, img)
|
||||
if upErr != nil {
|
||||
return nil, nil, upErr
|
||||
}
|
||||
imageRefs = append(imageRefs, map[string]any{
|
||||
"image": map[string]any{"id": uploadID, "type": "UPLOADED"},
|
||||
"strength": "MID",
|
||||
})
|
||||
}
|
||||
|
||||
promptEnhance := "AUTO"
|
||||
parameters := map[string]any{
|
||||
"height": height,
|
||||
"width": width,
|
||||
"prompt_enhance": promptEnhance,
|
||||
"quantity": 1,
|
||||
"style_ids": styleIDs,
|
||||
"prompt": prompt,
|
||||
}
|
||||
if len(imageRefs) > 0 {
|
||||
// Preserve the reference when image-guided (matches the web app).
|
||||
parameters["prompt_enhance"] = "OFF"
|
||||
parameters["guidances"] = map[string]any{"image_reference": imageRefs}
|
||||
}
|
||||
|
||||
// 1. submit
|
||||
genReq := map[string]any{
|
||||
"operationName": "Generate",
|
||||
"query": mGenerate,
|
||||
"variables": map[string]any{
|
||||
"request": map[string]any{
|
||||
"model": model,
|
||||
"public": true,
|
||||
"parameters": parameters,
|
||||
},
|
||||
},
|
||||
}
|
||||
payload, _ := json.Marshal(genReq)
|
||||
body, status, err := c.graphql(ctx, sess.AccessToken, payload)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return nil, nil, ErrAuth
|
||||
}
|
||||
if status != 200 {
|
||||
return nil, nil, fmt.Errorf("%w: generate http %d: %s", ErrTemporaryUpstream, status, clip(body, 200))
|
||||
}
|
||||
if e := graphqlError(body); e != nil {
|
||||
return nil, nil, e
|
||||
}
|
||||
var genResp struct {
|
||||
Data struct {
|
||||
Generate struct {
|
||||
GenerationID string `json:"generationId"`
|
||||
} `json:"generate"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &genResp); err != nil {
|
||||
return nil, nil, fmt.Errorf("%w: generate non-json", ErrTemporaryUpstream)
|
||||
}
|
||||
genID := strings.TrimSpace(genResp.Data.Generate.GenerationID)
|
||||
if genID == "" {
|
||||
return nil, nil, fmt.Errorf("%w: no generationId: %s", ErrTemporaryUpstream, clip(body, 200))
|
||||
}
|
||||
|
||||
// 2. poll until COMPLETE, then read the image url.
|
||||
imageURL, err := c.pollImage(ctx, sess.AccessToken, genID)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
// 3. download bytes
|
||||
data, err := c.downloadImage(ctx, imageURL)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
info := map[string]any{
|
||||
"generation_id": genID,
|
||||
"image_url": imageURL,
|
||||
"user_id": sess.UserID,
|
||||
}
|
||||
return data, info, nil
|
||||
}
|
||||
|
||||
// pollImage polls one generation until it reports COMPLETE (returning the first
|
||||
// image url) or FAILED (error). Honors ctx cancellation / deadline.
|
||||
func (c *Client) pollImage(ctx context.Context, accessToken, genID string) (string, error) {
|
||||
payload, _ := json.Marshal(map[string]any{
|
||||
"operationName": "GenerationImages",
|
||||
"query": qGenerationImages,
|
||||
"variables": map[string]any{
|
||||
"where": map[string]any{"id": map[string]any{"_in": []string{genID}}},
|
||||
},
|
||||
})
|
||||
|
||||
ticker := time.NewTicker(3 * time.Second)
|
||||
defer ticker.Stop()
|
||||
// Cap the wait independent of the parent deadline so a stuck job can't hang.
|
||||
deadline := time.Now().Add(5 * time.Minute)
|
||||
|
||||
for {
|
||||
body, status, err := c.graphql(ctx, accessToken, payload)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: poll: %s", ErrTemporaryUpstream, err.Error())
|
||||
}
|
||||
if status == 401 || status == 403 {
|
||||
return "", ErrAuth
|
||||
}
|
||||
if status == 200 {
|
||||
var pr struct {
|
||||
Data struct {
|
||||
Generations []struct {
|
||||
Status string `json:"status"`
|
||||
GeneratedImages []struct {
|
||||
URL string `json:"url"`
|
||||
} `json:"generated_images"`
|
||||
} `json:"generations"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &pr); err == nil && len(pr.Data.Generations) > 0 {
|
||||
g := pr.Data.Generations[0]
|
||||
switch strings.ToUpper(g.Status) {
|
||||
case "COMPLETE":
|
||||
for _, img := range g.GeneratedImages {
|
||||
if u := strings.TrimSpace(img.URL); u != "" {
|
||||
return u, nil
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf("%w: complete but no image url", ErrTemporaryUpstream)
|
||||
case "FAILED":
|
||||
return "", fmt.Errorf("%w: generation failed", ErrTemporaryUpstream)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if time.Now().After(deadline) {
|
||||
return "", fmt.Errorf("%w: generation timed out", ErrTemporaryUpstream)
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return "", ctx.Err()
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// graphqlError inspects a GraphQL response body for an "errors" array and maps the
|
||||
// first message to a classified sentinel (auth / quota / temporary). Returns nil
|
||||
// when there are no errors.
|
||||
func graphqlError(body []byte) error {
|
||||
var env struct {
|
||||
Errors []struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"errors"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &env); err != nil || len(env.Errors) == 0 {
|
||||
return nil
|
||||
}
|
||||
msg := strings.TrimSpace(env.Errors[0].Message)
|
||||
low := strings.ToLower(msg)
|
||||
switch {
|
||||
case strings.Contains(low, "unauthor") || strings.Contains(low, "jwt") || strings.Contains(low, "token is") || strings.Contains(low, "forbidden"):
|
||||
return ErrAuth
|
||||
case strings.Contains(low, "token") || strings.Contains(low, "credit") || strings.Contains(low, "quota") || strings.Contains(low, "insufficient") || strings.Contains(low, "not enough"):
|
||||
return ErrQuotaExhausted
|
||||
default:
|
||||
return fmt.Errorf("leonardo: %s", clip([]byte(msg), 200))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user