mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-18 12:14:24 +00:00
feat(tools): implement piping stream for Gemini File API from video URL
This commit is contained in:
1 parent
abb0fab94c
commit
fef98b4dd7
3 files changed
+276
-15
No files matched your search
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user