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