feat(upstream): 自定义上游支持 Vivid 协议(/v1/generate NDJSON)
- custom 上游新增 protocol 选项(openai|vivid);vivid 走 POST /v1/generate NDJSON 流:跳过心跳空行,取末行终态, 断流携 jobId 重连不重复扣费;参考图先 POST /v1/uploads 拿预签名再 PUT 直传;错误码 4011/4002/4293/5001/4001 分类 - 添加上游弹窗加协议二选一;账号 meta 存 protocol 并可回显 - 附 vivid 链路单测(成功流/URL 透传/失败帧/错误码/两步上传) Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
@@ -6,6 +6,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"backend/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
@@ -141,6 +142,7 @@ func (h *ProviderAdminHandler) ImportCustomAccount(c *gin.Context) {
|
||||
APIKey string `json:"api_key"`
|
||||
Models string `json:"models"`
|
||||
Name string `json:"name"`
|
||||
Protocol string `json:"protocol"` // openai (default) | vivid
|
||||
Weight int `json:"weight"`
|
||||
Concurrency int `json:"concurrency"`
|
||||
ID string `json:"id"`
|
||||
@@ -157,7 +159,7 @@ func (h *ProviderAdminHandler) ImportCustomAccount(c *gin.Context) {
|
||||
if key == "" {
|
||||
key = body.APIKey
|
||||
}
|
||||
item, err := h.tokens.ImportCustomAccount(c.Request.Context(), baseURL, key, body.Models, body.Name, body.Weight, body.Concurrency, body.ID)
|
||||
item, err := h.tokens.ImportCustomAccount(c.Request.Context(), baseURL, key, body.Models, body.Name, body.Protocol, body.Weight, body.Concurrency, body.ID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"detail": err.Error()})
|
||||
return
|
||||
@@ -395,13 +397,21 @@ func (h *ProviderAdminHandler) AccountsList(c *gin.Context) {
|
||||
isFree := plan == "" || plan == "free"
|
||||
switch planFilter {
|
||||
case "vip":
|
||||
if isFree { continue }
|
||||
if isFree {
|
||||
continue
|
||||
}
|
||||
case "free":
|
||||
if !isFree { continue }
|
||||
if !isFree {
|
||||
continue
|
||||
}
|
||||
case "sub":
|
||||
if !isSub { continue }
|
||||
if !isSub {
|
||||
continue
|
||||
}
|
||||
case "master":
|
||||
if isFree || isSub { continue }
|
||||
if isFree || isSub {
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
if q != "" {
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
package custom
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
@@ -120,6 +121,253 @@ func (c *Client) GenerateImage(ctx context.Context, baseURL, apiKey, model, prom
|
||||
return imageFromResponse(ctx, body, downloadResult)
|
||||
}
|
||||
|
||||
// ── Vivid 协议 ──────────────────────────────────────────────────────────────
|
||||
// Vivid 上游用自家 API:POST /v1/generate 返回 application/x-ndjson 流 —— 首行
|
||||
// {"jobId":...,"status":"running"},等待期间每 25s 一个空行心跳,末行为终态
|
||||
// (succeeded 带 result 预签名 URL / failed 带 error)。参考素材走两步:
|
||||
// POST /v1/uploads {kind,mime} → PUT 直传预签名 URL → 拿 key 放进 refs。
|
||||
// 断流可携 jobId 重发重连。错误体为 {"error":{"code":"4011|4001|4002|4293|5001"}}。
|
||||
|
||||
// GenerateVividImage submits an image job to a Vivid-protocol upstream.
|
||||
// quality is the tier string the upstream expects ("1K"/"2K"/"4K"), ratio the
|
||||
// "W:H" aspect. refs are pre-uploaded via /v1/uploads into key strings.
|
||||
func (c *Client) GenerateVividImage(ctx context.Context, baseURL, apiKey, model, prompt, ratio, quality string, refs [][]byte, downloadResult bool) ([]byte, string, error) {
|
||||
payload, err := c.vividPayload(ctx, baseURL, apiKey, model, prompt, ratio, quality, 0, refs)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return c.vividGenerate(ctx, baseURL, apiKey, payload, downloadResult)
|
||||
}
|
||||
|
||||
// GenerateVividVideo submits a video job (duration required by ranged models)
|
||||
// to a Vivid-protocol upstream.
|
||||
func (c *Client) GenerateVividVideo(ctx context.Context, baseURL, apiKey, model, prompt, ratio, quality string, durationSeconds int, refs [][]byte, downloadResult bool) ([]byte, string, error) {
|
||||
payload, err := c.vividPayload(ctx, baseURL, apiKey, model, prompt, ratio, quality, durationSeconds, refs)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return c.vividGenerate(ctx, baseURL, apiKey, payload, downloadResult)
|
||||
}
|
||||
|
||||
func (c *Client) vividPayload(ctx context.Context, baseURL, apiKey, model, prompt, ratio, quality string, durationSeconds int, refs [][]byte) (map[string]any, error) {
|
||||
baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
|
||||
if baseURL == "" || apiKey == "" {
|
||||
return nil, ErrAuth
|
||||
}
|
||||
payload := map[string]any{"model": model, "prompt": prompt}
|
||||
if quality != "" {
|
||||
payload["quality"] = quality
|
||||
}
|
||||
if ratio != "" {
|
||||
payload["ratio"] = ratio
|
||||
}
|
||||
if durationSeconds > 0 {
|
||||
payload["duration"] = durationSeconds
|
||||
}
|
||||
if len(refs) > 0 {
|
||||
keys := make([]string, 0, len(refs))
|
||||
for _, ref := range refs {
|
||||
key, err := c.vividUploadRef(ctx, baseURL, apiKey, ref)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
keys = append(keys, key)
|
||||
}
|
||||
payload["refs"] = keys
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
// vividUploadRef runs the two-step upload: POST /v1/uploads {kind:"image",mime}
|
||||
// → PUT the bytes to the presigned url → returns the key for refs[].
|
||||
func (c *Client) vividUploadRef(ctx context.Context, baseURL, apiKey string, data []byte) (string, error) {
|
||||
mime := http.DetectContentType(data)
|
||||
if mime == "" || mime == "application/octet-stream" {
|
||||
mime = "image/png"
|
||||
}
|
||||
raw, _ := json.Marshal(map[string]any{"kind": "image", "mime": mime})
|
||||
up, err := c.doJSON(ctx, http.MethodPost, baseURL+"/v1/uploads", apiKey, raw)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
putURL := strings.TrimSpace(stringValue(up["url"]))
|
||||
key := strings.TrimSpace(stringValue(up["key"]))
|
||||
if putURL == "" || key == "" {
|
||||
return "", fmt.Errorf("%w: upload session missing url/key", ErrTemporaryUpstream)
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPut, putURL, bytes.NewReader(data))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header.Set("Content-Type", mime) // must match the signed mime exactly
|
||||
resp, err := httpClient().Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err))
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return "", fmt.Errorf("%w: ref upload %d %s", ErrTemporaryUpstream, resp.StatusCode, clip(body, 120))
|
||||
}
|
||||
return key, nil
|
||||
}
|
||||
|
||||
// vividGenerate POSTs /v1/generate and reads the NDJSON stream to the terminal
|
||||
// frame. If the stream drops before terminal, it reconnects once with the
|
||||
// jobId (same-request resume — no new task, no double charge).
|
||||
func (c *Client) vividGenerate(ctx context.Context, baseURL, apiKey string, payload map[string]any, downloadResult bool) ([]byte, string, error) {
|
||||
var last map[string]any
|
||||
for attempt := 0; attempt < 3; attempt++ {
|
||||
if attempt > 0 {
|
||||
if sleepCtx(ctx, 2*time.Second) != nil {
|
||||
return nil, "", ctx.Err()
|
||||
}
|
||||
}
|
||||
term, err := c.vividStreamOnce(ctx, baseURL, apiKey, payload)
|
||||
if err == nil && term != nil {
|
||||
last = term
|
||||
break
|
||||
}
|
||||
if err != nil && !errors.Is(err, ErrTemporaryUpstream) {
|
||||
return nil, "", err
|
||||
}
|
||||
// Temporary error mid-stream: if we hold a jobId, reconnect mode retries
|
||||
// without creating a new task.
|
||||
if job := strings.TrimSpace(stringValue(payload["jobId"])); job != "" {
|
||||
payload = map[string]any{"jobId": job}
|
||||
continue
|
||||
}
|
||||
if attempt == 2 {
|
||||
return nil, "", err
|
||||
}
|
||||
}
|
||||
if last == nil {
|
||||
return nil, "", fmt.Errorf("%w: no terminal frame", ErrTemporaryUpstream)
|
||||
}
|
||||
status := strings.ToLower(strings.TrimSpace(stringValue(last["status"])))
|
||||
switch status {
|
||||
case "succeeded", "completed", "success":
|
||||
resultURL := strings.TrimSpace(stringValue(last["result"]))
|
||||
if resultURL == "" {
|
||||
return nil, "", fmt.Errorf("%w: succeeded frame missing result", ErrTemporaryUpstream)
|
||||
}
|
||||
if !downloadResult {
|
||||
return nil, resultURL, nil
|
||||
}
|
||||
data, err := c.vividDownload(ctx, resultURL)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return data, resultURL, nil
|
||||
default:
|
||||
reason := strings.TrimSpace(stringValue(last["error"]))
|
||||
if reason == "" {
|
||||
reason = "upstream job " + status
|
||||
}
|
||||
if isCreditError(reason) || strings.Contains(reason, "积分不足") || strings.Contains(reason, "余额") {
|
||||
return nil, "", fmt.Errorf("%w: %s", ErrQuotaExhausted, clip([]byte(reason), 160))
|
||||
}
|
||||
return nil, "", errors.New("custom: " + clip([]byte(reason), 160))
|
||||
}
|
||||
}
|
||||
|
||||
// vividStreamOnce performs one /v1/generate request and scans the NDJSON body
|
||||
// to the last frame. Returns the terminal frame, or a temporary error on
|
||||
// transport failure (caller's reconnect logic decides what to do with jobId).
|
||||
func (c *Client) vividStreamOnce(ctx context.Context, baseURL, apiKey string, payload map[string]any) (map[string]any, error) {
|
||||
raw, _ := json.Marshal(payload)
|
||||
req, err := http.NewRequest(http.MethodPost, baseURL+"/v1/generate", bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(ctx)
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("X-Wait-Seconds", "7200")
|
||||
resp, err := httpClient().Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err))
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return nil, vividError(resp.StatusCode, body)
|
||||
}
|
||||
var last map[string]any
|
||||
sc := bufio.NewScanner(resp.Body)
|
||||
sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
|
||||
for sc.Scan() {
|
||||
line := bytes.TrimSpace(sc.Bytes())
|
||||
if len(line) == 0 {
|
||||
continue // 25s heartbeat empty line
|
||||
}
|
||||
var frame map[string]any
|
||||
if json.Unmarshal(line, &frame) != nil {
|
||||
continue
|
||||
}
|
||||
last = frame
|
||||
// Remember the jobId for potential reconnect.
|
||||
if job := strings.TrimSpace(stringValue(frame["jobId"])); job != "" {
|
||||
payload["jobId"] = job
|
||||
}
|
||||
}
|
||||
if err := sc.Err(); err != nil {
|
||||
return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err))
|
||||
}
|
||||
return last, nil
|
||||
}
|
||||
|
||||
// vividError maps the {"error":{"code":"NNNN"}} envelope to our error classes.
|
||||
func vividError(status int, body []byte) error {
|
||||
var e struct {
|
||||
Error struct {
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
code := ""
|
||||
msg := clip(body, 160)
|
||||
if json.Unmarshal(body, &e) == nil {
|
||||
code = strings.TrimSpace(e.Error.Code)
|
||||
if m := strings.TrimSpace(e.Error.Message); m != "" {
|
||||
msg = clip([]byte(m), 160)
|
||||
}
|
||||
}
|
||||
switch code {
|
||||
case "4011":
|
||||
return fmt.Errorf("%w: %s", ErrAuth, msg)
|
||||
case "4002":
|
||||
return fmt.Errorf("%w: %s", ErrQuotaExhausted, msg)
|
||||
case "4293", "5001":
|
||||
return fmt.Errorf("%w: %s", ErrTemporaryUpstream, msg)
|
||||
case "4001":
|
||||
return errors.New("custom: " + msg) // 参数/业务拒绝 → 永久错,不换号
|
||||
}
|
||||
return mapStatus(status, body)
|
||||
}
|
||||
|
||||
// vividDownload fetches the presigned result URL — no auth header (it's signed).
|
||||
func (c *Client) vividDownload(ctx context.Context, url string) ([]byte, error) {
|
||||
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
resp, err := httpClient().Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%w: %s", ErrTemporaryUpstream, sanitizeErr(err))
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("%w: result 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 result", ErrTemporaryUpstream)
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// GenerateVideo drives the upstream Sora-style async video API:
|
||||
// POST /v1/videos → poll GET /v1/videos/{id} → GET /v1/videos/{id}/content.
|
||||
// Reference frames (image-to-video / first-last frames) are sent as multipart
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
package custom
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// vividServer serves /v1/uploads, /v1/generate (ndjson) and a fake presigned
|
||||
// result URL, all under one httptest server.
|
||||
func vividServer(t *testing.T, frames []string, httpStatus int, errCode string) *httptest.Server {
|
||||
t.Helper()
|
||||
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/v1/uploads":
|
||||
var req struct {
|
||||
Kind string `json:"kind"`
|
||||
Mime string `json:"mime"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||
if req.Kind != "image" || req.Mime == "" {
|
||||
t.Errorf("bad upload request: %+v", req)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprintf(w, `{"uuid":"u1","key":"u/refs/image/u1.png","url":"%s","expiresAt":9999999999}`,
|
||||
// presigned URL points back at this server's /put endpoint
|
||||
"http://"+r.Host+"/put")
|
||||
case "/put":
|
||||
if r.Method != http.MethodPut {
|
||||
t.Errorf("expected PUT, got %s", r.Method)
|
||||
}
|
||||
if ct := r.Header.Get("Content-Type"); !strings.HasPrefix(ct, "image/") {
|
||||
t.Errorf("bad content-type %q", ct)
|
||||
}
|
||||
io.Copy(io.Discard, r.Body)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
case "/v1/generate":
|
||||
if httpStatus != 0 {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(httpStatus)
|
||||
fmt.Fprintf(w, `{"error":{"code":%q,"message":"boom"}}`, errCode)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/x-ndjson")
|
||||
for _, f := range frames {
|
||||
fmt.Fprintln(w, f)
|
||||
fmt.Fprintln(w) // heartbeat line
|
||||
}
|
||||
case "/result.png":
|
||||
w.Write([]byte("PNGDATA"))
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
func TestVividGenerateImageSuccess(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/v1/generate":
|
||||
w.Header().Set("Content-Type", "application/x-ndjson")
|
||||
fmt.Fprintln(w, `{"jobId":"j1","status":"running"}`)
|
||||
fmt.Fprintln(w) // heartbeat
|
||||
fmt.Fprintf(w, `{"jobId":"j1","status":"succeeded","result":"http://%s/result.png","creditsCharged":10}`+"\n", r.Host)
|
||||
case "/result.png":
|
||||
w.Write([]byte("PNGDATA"))
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
data, url, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", nil, true)
|
||||
if err != nil {
|
||||
t.Fatalf("generate: %v", err)
|
||||
}
|
||||
if string(data) != "PNGDATA" {
|
||||
t.Fatalf("data=%q", data)
|
||||
}
|
||||
if url == "" {
|
||||
t.Fatal("empty url")
|
||||
}
|
||||
}
|
||||
|
||||
func TestVividGenerateImageURLPassthrough(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/x-ndjson")
|
||||
fmt.Fprintln(w, `{"jobId":"j1","status":"running"}`)
|
||||
fmt.Fprintln(w, `{"jobId":"j1","status":"succeeded","result":"https://fs.example.com/x.png"}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
data, url, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", nil, false)
|
||||
if err != nil {
|
||||
t.Fatalf("generate: %v", err)
|
||||
}
|
||||
if data != nil {
|
||||
t.Fatalf("expected no bytes, got %d", len(data))
|
||||
}
|
||||
if url != "https://fs.example.com/x.png" {
|
||||
t.Fatalf("url=%q", url)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVividGenerateFailedFrame(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/x-ndjson")
|
||||
fmt.Fprintln(w, `{"jobId":"j1","status":"running"}`)
|
||||
fmt.Fprintln(w, `{"jobId":"j1","status":"failed","error":"积分不足","refunded":true}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
_, _, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", nil, false)
|
||||
if !errors.Is(err, ErrQuotaExhausted) {
|
||||
t.Fatalf("want quota, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVividErrorCodes(t *testing.T) {
|
||||
cases := []struct {
|
||||
status int
|
||||
code string
|
||||
want error
|
||||
}{
|
||||
{401, "4011", ErrAuth},
|
||||
{402, "4002", ErrQuotaExhausted},
|
||||
{429, "4293", ErrTemporaryUpstream},
|
||||
{500, "5001", ErrTemporaryUpstream},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
srv := vividServer(t, nil, tc.status, tc.code)
|
||||
_, _, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", nil, false)
|
||||
srv.Close()
|
||||
if !errors.Is(err, tc.want) {
|
||||
t.Fatalf("code %s: want %v, got %v", tc.code, tc.want, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestVividUploadRefFlow(t *testing.T) {
|
||||
var putOK bool
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/v1/uploads":
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprintf(w, `{"key":"k123","url":"http://%s/up"}`, r.Host)
|
||||
case "/up":
|
||||
if r.Method == http.MethodPut {
|
||||
putOK = true
|
||||
}
|
||||
w.WriteHeader(http.StatusOK)
|
||||
case "/v1/generate":
|
||||
var req map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||
refs, _ := req["refs"].([]any)
|
||||
if len(refs) != 1 || refs[0] != "k123" {
|
||||
t.Errorf("refs=%v", req["refs"])
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/x-ndjson")
|
||||
fmt.Fprintln(w, `{"jobId":"j1","status":"succeeded","result":"https://x/1.png"}`)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
png := []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A} // fake png magic
|
||||
_, url, err := NewClient().GenerateVividImage(context.Background(), srv.URL, "vk_test", "m1", "hi", "1:1", "1K", [][]byte{png}, false)
|
||||
if err != nil {
|
||||
t.Fatalf("generate: %v", err)
|
||||
}
|
||||
if !putOK {
|
||||
t.Fatal("PUT never happened")
|
||||
}
|
||||
if url == "" {
|
||||
t.Fatal("no url")
|
||||
}
|
||||
}
|
||||
@@ -2,24 +2,24 @@ package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"backend/internal/model"
|
||||
"backend/internal/provider/adobe"
|
||||
"backend/internal/provider/chatgpt"
|
||||
"backend/internal/provider/creativefabrica"
|
||||
"backend/internal/provider/grok"
|
||||
"backend/internal/provider/imagine"
|
||||
"backend/internal/provider/krea"
|
||||
"backend/internal/provider/leonardo"
|
||||
"backend/internal/provider/grok"
|
||||
"backend/internal/provider/runway"
|
||||
"backend/internal/repo"
|
||||
|
||||
@@ -28,14 +28,14 @@ import (
|
||||
)
|
||||
|
||||
var validTokenPools = map[string]string{
|
||||
"chatgpt": "openai",
|
||||
"adobe": "adobe",
|
||||
"runway": "runway",
|
||||
"leonardo": "leonardo",
|
||||
"krea": "krea",
|
||||
"imagine": "imagine",
|
||||
"grok": "grok",
|
||||
"custom": "custom",
|
||||
"chatgpt": "openai",
|
||||
"adobe": "adobe",
|
||||
"runway": "runway",
|
||||
"leonardo": "leonardo",
|
||||
"krea": "krea",
|
||||
"imagine": "imagine",
|
||||
"grok": "grok",
|
||||
"custom": "custom",
|
||||
"creativefabrica": "creativefabrica",
|
||||
}
|
||||
|
||||
@@ -1036,7 +1036,6 @@ func (s *TokenService) ImportGrokToken(ctx context.Context, ssoToken, tokenID st
|
||||
return item, nil
|
||||
}
|
||||
|
||||
|
||||
func (s *TokenService) checkPendingGrok(tokenID, ssoToken string) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
@@ -1217,9 +1216,10 @@ func (s *TokenService) checkPendingCreativeFabrica(tokenID, cookie string) {
|
||||
// csv list of model ids it serves (empty = all), plus optional weight and
|
||||
// per-account concurrency. No probe — the account goes active immediately and is
|
||||
// matched to custom models by id at generation time. Calls go direct (no proxy).
|
||||
func (s *TokenService) ImportCustomAccount(ctx context.Context, baseURL, apiKey, models, name string, weight, concurrency int, tokenID string) (*model.TokenAccount, error) {
|
||||
func (s *TokenService) ImportCustomAccount(ctx context.Context, baseURL, apiKey, models, name, protocol string, weight, concurrency int, tokenID string) (*model.TokenAccount, error) {
|
||||
baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
|
||||
apiKey = strings.TrimSpace(apiKey)
|
||||
protocol = strings.ToLower(strings.TrimSpace(protocol))
|
||||
// Edit mode: tokenID points at an existing custom account. base_url required;
|
||||
// a blank key keeps the stored one.
|
||||
if strings.TrimSpace(tokenID) != "" {
|
||||
@@ -1230,7 +1230,14 @@ func (s *TokenService) ImportCustomAccount(ctx context.Context, baseURL, apiKey,
|
||||
if baseURL == "" {
|
||||
return nil, errors.New("base_url required")
|
||||
}
|
||||
meta := datatypes.JSONMap{"base_url": baseURL}
|
||||
proto := protocol
|
||||
if proto == "" {
|
||||
proto = strings.ToLower(strings.TrimSpace(stringValue(existing.Meta["protocol"])))
|
||||
}
|
||||
if proto != "vivid" {
|
||||
proto = "openai"
|
||||
}
|
||||
meta := datatypes.JSONMap{"base_url": baseURL, "protocol": proto}
|
||||
if m := strings.TrimSpace(models); m != "" {
|
||||
meta["models"] = m
|
||||
}
|
||||
@@ -1248,7 +1255,10 @@ func (s *TokenService) ImportCustomAccount(ctx context.Context, baseURL, apiKey,
|
||||
if baseURL == "" || apiKey == "" {
|
||||
return nil, errors.New("base_url and key required")
|
||||
}
|
||||
meta := datatypes.JSONMap{"base_url": baseURL}
|
||||
if protocol != "vivid" {
|
||||
protocol = "openai"
|
||||
}
|
||||
meta := datatypes.JSONMap{"base_url": baseURL, "protocol": protocol}
|
||||
if m := strings.TrimSpace(models); m != "" {
|
||||
meta["models"] = m
|
||||
}
|
||||
@@ -1952,6 +1962,7 @@ func accountRow(item model.TokenAccount, inFlight int64) map[string]any {
|
||||
"weight": item.Weight,
|
||||
"concurrency": item.Concurrency,
|
||||
"base_url": emptyToNil(strings.TrimSpace(stringValue(item.Meta["base_url"]))),
|
||||
"protocol": emptyToNil(strings.TrimSpace(stringValue(item.Meta["protocol"]))),
|
||||
"models": strings.TrimSpace(stringValue(item.Meta["models"])),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -23,8 +23,8 @@ import (
|
||||
"backend/internal/model"
|
||||
"backend/internal/provider/adobe"
|
||||
"backend/internal/provider/chatgpt"
|
||||
"backend/internal/provider/custom"
|
||||
"backend/internal/provider/creativefabrica"
|
||||
"backend/internal/provider/custom"
|
||||
"backend/internal/provider/grok"
|
||||
"backend/internal/provider/imagine"
|
||||
"backend/internal/provider/krea"
|
||||
@@ -32,6 +32,7 @@ import (
|
||||
"backend/internal/provider/runway"
|
||||
"backend/internal/repo"
|
||||
"backend/internal/storage"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
@@ -2228,7 +2229,14 @@ func (s *V1Service) generateCustomImage(ctx context.Context, eventID string, mod
|
||||
_ = s.events.SetAccount(ctx, eventID, token.ID, token.AccountEmail)
|
||||
_ = s.tokens.TouchLastUsed(ctx, token.ID)
|
||||
baseURL := stringValue(token.Meta["base_url"])
|
||||
d, u, genErr := s.custom.GenerateImage(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, size, quality, refs, !urlOnly)
|
||||
var d []byte
|
||||
var u string
|
||||
var genErr error
|
||||
if strings.EqualFold(stringValue(token.Meta["protocol"]), "vivid") {
|
||||
d, u, genErr = s.custom.GenerateVividImage(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, aspectRatio, vividQuality(resolution), refs, !urlOnly)
|
||||
} else {
|
||||
d, u, genErr = s.custom.GenerateImage(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, size, quality, refs, !urlOnly)
|
||||
}
|
||||
if genErr == nil {
|
||||
_, _ = s.tokens.Update(ctx, "custom", token.ID, map[string]any{
|
||||
"last_used_at": time.Now(), "success_total": gorm.Expr("success_total + 1"), "fails": 0,
|
||||
@@ -2303,7 +2311,14 @@ func (s *V1Service) generateCustomVideo(ctx context.Context, eventID string, mod
|
||||
_ = s.events.SetAccount(ctx, eventID, token.ID, token.AccountEmail)
|
||||
_ = s.tokens.TouchLastUsed(ctx, token.ID)
|
||||
baseURL := stringValue(token.Meta["base_url"])
|
||||
d, url, genErr := s.custom.GenerateVideo(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, size, durationSeconds, frames, downloadResult)
|
||||
var d []byte
|
||||
var url string
|
||||
var genErr error
|
||||
if strings.EqualFold(stringValue(token.Meta["protocol"]), "vivid") {
|
||||
d, url, genErr = s.custom.GenerateVividVideo(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, aspectRatio, vividQuality(resolution), durationSeconds, frames, downloadResult)
|
||||
} else {
|
||||
d, url, genErr = s.custom.GenerateVideo(ctx, baseURL, token.Value, modelItem.ID, in.Prompt, size, durationSeconds, frames, downloadResult)
|
||||
}
|
||||
if genErr == nil {
|
||||
_, _ = s.tokens.Update(ctx, "custom", token.ID, map[string]any{
|
||||
"last_used_at": time.Now(), "success_total": gorm.Expr("success_total + 1"), "fails": 0,
|
||||
@@ -2401,6 +2416,16 @@ func upstreamVideoSize(aspectRatio, resolution string) string {
|
||||
return fmt.Sprintf("%dx%d", short, short*h/w)
|
||||
}
|
||||
|
||||
// vividQuality maps a resolution tier to the quality string a Vivid-protocol
|
||||
// upstream expects ("1K"/"2K"/"4K"); empty → "1K" (required field).
|
||||
func vividQuality(resolution string) string {
|
||||
switch strings.ToUpper(strings.TrimSpace(resolution)) {
|
||||
case "2K", "4K":
|
||||
return strings.ToUpper(strings.TrimSpace(resolution))
|
||||
}
|
||||
return "1K"
|
||||
}
|
||||
|
||||
// upstreamQuality maps a resolution tier to the OpenAI quality enum.
|
||||
func upstreamQuality(resolution string) string {
|
||||
switch strings.ToUpper(strings.TrimSpace(resolution)) {
|
||||
@@ -2975,8 +3000,8 @@ const leonardoPrivateSuffix = "-不卡人脸"
|
||||
type leonardoVideoSpec struct {
|
||||
upstream string
|
||||
// long/short 是长边/短边像素,按比例组合成 16:9 或 9:16。
|
||||
long int
|
||||
short int
|
||||
long int
|
||||
short int
|
||||
maxImages int
|
||||
maxAudios int
|
||||
maxVideos int
|
||||
|
||||
@@ -9,6 +9,7 @@ const emit = defineEmits(['close', 'imported'])
|
||||
const isEdit = !!props.account
|
||||
const name = ref(props.account?.email || '')
|
||||
const baseUrl = ref(props.account?.base_url || '')
|
||||
const protocol = ref(props.account?.protocol === 'vivid' ? 'vivid' : 'openai')
|
||||
const key = ref('') // edit: blank = keep existing key
|
||||
const allModels = ref([]) // existing models to pick from
|
||||
const selected = ref(props.account?.models ? String(props.account.models).split(',').map((x) => x.trim()).filter(Boolean) : [])
|
||||
@@ -43,6 +44,7 @@ async function submit() {
|
||||
base_url: baseUrl.value.trim(),
|
||||
key: key.value.trim(), // blank in edit = keep existing
|
||||
models: selected.value.join(','),
|
||||
protocol: protocol.value,
|
||||
weight: Number(weight.value) || 0,
|
||||
concurrency: Number(concurrency.value) || 1,
|
||||
}))
|
||||
@@ -65,7 +67,7 @@ async function submit() {
|
||||
@click.self="emit('close')">
|
||||
<div class="card !shadow-xl mt-14 mb-14 w-full max-w-lg">
|
||||
<div class="px-5 py-4 border-b border-slate-100 flex items-center justify-between">
|
||||
<h2 class="text-sm font-semibold">{{ isEdit ? '编辑上游' : '添加上游(自定义 OpenAI 兼容)' }}</h2>
|
||||
<h2 class="text-sm font-semibold">{{ isEdit ? '编辑上游' : '添加上游' }}</h2>
|
||||
<button @click="emit('close')" class="text-slate-400 hover:text-slate-700 transition-colors">
|
||||
<Icon name="close" class="w-5 h-5" />
|
||||
</button>
|
||||
@@ -79,6 +81,18 @@ async function submit() {
|
||||
<label class="text-xs text-slate-500">备注名</label>
|
||||
<input v-model="name" class="field" placeholder="例如:我的中转 / xx-api" />
|
||||
</div>
|
||||
<div>
|
||||
<label class="text-xs text-slate-500">协议</label>
|
||||
<div class="flex gap-2">
|
||||
<button v-for="p in [{id:'openai',label:'OpenAI',desc:'/v1/images/* 标准格式'},{id:'vivid',label:'Vivid',desc:'/v1/generate NDJSON 流'}]" :key="p.id" type="button"
|
||||
@click="protocol = p.id"
|
||||
class="flex-1 rounded-lg ring-1 px-3 py-2 text-left transition-colors"
|
||||
:class="protocol === p.id ? 'ring-sky-400 bg-sky-500/10' : 'ring-slate-200 bg-white hover:ring-slate-300'">
|
||||
<div class="text-xs font-medium" :class="protocol === p.id ? 'text-sky-700' : 'text-slate-700'">{{ p.label }}</div>
|
||||
<div class="text-[11px] text-slate-400 font-mono mt-0.5">{{ p.desc }}</div>
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
<div>
|
||||
<label class="text-xs text-slate-500">v1 URL <span class="text-rose-500">*</span></label>
|
||||
<input v-model="baseUrl" class="field font-mono text-xs" placeholder="https://api.example.com(无需 /v1 结尾)" />
|
||||
|
||||
Reference in New Issue
Block a user