From fef98b4dd76a494315309b65a0727beed8898c3b Mon Sep 17 00:00:00 2001 From: thotam Date: Sat, 6 Jun 2026 19:14:13 +0700 Subject: [PATCH] feat(tools): implement piping stream for Gemini File API from video URL --- internal/tools/gemini_file_api.go | 150 +++++++++++++++++++++++++++ internal/tools/read_video_resolve.go | 73 ++++++++++++- internal/tools/read_video_test.go | 68 ++++++++++-- 3 files changed, 276 insertions(+), 15 deletions(-) diff --git a/internal/tools/gemini_file_api.go b/internal/tools/gemini_file_api.go index fe96b941..4ab1dd64 100644 --- a/internal/tools/gemini_file_api.go +++ b/internal/tools/gemini_file_api.go @@ -264,3 +264,153 @@ func parseGeminiResponse(respBody []byte) (*providers.ChatResponse, error) { }, }, nil } + +// geminiFileUploadStream uploads a file stream to Gemini File API using resumable upload protocol. +// Returns the file name (e.g. "files/abc123") and file URI for use in generateContent. +func geminiFileUploadStream(ctx context.Context, apiKey, displayName string, reader io.Reader, contentLength int64, mime string) (fileName, fileURI string, err error) { + // Step 1: Initiate resumable upload. + initBody, _ := json.Marshal(map[string]any{ + "file": map[string]string{"display_name": displayName}, + }) + initReq, err := http.NewRequestWithContext(ctx, "POST", geminiUploadBase+"?key="+apiKey, bytes.NewReader(initBody)) + if err != nil { + return "", "", fmt.Errorf("create init request: %w", err) + } + initReq.Header.Set("Content-Type", "application/json") + initReq.Header.Set("X-Goog-Upload-Protocol", "resumable") + initReq.Header.Set("X-Goog-Upload-Command", "start") + initReq.Header.Set("X-Goog-Upload-Header-Content-Length", fmt.Sprintf("%d", contentLength)) + initReq.Header.Set("X-Goog-Upload-Header-Content-Type", mime) + + client := &http.Client{Timeout: 60 * time.Second} + initResp, err := client.Do(initReq) + if err != nil { + return "", "", fmt.Errorf("init upload: %w", err) + } + defer initResp.Body.Close() + io.ReadAll(initResp.Body) // drain + + if initResp.StatusCode != 200 { + return "", "", fmt.Errorf("init upload HTTP %d", initResp.StatusCode) + } + + uploadURL := initResp.Header.Get("X-Goog-Upload-URL") + if uploadURL == "" { + return "", "", fmt.Errorf("no upload URL in response headers") + } + + // Step 2: Upload file stream. + uploadReq, err := http.NewRequestWithContext(ctx, "POST", uploadURL, reader) + if err != nil { + return "", "", fmt.Errorf("create upload request: %w", err) + } + uploadReq.ContentLength = contentLength + uploadReq.Header.Set("Content-Length", fmt.Sprintf("%d", contentLength)) + uploadReq.Header.Set("X-Goog-Upload-Offset", "0") + uploadReq.Header.Set("X-Goog-Upload-Command", "upload, finalize") + + // Do not set global Timeout on the HTTP client here because we are piping a potentially large stream. + uploadClient := &http.Client{} + uploadResp, err := uploadClient.Do(uploadReq) + if err != nil { + return "", "", fmt.Errorf("upload stream: %w", err) + } + defer uploadResp.Body.Close() + + respBody, err := io.ReadAll(uploadResp.Body) + if err != nil { + return "", "", fmt.Errorf("read upload response: %w", err) + } + if uploadResp.StatusCode != 200 { + return "", "", fmt.Errorf("upload HTTP %d: %s", uploadResp.StatusCode, truncateStr(string(respBody), 500)) + } + + var uploadResult struct { + File struct { + Name string `json:"name"` + URI string `json:"uri"` + State string `json:"state"` + } `json:"file"` + } + if err := json.Unmarshal(respBody, &uploadResult); err != nil { + return "", "", fmt.Errorf("parse upload response: %w", err) + } + + // Only return URI if file is already ACTIVE; otherwise caller must poll. + if uploadResult.File.State == "ACTIVE" { + return uploadResult.File.Name, uploadResult.File.URI, nil + } + return uploadResult.File.Name, "", nil +} + +// geminiFileAPICallStream uploads a file stream via Gemini File API, polls until ready, +// then calls generateContent with file_data reference. +func geminiFileAPICallStream(ctx context.Context, apiKey, model, prompt string, reader io.Reader, contentLength int64, mime string, httpTimeout time.Duration) (*providers.ChatResponse, error) { + displayName := fmt.Sprintf("goclaw_%d", time.Now().UnixNano()) + + slog.Info("gemini file api: uploading stream", "size", contentLength, "mime", mime) + fileName, fileURI, err := geminiFileUploadStream(ctx, apiKey, displayName, reader, contentLength, mime) + if err != nil { + return nil, fmt.Errorf("upload: %w", err) + } + slog.Info("gemini file api: uploaded stream", "name", fileName) + + // If file URI not returned directly, poll for it. + if fileURI == "" { + slog.Info("gemini file api: polling for active state", "name", fileName) + fileURI, err = geminiFilePoll(ctx, apiKey, fileName) + if err != nil { + return nil, fmt.Errorf("poll: %w", err) + } + } + slog.Info("gemini file api: file active", "uri", fileURI) + + // Call generateContent with file_data reference. + body := map[string]any{ + "contents": []map[string]any{ + { + "parts": []map[string]any{ + {"file_data": map[string]any{"mime_type": mime, "file_uri": fileURI}}, + {"text": prompt}, + }, + }, + }, + "generationConfig": map[string]any{ + "maxOutputTokens": 16384, + "temperature": 0.2, + }, + } + + bodyJSON, err := json.Marshal(body) + if err != nil { + return nil, fmt.Errorf("marshal request: %w", err) + } + + url := fmt.Sprintf("https://generativelanguage.googleapis.com/v1beta/models/%s:generateContent?key=%s", model, apiKey) + httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(bodyJSON)) + if err != nil { + return nil, fmt.Errorf("create request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + if httpTimeout == 0 { + httpTimeout = 120 * time.Second + } + client := &http.Client{Timeout: httpTimeout} + httpResp, err := client.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("HTTP request: %w", err) + } + defer httpResp.Body.Close() + + respBody, err := io.ReadAll(httpResp.Body) + if err != nil { + return nil, fmt.Errorf("read response: %w", err) + } + if httpResp.StatusCode != 200 { + return nil, fmt.Errorf("HTTP %d: %s", httpResp.StatusCode, truncateStr(string(respBody), 500)) + } + + // Parse — same response format as geminiNativeDocumentCall. + return parseGeminiResponse(respBody) +} diff --git a/internal/tools/read_video_resolve.go b/internal/tools/read_video_resolve.go index 53a49c0e..8e371bd1 100644 --- a/internal/tools/read_video_resolve.go +++ b/internal/tools/read_video_resolve.go @@ -5,7 +5,9 @@ import ( "encoding/base64" "fmt" "log/slog" + "net/http" "path/filepath" + "strings" "time" "github.com/nextlevelbuilder/goclaw/internal/providers" @@ -70,11 +72,9 @@ func (t *ReadVideoTool) callProvider(ctx context.Context, cp credentialProvider, // Gemini: use File API (requires credentials). ptype := GetParamString(params, "_provider_type", providerTypeFromName(providerName)) if cp != nil && ptype == "gemini" { - if videoURL != "" { - return nil, nil, fmt.Errorf("provider %q does not support analyzing videos directly from a URL", providerName) - } + var resp *providers.ChatResponse + var err error - slog.Info("read_video: using gemini file API", "provider", providerName, "model", model, "size", len(data), "mime", mime) chatReq := providers.ChatRequest{ Messages: []providers.Message{{Role: "user", Content: prompt}}, Model: model, @@ -84,7 +84,70 @@ func (t *ReadVideoTool) callProvider(ctx context.Context, cp credentialProvider, if reserveErr != nil { return nil, nil, reserveErr } - resp, err := geminiFileAPICall(ctx, cp.APIKey(), model, prompt, data, mime, 180*time.Second) + + if videoURL != "" { + slog.Info("read_video: streaming URL directly to Gemini File API", "provider", providerName, "model", model, "url", videoURL) + + // Send GET request to fetch the stream. + req, getErr := http.NewRequestWithContext(ctx, "GET", videoURL, nil) + if getErr != nil { + if reservation != nil { + reservation.Reconcile(ctx, nil, getErr) + } + return nil, nil, fmt.Errorf("failed to create GET request for video URL: %w", getErr) + } + + // No global Timeout on http.Client to allow piping large streams. + client := &http.Client{} + httpResp, getErr := client.Do(req) + if getErr != nil { + if reservation != nil { + reservation.Reconcile(ctx, nil, getErr) + } + return nil, nil, fmt.Errorf("failed to fetch video URL: %w", getErr) + } + defer httpResp.Body.Close() + + if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { + statusErr := fmt.Errorf("video URL returned status code %d", httpResp.StatusCode) + if reservation != nil { + reservation.Reconcile(ctx, nil, statusErr) + } + return nil, nil, statusErr + } + + // Validate Content-Length + contentLength := httpResp.ContentLength + if contentLength <= 0 { + invalidLenErr := fmt.Errorf("URL does not support static streaming (missing or invalid Content-Length: %d)", contentLength) + if reservation != nil { + reservation.Reconcile(ctx, nil, invalidLenErr) + } + return nil, nil, invalidLenErr + } + + // Check limits: maximum 2 GB + const videoMaxStreamBytes = 2 * 1024 * 1024 * 1024 // 2 GB + if contentLength > videoMaxStreamBytes { + limitErr := fmt.Errorf("video stream size (%d bytes) exceeds the maximum limit of 2 GB", contentLength) + if reservation != nil { + reservation.Reconcile(ctx, nil, limitErr) + } + return nil, nil, limitErr + } + + // Extract MIME type from Content-Type header if valid video type, otherwise use the inferred/passed one + contentType := httpResp.Header.Get("Content-Type") + if contentType != "" && strings.HasPrefix(contentType, "video/") { + mime = contentType + } + + resp, err = geminiFileAPICallStream(ctx, cp.APIKey(), model, prompt, httpResp.Body, contentLength, mime, 300*time.Second) + } else { + slog.Info("read_video: using gemini file API", "provider", providerName, "model", model, "size", len(data), "mime", mime) + resp, err = geminiFileAPICall(ctx, cp.APIKey(), model, prompt, data, mime, 180*time.Second) + } + if reservation != nil { reservation.Reconcile(ctx, resp, err) } diff --git a/internal/tools/read_video_test.go b/internal/tools/read_video_test.go index 73999222..31d9ba99 100644 --- a/internal/tools/read_video_test.go +++ b/internal/tools/read_video_test.go @@ -2,6 +2,8 @@ package tools import ( "context" + "net/http" + "net/http/httptest" "strings" "testing" ) @@ -32,23 +34,69 @@ func TestReadVideo_BothMediaIdAndUrl_Error(t *testing.T) { } } -func TestReadVideo_GeminiURL_Error(t *testing.T) { - tool := NewReadVideoTool(nil, nil) +func TestReadVideo_GeminiURL_Validation(t *testing.T) { + // 1. Trường hợp không có Content-Length (ContentLength <= 0) + ts1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Transfer-Encoding", "chunked") + w.Write([]byte("chunked data mock video")) + })) + defer ts1.Close() - params := map[string]any{ + tool := NewReadVideoTool(nil, nil) + cp := &mockCredentialProvider{apiKey: "test-key"} + + params1 := map[string]any{ "prompt": "describe this video", - "url": "https://example.com/video.mp4", + "url": ts1.URL, "_provider_type": "gemini", } - cp := &mockCredentialProvider{apiKey: "test-key"} - - _, _, err := tool.callProvider(context.Background(), cp, "gemini", "gemini-2.5-flash", params) + _, _, err := tool.callProvider(context.Background(), cp, "gemini", "gemini-2.5-flash", params1) if err == nil { - t.Fatalf("expected error for gemini native provider with video URL") + t.Fatalf("expected error for missing Content-Length") + } + if !strings.Contains(err.Error(), "URL does not support static streaming") { + t.Errorf("unexpected error for missing Content-Length: %v", err) } - if !strings.Contains(err.Error(), "does not support analyzing videos directly from a URL") { - t.Errorf("unexpected error message: %v", err) + // 2. Trường hợp Content-Length vượt quá 2 GB + ts2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Length", "2147483649") // 2GB + 1 byte + w.WriteHeader(http.StatusOK) + })) + defer ts2.Close() + + params2 := map[string]any{ + "prompt": "describe this video", + "url": ts2.URL, + "_provider_type": "gemini", + } + + _, _, err = tool.callProvider(context.Background(), cp, "gemini", "gemini-2.5-flash", params2) + if err == nil { + t.Fatalf("expected error for Content-Length exceeding 2GB") + } + if !strings.Contains(err.Error(), "exceeds the maximum limit of 2 GB") { + t.Errorf("unexpected error for limit exceed: %v", err) + } + + // 3. Trường hợp HTTP status code lỗi (ví dụ 404) + ts3 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer ts3.Close() + + params3 := map[string]any{ + "prompt": "describe this video", + "url": ts3.URL, + "_provider_type": "gemini", + } + + _, _, err = tool.callProvider(context.Background(), cp, "gemini", "gemini-2.5-flash", params3) + if err == nil { + t.Fatalf("expected error for HTTP 404 status code") + } + if !strings.Contains(err.Error(), "video URL returned status code 404") { + t.Errorf("unexpected error for HTTP status code: %v", err) } }