feat: 自定义 OpenAI 兼容上游聚合 + 权重/并发调度

- custom provider:把任意 v1 端点当账号(base_url+key),按 model id 路由, 直连不走代理;支持 images/generations、images/edits、Sora 式异步视频 - 调度器全局权重优先 + 每账号并发感知(上游账号可配并发,其余系统固定) - 后台:添加/编辑上游、自定义模型表单(每档普通/代理价,留空=不支持)、账号权重/并发列与导入权重 - grok-video / nano-banana-2(runway) 写入硬编码目录(开源自带); 去掉 540p,Adobe 视频最低 720p - UI:SelectMenu 与输入框等高、测试耗时改用秒 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-28 21:59:46 +08:00
parent d289ec378d
commit 3ef84bb9b3
17 changed files with 1149 additions and 64 deletions
@@ -0,0 +1,298 @@
+// Package custom implements a generic OpenAI-compatible upstream client. A
+// "custom" model forwards generation to any OpenAI-compatible API: the upstream
+// base_url + api_key live on a custom account (pool="custom"), the upstream model
+// name on the model config (UpstreamModel). Calls go DIRECT (no tls-client, no
+// proxy) — the upstream is a normal API with no anti-bot.
+package custom
+
+import (
+	"bytes"
+	"context"
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"mime/multipart"
+	"net/http"
+	"strings"
+	"time"
+)
+
+var (
+	ErrAuth              = errors.New("custom upstream auth failed")
+	ErrQuotaExhausted    = errors.New("custom upstream quota exhausted")
+	ErrTemporaryUpstream = errors.New("custom upstream temporary error")
+)
+
+type Client struct{}
+
+func NewClient() *Client { return &Client{} }
+
+func httpClient() *http.Client { return &http.Client{Timeout: 10 * time.Minute} }
+
+// GenerateImage calls the upstream OpenAI image API. With reference images it
+// uses /v1/images/edits (multipart); otherwise /v1/images/generations. Returns
+// the raw image bytes (decoded from b64_json, or downloaded from url).
+func (c *Client) GenerateImage(ctx context.Context, baseURL, apiKey, model, prompt, size, quality string, refs [][]byte) ([]byte, error) {
+	baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
+	if baseURL == "" || apiKey == "" {
+		return nil, ErrAuth
+	}
+	var req *http.Request
+	var err error
+	if len(refs) > 0 {
+		body := &bytes.Buffer{}
+		w := multipart.NewWriter(body)
+		_ = w.WriteField("model", model)
+		_ = w.WriteField("prompt", prompt)
+		if size != "" {
+			_ = w.WriteField("size", size)
+		}
+		for i, r := range refs {
+			fw, e := w.CreateFormFile("image[]", fmt.Sprintf("ref_%d.png", i+1))
+			if e != nil {
+				return nil, e
+			}
+			_, _ = fw.Write(r)
+		}
+		_ = w.Close()
+		req, err = http.NewRequest(http.MethodPost, baseURL+"/v1/images/edits", body)
+		if err != nil {
+			return nil, err
+		}
+		req.Header.Set("Content-Type", w.FormDataContentType())
+	} else {
+		payload := map[string]any{"model": model, "prompt": prompt, "n": 1}
+		if size != "" {
+			payload["size"] = size
+		}
+		if quality != "" {
+			payload["quality"] = quality
+		}
+		raw, _ := json.Marshal(payload)
+		req, err = http.NewRequest(http.MethodPost, baseURL+"/v1/images/generations", bytes.NewReader(raw))
+		if err != nil {
+			return nil, err
+		}
+		req.Header.Set("Content-Type", "application/json")
+	}
+	req = req.WithContext(ctx)
+	req.Header.Set("Authorization", "Bearer "+apiKey)
+
+	resp, err := httpClient().Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
+	}
+	defer resp.Body.Close()
+	body, _ := io.ReadAll(resp.Body)
+	if e := mapStatus(resp.StatusCode, body); e != nil {
+		return nil, e
+	}
+	return imageBytesFromResponse(ctx, body)
+}
+
+// GenerateVideo drives the upstream Sora-style async video API:
+// POST /v1/videos → poll GET /v1/videos/{id} → GET /v1/videos/{id}/content.
+// When downloadResult is false it returns the upstream content URL instead.
+func (c *Client) GenerateVideo(ctx context.Context, baseURL, apiKey, model, prompt, size string, seconds int, downloadResult bool) ([]byte, string, error) {
+	baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
+	if baseURL == "" || apiKey == "" {
+		return nil, "", ErrAuth
+	}
+	payload := map[string]any{"model": model, "prompt": prompt}
+	if size != "" {
+		payload["size"] = size
+	}
+	if seconds > 0 {
+		payload["seconds"] = fmt.Sprintf("%d", seconds)
+	}
+	raw, _ := json.Marshal(payload)
+	created, err := c.doJSON(ctx, http.MethodPost, baseURL+"/v1/videos", apiKey, raw)
+	if err != nil {
+		return nil, "", err
+	}
+	jobID := strings.TrimSpace(stringValue(created["id"]))
+	if jobID == "" {
+		return nil, "", fmt.Errorf("%w: video create missing id", ErrTemporaryUpstream)
+	}
+	// Poll until terminal.
+	for {
+		if err := ctx.Err(); err != nil {
+			return nil, "", err
+		}
+		job, err := c.doJSON(ctx, http.MethodGet, baseURL+"/v1/videos/"+jobID, apiKey, nil)
+		if err != nil {
+			if errors.Is(err, ErrTemporaryUpstream) {
+				if sleepCtx(ctx, 5*time.Second) != nil {
+					return nil, "", ctx.Err()
+				}
+				continue
+			}
+			return nil, "", err
+		}
+		switch strings.ToLower(strings.TrimSpace(stringValue(job["status"]))) {
+		case "completed", "succeeded", "success":
+			contentURL := baseURL + "/v1/videos/" + jobID + "/content"
+			if !downloadResult {
+				return nil, contentURL, nil
+			}
+			data, err := c.download(ctx, contentURL, apiKey)
+			if err != nil {
+				return nil, "", err
+			}
+			return data, contentURL, nil
+		case "failed", "error", "canceled", "cancelled":
+			reason := stringValue(job["error"])
+			if isCreditError(reason) {
+				return nil, "", fmt.Errorf("%w: %s", ErrQuotaExhausted, clip([]byte(reason), 160))
+			}
+			return nil, "", fmt.Errorf("custom: video %s", clip([]byte(reason), 160))
+		}
+		if sleepCtx(ctx, 5*time.Second) != nil {
+			return nil, "", ctx.Err()
+		}
+	}
+}
+
+func (c *Client) doJSON(ctx context.Context, method, url, apiKey string, body []byte) (map[string]any, error) {
+	var reader io.Reader
+	if body != nil {
+		reader = bytes.NewReader(body)
+	}
+	req, err := http.NewRequest(method, url, reader)
+	if err != nil {
+		return nil, err
+	}
+	req = req.WithContext(ctx)
+	req.Header.Set("Authorization", "Bearer "+apiKey)
+	if body != nil {
+		req.Header.Set("Content-Type", "application/json")
+	}
+	resp, err := httpClient().Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
+	}
+	defer resp.Body.Close()
+	raw, _ := io.ReadAll(resp.Body)
+	if e := mapStatus(resp.StatusCode, raw); e != nil {
+		return nil, e
+	}
+	var out map[string]any
+	if len(raw) == 0 {
+		return map[string]any{}, nil
+	}
+	if err := json.Unmarshal(raw, &out); err != nil {
+		return nil, fmt.Errorf("%w: non-json: %s", ErrTemporaryUpstream, clip(raw, 120))
+	}
+	return out, nil
+}
+
+func (c *Client) download(ctx context.Context, url, apiKey string) ([]byte, error) {
+	req, _ := http.NewRequest(http.MethodGet, url, nil)
+	req = req.WithContext(ctx)
+	req.Header.Set("Authorization", "Bearer "+apiKey)
+	resp, err := httpClient().Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return nil, fmt.Errorf("%w: download %d", ErrTemporaryUpstream, resp.StatusCode)
+	}
+	data, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+	if len(data) == 0 {
+		return nil, fmt.Errorf("%w: empty download", ErrTemporaryUpstream)
+	}
+	return data, nil
+}
+
+// imageBytesFromResponse extracts image bytes from an OpenAI images response:
+// data[0].b64_json (preferred) or data[0].url (downloaded).
+func imageBytesFromResponse(ctx context.Context, body []byte) ([]byte, error) {
+	var out struct {
+		Data []struct {
+			B64JSON string `json:"b64_json"`
+			URL     string `json:"url"`
+		} `json:"data"`
+	}
+	if err := json.Unmarshal(body, &out); err != nil || len(out.Data) == 0 {
+		return nil, fmt.Errorf("%w: bad image response: %s", ErrTemporaryUpstream, clip(body, 160))
+	}
+	d := out.Data[0]
+	if d.B64JSON != "" {
+		raw, err := base64.StdEncoding.DecodeString(strings.TrimSpace(d.B64JSON))
+		if err != nil {
+			return nil, fmt.Errorf("%w: bad b64: %v", ErrTemporaryUpstream, err)
+		}
+		return raw, nil
+	}
+	if d.URL != "" {
+		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, d.URL, nil)
+		resp, err := httpClient().Do(req)
+		if err != nil {
+			return nil, fmt.Errorf("%w: %v", ErrTemporaryUpstream, err)
+		}
+		defer resp.Body.Close()
+		return io.ReadAll(resp.Body)
+	}
+	return nil, fmt.Errorf("%w: image response had no b64/url", ErrTemporaryUpstream)
+}
+
+func mapStatus(status int, body []byte) error {
+	switch {
+	case status >= 200 && status < 300:
+		return nil
+	case status == 401 || status == 403:
+		return fmt.Errorf("%w: %d %s", ErrAuth, status, clip(body, 160))
+	case status == 429:
+		return fmt.Errorf("%w: 429 %s", ErrQuotaExhausted, clip(body, 160))
+	case status >= 500:
+		return fmt.Errorf("%w: %d %s", ErrTemporaryUpstream, status, clip(body, 160))
+	default:
+		if isCreditError(string(body)) {
+			return fmt.Errorf("%w: %s", ErrQuotaExhausted, clip(body, 160))
+		}
+		return fmt.Errorf("custom: %d %s", status, clip(body, 160))
+	}
+}
+
+func isCreditError(s string) bool {
+	s = strings.ToLower(s)
+	return strings.Contains(s, "insufficient") || strings.Contains(s, "quota") ||
+		strings.Contains(s, "credit") || strings.Contains(s, "balance")
+}
+
+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 clip(b []byte, n int) string {
+	s := strings.TrimSpace(string(b))
+	if len(s) > n {
+		return s[:n]
+	}
+	return s
+}
+
+func sleepCtx(ctx context.Context, d time.Duration) error {
+	t := time.NewTimer(d)
+	defer t.Stop()
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-t.C:
+		return nil
+	}
+}