2
0
Seefs 1 сар өмнө
parent
commit
9629c8a771

+ 24 - 2
controller/video_proxy.go

@@ -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)

+ 158 - 0
controller/video_proxy_gemini.go

@@ -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)
+}

+ 44 - 9
model/task.go

@@ -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) {
@@ -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
@@ -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
 }

+ 50 - 2
relay/channel/task/gemini/adaptor.go

@@ -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"
@@ -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
 // ============================
@@ -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 - 0
relay/common/relay_info.go

@@ -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"`      // 用于按倍率计费

+ 1 - 1
relay/relay_task.go

@@ -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]