fix veo3 (#2140)

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