Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 24 additions & 2 deletions controller/video_proxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -92,8 +92,30 @@ func VideoProxy(c *gin.Context) {
}

if channel.Type == constant.ChannelTypeGemini {
videoURL = fmt.Sprintf("%s&key=%s", c.Query("url"), channel.Key)
req.Header.Set("x-goog-api-key", channel.Key)
apiKey := task.PrivateData.Key
if apiKey == "" {
logger.LogError(c.Request.Context(), fmt.Sprintf("Missing stored API key for Gemini task %s", taskID))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "API key not stored for task",
"type": "server_error",
},
})
return
}

videoURL, err = getGeminiVideoURL(channel, task, apiKey)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to resolve Gemini video URL for task %s: %s", taskID, err.Error()))
c.JSON(http.StatusBadGateway, gin.H{
"error": gin.H{
"message": "Failed to resolve Gemini video URL",
"type": "server_error",
},
})
return
}
req.Header.Set("x-goog-api-key", apiKey)
} else {
// Default (Sora, etc.): Use original logic
videoURL = fmt.Sprintf("%s/v1/videos/%s/content", baseURL, task.TaskID)
Expand Down
158 changes: 158 additions & 0 deletions controller/video_proxy_gemini.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,158 @@
package controller

import (
"encoding/json"
"fmt"
"io"
"strconv"
"strings"

"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay"
)

func getGeminiVideoURL(channel *model.Channel, task *model.Task, apiKey string) (string, error) {
if channel == nil || task == nil {
return "", fmt.Errorf("invalid channel or task")
}

if url := extractGeminiVideoURLFromTaskData(task); url != "" {
return ensureAPIKey(url, apiKey), nil
}

baseURL := constant.ChannelBaseURLs[channel.Type]
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}

adaptor := relay.GetTaskAdaptor(constant.TaskPlatform(strconv.Itoa(channel.Type)))
if adaptor == nil {
return "", fmt.Errorf("gemini task adaptor not found")
}

if apiKey == "" {
return "", fmt.Errorf("api key not available for task")
}

resp, err := adaptor.FetchTask(baseURL, apiKey, map[string]any{
"task_id": task.TaskID,
"action": task.Action,
})
if err != nil {
return "", fmt.Errorf("fetch task failed: %w", err)
}
defer resp.Body.Close()

body, err := io.ReadAll(resp.Body)
if err != nil {
return "", fmt.Errorf("read task response failed: %w", err)
}

taskInfo, parseErr := adaptor.ParseTaskResult(body)
if parseErr == nil && taskInfo != nil && taskInfo.RemoteUrl != "" {
return ensureAPIKey(taskInfo.RemoteUrl, apiKey), nil
}

if url := extractGeminiVideoURLFromPayload(body); url != "" {
return ensureAPIKey(url, apiKey), nil
}

if parseErr != nil {
return "", fmt.Errorf("parse task result failed: %w", parseErr)
}

return "", fmt.Errorf("gemini video url not found")
}

func extractGeminiVideoURLFromTaskData(task *model.Task) string {
if task == nil || len(task.Data) == 0 {
return ""
}
var payload map[string]any
if err := json.Unmarshal(task.Data, &payload); err != nil {
return ""
}
return extractGeminiVideoURLFromMap(payload)
}

func extractGeminiVideoURLFromPayload(body []byte) string {
var payload map[string]any
if err := json.Unmarshal(body, &payload); err != nil {
return ""
}
return extractGeminiVideoURLFromMap(payload)
}

func extractGeminiVideoURLFromMap(payload map[string]any) string {
if payload == nil {
return ""
}
if uri, ok := payload["uri"].(string); ok && uri != "" {
return uri
}
if resp, ok := payload["response"].(map[string]any); ok {
if uri := extractGeminiVideoURLFromResponse(resp); uri != "" {
return uri
}
}
return ""
}

func extractGeminiVideoURLFromResponse(resp map[string]any) string {
if resp == nil {
return ""
}
if gvr, ok := resp["generateVideoResponse"].(map[string]any); ok {
if uri := extractGeminiVideoURLFromGeneratedSamples(gvr); uri != "" {
return uri
}
}
if videos, ok := resp["videos"].([]any); ok {
for _, video := range videos {
if vm, ok := video.(map[string]any); ok {
if uri, ok := vm["uri"].(string); ok && uri != "" {
return uri
}
}
}
}
if uri, ok := resp["video"].(string); ok && uri != "" {
return uri
}
if uri, ok := resp["uri"].(string); ok && uri != "" {
return uri
}
return ""
}

func extractGeminiVideoURLFromGeneratedSamples(gvr map[string]any) string {
if gvr == nil {
return ""
}
if samples, ok := gvr["generatedSamples"].([]any); ok {
for _, sample := range samples {
if sm, ok := sample.(map[string]any); ok {
if video, ok := sm["video"].(map[string]any); ok {
if uri, ok := video["uri"].(string); ok && uri != "" {
return uri
}
}
}
}
}
return ""
}

func ensureAPIKey(uri, key string) string {
if key == "" || uri == "" {
return uri
}
if strings.Contains(uri, "key=") {
return uri
}
if strings.Contains(uri, "?") {
return fmt.Sprintf("%s&key=%s", uri, key)
}
return fmt.Sprintf("%s?key=%s", uri, key)
}
53 changes: 44 additions & 9 deletions model/task.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,8 +57,9 @@ type Task struct {
FinishTime int64 `json:"finish_time" gorm:"index"`
Progress string `json:"progress" gorm:"type:varchar(20);index"`
Properties Properties `json:"properties" gorm:"type:json"`

Data json.RawMessage `json:"data" gorm:"type:json"`
// 禁止返回给用户,内部可能包含key等隐私信息
PrivateData TaskPrivateData `json:"-" gorm:"column:private_data;type:json"`
Data json.RawMessage `json:"data" gorm:"type:json"`
}

func (t *Task) SetData(data any) {
Expand All @@ -77,13 +78,39 @@ type Properties struct {

func (m *Properties) Scan(val interface{}) error {
bytesValue, _ := val.([]byte)
if len(bytesValue) == 0 {
m.Input = ""
return nil
}
return json.Unmarshal(bytesValue, m)
}

func (m Properties) Value() (driver.Value, error) {
if m.Input == "" {
return nil, nil
}
return json.Marshal(m)
}

type TaskPrivateData struct {
Key string `json:"key,omitempty"`
}

func (p *TaskPrivateData) Scan(val interface{}) error {
bytesValue, _ := val.([]byte)
if len(bytesValue) == 0 {
return nil
}
return json.Unmarshal(bytesValue, p)
}

func (p TaskPrivateData) Value() (driver.Value, error) {
if (p == TaskPrivateData{}) {
return nil, nil
}
return json.Marshal(p)
}

// SyncTaskQueryParams 用于包含所有搜索条件的结构体,可以根据需求添加更多字段
type SyncTaskQueryParams struct {
Platform constant.TaskPlatform
Expand All @@ -98,14 +125,22 @@ type SyncTaskQueryParams struct {
}

func InitTask(platform constant.TaskPlatform, relayInfo *commonRelay.RelayInfo) *Task {
properties := Properties{}
privateData := TaskPrivateData{}
if relayInfo != nil && relayInfo.ChannelMeta != nil && relayInfo.ChannelMeta.ChannelType == constant.ChannelTypeGemini {
privateData.Key = relayInfo.ChannelMeta.ApiKey
}

t := &Task{
UserId: relayInfo.UserId,
Group: relayInfo.UsingGroup,
SubmitTime: time.Now().Unix(),
Status: TaskStatusNotStart,
Progress: "0%",
ChannelId: relayInfo.ChannelId,
Platform: platform,
UserId: relayInfo.UserId,
Group: relayInfo.UsingGroup,
SubmitTime: time.Now().Unix(),
Status: TaskStatusNotStart,
Progress: "0%",
ChannelId: relayInfo.ChannelId,
Platform: platform,
Properties: properties,
PrivateData: privateData,
}
return t
}
Expand Down
52 changes: 50 additions & 2 deletions relay/channel/task/gemini/adaptor.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,11 @@ import (
"fmt"
"io"
"net/http"
"regexp"
"strings"
"time"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
Expand Down Expand Up @@ -248,17 +250,45 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
ti.Status = model.TaskStatusSuccess
ti.Progress = "100%"

taskID := encodeLocalTaskID(op.Name)
ti.TaskID = taskID
ti.Url = fmt.Sprintf("%s/v1/videos/%s/content", system_setting.ServerAddress, taskID)

// Extract URL from generateVideoResponse if available
if len(op.Response.GenerateVideoResponse.GeneratedSamples) > 0 {
if uri := op.Response.GenerateVideoResponse.GeneratedSamples[0].Video.URI; uri != "" {
taskID := encodeLocalTaskID(op.Name)
ti.Url = fmt.Sprintf("%s/v1/videos/%s/content?url=%s", system_setting.ServerAddress, taskID, uri)
ti.RemoteUrl = uri
}
}

return ti, nil
}

func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
upstreamName, err := decodeLocalTaskID(task.TaskID)
if err != nil {
upstreamName = ""
}
modelName := extractModelFromOperationName(upstreamName)
if strings.TrimSpace(modelName) == "" {
modelName = "veo-3.0-generate-001"
}

video := dto.NewOpenAIVideo()
video.ID = task.TaskID
video.Model = modelName
video.Status = task.Status.ToVideoStatus()
video.SetProgressStr(task.Progress)
video.CreatedAt = task.CreatedAt
if task.FinishTime > 0 {
video.CompletedAt = task.FinishTime
} else if task.UpdatedAt > 0 {
video.CompletedAt = task.UpdatedAt
}

return common.Marshal(video)
}

// ============================
// helpers
// ============================
Expand All @@ -274,3 +304,21 @@ func decodeLocalTaskID(local string) (string, error) {
}
return string(b), nil
}

var modelRe = regexp.MustCompile(`models/([^/]+)/operations/`)

func extractModelFromOperationName(name string) string {
if name == "" {
return ""
}
if m := modelRe.FindStringSubmatch(name); len(m) == 2 {
return m[1]
}
if idx := strings.Index(name, "models/"); idx >= 0 {
s := name[idx+len("models/"):]
if p := strings.Index(s, "/operations/"); p > 0 {
return s[:p]
}
}
return ""
}
1 change: 1 addition & 0 deletions relay/common/relay_info.go
Original file line number Diff line number Diff line change
Expand Up @@ -509,6 +509,7 @@ type TaskInfo struct {
Status string `json:"status"`
Reason string `json:"reason,omitempty"`
Url string `json:"url,omitempty"`
RemoteUrl string `json:"remote_url,omitempty"`
Progress string `json:"progress,omitempty"`
CompletionTokens int `json:"completion_tokens,omitempty"` // 用于按倍率计费
TotalTokens int `json:"total_tokens,omitempty"` // 用于按倍率计费
Expand Down
2 changes: 1 addition & 1 deletion relay/relay_task.go
Original file line number Diff line number Diff line change
Expand Up @@ -319,7 +319,7 @@ func videoFetchByIDRespBodyBuilder(c *gin.Context) (respBody []byte, taskResp *d
if err2 != nil {
return
}
if channelModel.Type != constant.ChannelTypeVertexAi {
if channelModel.Type != constant.ChannelTypeVertexAi && channelModel.Type != constant.ChannelTypeGemini {
return
}
baseURL := constant.ChannelBaseURLs[channelModel.Type]
Expand Down