From 6809081f37f5926b756bb9abda37ba9ce3bd4351 Mon Sep 17 00:00:00 2001 From: qian Date: Sat, 13 Jun 2026 09:21:11 +0800 Subject: [PATCH] =?UTF-8?q?feat(model):veo=E8=A7=86=E9=A2=91=E6=A8=A1?= =?UTF-8?q?=E5=9E=8B=E9=80=82=E9=85=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/logic/task/async_polling.go | 14 +- relay/constant/task.go | 3 + relay/handler/task_handler.go | 1 + relay/taskchannel/gemini/adaptor.go | 302 +++++++++++++++++++++++++++ 4 files changed, 319 insertions(+), 1 deletion(-) create mode 100644 relay/taskchannel/gemini/adaptor.go diff --git a/internal/logic/task/async_polling.go b/internal/logic/task/async_polling.go index 5068b9a..4e5fba1 100644 --- a/internal/logic/task/async_polling.go +++ b/internal/logic/task/async_polling.go @@ -251,7 +251,8 @@ func pollSingleTask(ctx context.Context, adaptor common.TaskAdaptor, channel *co // 查询上游状态 taskData, _ := json.Marshal(map[string]any{ - "task_id": pd.UpstreamTaskID, + "task_id": pd.UpstreamTaskID, + "use_proxy": parseChannelUseProxy(channel.Settings), }) resp, err := adaptor.FetchTask(channel.BaseURL, channel.ApiKey, taskData) @@ -478,3 +479,14 @@ func upstreamRespHeaders(resp *http.Response) map[string]string { } return headers } + +// parseChannelUseProxy 从渠道 Settings JSONB 中提取 use_proxy 字段 +func parseChannelUseProxy(settings json.RawMessage) bool { + var s struct { + UseProxy bool `json:"use_proxy"` + } + if len(settings) > 0 { + json.Unmarshal(settings, &s) + } + return s.UseProxy +} diff --git a/relay/constant/task.go b/relay/constant/task.go index 5725421..151c087 100644 --- a/relay/constant/task.go +++ b/relay/constant/task.go @@ -10,6 +10,7 @@ const ( TaskPlatformMidjourney TaskPlatform = "midjourney" TaskPlatformVolcengine TaskPlatform = "volcengine" TaskPlatformAli TaskPlatform = "ali" + TaskPlatformGemini TaskPlatform = "gemini" ) // TaskAction 异步任务动作类型 @@ -84,6 +85,8 @@ func ProviderTypeToTaskPlatform(p ProviderType) (TaskPlatform, bool) { return TaskPlatformVolcengine, true case ProviderAli: return TaskPlatformAli, true + case ProviderGemini: + return TaskPlatformGemini, true default: return "", false } diff --git a/relay/handler/task_handler.go b/relay/handler/task_handler.go index abe5075..fcfc968 100644 --- a/relay/handler/task_handler.go +++ b/relay/handler/task_handler.go @@ -15,6 +15,7 @@ import ( "github.com/qianfree/team-api/relay/constant" "github.com/qianfree/team-api/relay/taskchannel" _ "github.com/qianfree/team-api/relay/taskchannel/ali" + _ "github.com/qianfree/team-api/relay/taskchannel/gemini" _ "github.com/qianfree/team-api/relay/taskchannel/kling" "github.com/qianfree/team-api/relay/taskchannel/midjourney" _ "github.com/qianfree/team-api/relay/taskchannel/sora" diff --git a/relay/taskchannel/gemini/adaptor.go b/relay/taskchannel/gemini/adaptor.go new file mode 100644 index 0000000..85b324f --- /dev/null +++ b/relay/taskchannel/gemini/adaptor.go @@ -0,0 +1,302 @@ +package gemini + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + + "github.com/qianfree/team-api/relay/common" + "github.com/qianfree/team-api/relay/constant" + "github.com/qianfree/team-api/relay/taskchannel" +) + +const channelName = "Gemini" + +func init() { + taskchannel.Register(constant.ProviderGemini, func() common.TaskAdaptor { + return &GeminiAdaptor{} + }) +} + +type GeminiAdaptor struct { + info *common.RelayInfo +} + +func (a *GeminiAdaptor) Init(info *common.RelayInfo) { + a.info = info +} + +func (a *GeminiAdaptor) ValidateRequest(_ context.Context, _ *common.RelayInfo, body []byte) *common.TaskError { + var req struct { + Prompt string `json:"prompt"` + } + if err := json.Unmarshal(body, &req); err != nil { + return &common.TaskError{StatusCode: http.StatusBadRequest, Message: "invalid request body"} + } + if req.Prompt == "" { + return &common.TaskError{StatusCode: http.StatusBadRequest, Message: "prompt is required"} + } + return nil +} + +func (a *GeminiAdaptor) EstimateBilling(_ context.Context, _ *common.RelayInfo, body []byte) map[string]float64 { + ratios := map[string]float64{"base": 1.0} + var req openAIVideoRequest + if json.Unmarshal(body, &req) == nil { + resolution := parseResolution(req.Metadata.Resolution) + duration := req.Metadata.Duration + if resolution == "1080p" { + ratios["resolution"] = 1.5 + } else if resolution == "4k" { + ratios["resolution"] = 2.0 + } + if duration >= 8 { + ratios["duration"] = 1.3 + } + } + return ratios +} + +func (a *GeminiAdaptor) AdjustBillingOnSubmit(_ *common.RelayInfo, _ []byte) map[string]float64 { + return nil +} + +func (a *GeminiAdaptor) BuildRequestURL(info *common.RelayInfo) (string, error) { + baseURL := strings.TrimRight(info.ChannelMeta.BaseURL, "/") + model := info.ChannelMeta.UpstreamModelName + if model == "" { + model = info.OriginModelName + } + return fmt.Sprintf("%s/v1beta/models/%s:predictLongRunning", baseURL, model), nil +} + +func (a *GeminiAdaptor) BuildRequestHeader(header http.Header, info *common.RelayInfo) error { + header.Set("x-goog-api-key", info.ChannelMeta.ApiKey) + header.Set("Content-Type", "application/json") + return nil +} + +func (a *GeminiAdaptor) BuildRequestBody(_ context.Context, info *common.RelayInfo, body []byte) (io.Reader, error) { + var req openAIVideoRequest + if err := json.Unmarshal(body, &req); err != nil { + return strings.NewReader(string(body)), nil + } + + // 构建 Gemini predictLongRunning 格式 + geminiReq := map[string]any{ + "instances": []map[string]any{ + {"prompt": req.Prompt}, + }, + } + + params := map[string]any{} + if resolution := parseResolution(req.Metadata.Resolution); resolution != "" { + params["resolution"] = resolution + // 竖屏分辨率（WxH 中 H > W） + if strings.Contains(req.Metadata.Resolution, "x") { + parts := strings.Split(req.Metadata.Resolution, "x") + if len(parts) == 2 && parts[1] > parts[0] { + params["aspectRatio"] = "9:16" + } else { + params["aspectRatio"] = "16:9" + } + } + } + if req.Metadata.Duration > 0 { + params["durationSeconds"] = fmt.Sprintf("%d", req.Metadata.Duration) + } + if len(params) > 0 { + geminiReq["parameters"] = params + } + + data, _ := json.Marshal(geminiReq) + return strings.NewReader(string(data)), nil +} + +func (a *GeminiAdaptor) DoRequest(_ context.Context, info *common.RelayInfo, requestBody io.Reader) (*http.Response, error) { + url, err := a.BuildRequestURL(info) + if err != nil { + return nil, err + } + req, err := http.NewRequest("POST", url, requestBody) + if err != nil { + return nil, err + } + a.BuildRequestHeader(req.Header, info) + client := common.NewPooledClient(120, info.ChannelMeta.Settings.UseProxy) + return client.Do(req) +} + +func (a *GeminiAdaptor) DoResponse(_ context.Context, resp *http.Response, _ *common.RelayInfo) (string, []byte, *common.TaskError) { + body, err := io.ReadAll(resp.Body) + if err != nil { + return "", nil, &common.TaskError{StatusCode: 500, Message: "read response failed"} + } + + if resp.StatusCode != http.StatusOK { + return "", body, &common.TaskError{StatusCode: resp.StatusCode, Message: parseGeminiError(body)} + } + + var result struct { + Name string `json:"name"` + Done bool `json:"done"` + Error *struct { + Code int `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal(body, &result); err != nil { + return "", body, &common.TaskError{StatusCode: 500, Message: "parse response failed"} + } + if result.Error != nil { + return "", body, &common.TaskError{ + StatusCode: result.Error.Code, + Message: result.Error.Message, + } + } + if result.Name == "" { + return "", body, &common.TaskError{StatusCode: 500, Message: "upstream returned empty operation name"} + } + return result.Name, body, nil +} + +func (a *GeminiAdaptor) FetchTask(baseURL, apiKey string, taskData []byte) (*http.Response, error) { + var data struct { + TaskID string `json:"task_id"` + UseProxy bool `json:"use_proxy"` + } + if err := json.Unmarshal(taskData, &data); err != nil || data.TaskID == "" { + return nil, fmt.Errorf("gemini: invalid task data for polling") + } + + url := fmt.Sprintf("%s/v1beta/%s", strings.TrimRight(baseURL, "/"), data.TaskID) + req, err := http.NewRequest("GET", url, nil) + if err != nil { + return nil, err + } + req.Header.Set("x-goog-api-key", apiKey) + client := common.NewPooledClient(30, data.UseProxy) + return client.Do(req) +} + +func (a *GeminiAdaptor) ParseTaskResult(body []byte) (*common.TaskInfo, error) { + var resp struct { + Name string `json:"name"` + Done bool `json:"done"` + Error *struct { + Code int `json:"code"` + Message string `json:"message"` + } `json:"error"` + Response *struct { + GenerateVideoResponse *struct { + GeneratedSamples []struct { + Video struct { + URI string `json:"uri"` + VideoBytes string `json:"videoBytes"` + } `json:"video"` + } `json:"generatedSamples"` + } `json:"generateVideoResponse"` + } `json:"response"` + } + + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("gemini: parse task result: %w", err) + } + + info := &common.TaskInfo{Data: body} + + if resp.Error != nil { + info.Status = common.TaskStatusFailure + info.FailReason = resp.Error.Message + return info, nil + } + + if !resp.Done { + info.Status = common.TaskStatusInProgress + info.Progress = "50%" + return info, nil + } + + // done=true：提取视频 URL + info.Status = common.TaskStatusSuccess + info.Progress = "100%" + + if resp.Response != nil && resp.Response.GenerateVideoResponse != nil { + samples := resp.Response.GenerateVideoResponse.GeneratedSamples + if len(samples) > 0 { + info.ResultURL = samples[0].Video.URI + if len(samples) > 1 { + for i, s := range samples { + info.SubTasks = append(info.SubTasks, common.SubTask{ + Index: i, + Status: common.TaskStatusSuccess, + ResultURL: s.Video.URI, + }) + } + } + } + } + + return info, nil +} + +func (a *GeminiAdaptor) GetModelList() []string { + return []string{ + "veo-3.1-generate-001", + "veo-3.1-fast-generate-001", + "veo-3.1-generate-preview", + "veo-3.1-fast-generate-preview", + "veo-3.0-generate-001", + "veo-3.0-fast-generate-001", + "veo-2.0-generate-001", + } +} + +func (a *GeminiAdaptor) GetChannelName() string { + return channelName +} + +// parseGeminiError 从 Gemini 错误响应中提取可读消息 +func parseGeminiError(body []byte) string { + var errResp struct { + Error *struct { + Message string `json:"message"` + Status string `json:"status"` + } `json:"error"` + } + if json.Unmarshal(body, &errResp) == nil && errResp.Error != nil { + if errResp.Error.Message != "" { + return errResp.Error.Message + } + return errResp.Error.Status + } + return string(body) +} + +// openAIVideoRequest OpenAI 兼容格式的视频生成请求 +type openAIVideoRequest struct { + Model string `json:"model"` + Prompt string `json:"prompt"` + Metadata struct { + Resolution string `json:"resolution"` // "1280x720", "1920x1080" 等 + Duration int `json:"duration"` // 秒数 + } `json:"metadata"` +} + +// parseResolution 从 WxH 格式提取分辨率等级 +func parseResolution(res string) string { + switch res { + case "1920x1080", "1080x1920": + return "1080p" + case "1280x720", "720x1280": + return "720p" + case "854x480", "480x854": + return "480p" + case "3840x2160", "2160x3840": + return "4k" + } + return "" +}