Bläddra i källkod

feat: support Gemini resumable upload protocol

- Add support for X-Goog-Upload-Protocol: resumable
- Handle both resumable and multipart upload methods
- Forward all X-Goog-Upload-* headers for resumable uploads
- Properly set MIME type via X-Goog-Upload-Header-Content-Type
- Maintain backward compatibility with multipart uploads
- Keep Gemini File API completely isolated
supeng 1 månad sedan
förälder
incheckning
56494f9cad
2 ändrade filer med 74 tillägg och 22 borttagningar
  1. 67 15
      controller/relay_gemini_file.go
  2. 7 7
      docker-compose.yml

+ 67 - 15
controller/relay_gemini_file.go

@@ -14,7 +14,74 @@ import (
 )
 
 // RelayGeminiFileUpload handles file upload to Gemini File API
+// Supports both simple multipart upload and resumable upload protocol
 func RelayGeminiFileUpload(c *gin.Context) {
+	// Get API key from channel context
+	apiKey := common.GetContextKeyString(c, constant.ContextKeyChannelKey)
+	if apiKey == "" {
+		logger.LogError(c, "Failed to get Gemini channel API key")
+		c.JSON(http.StatusServiceUnavailable, gin.H{
+			"error": gin.H{
+				"message": "No available Gemini channel found",
+				"type":    "service_unavailable_error",
+				"code":    "no_available_channel",
+			},
+		})
+		return
+	}
+
+	// Check if this is a resumable upload request
+	uploadProtocol := c.GetHeader("X-Goog-Upload-Protocol")
+	if uploadProtocol == "resumable" {
+		// Handle resumable upload
+		handleResumableUpload(c, apiKey)
+		return
+	}
+
+	// Handle standard multipart upload
+	handleMultipartUpload(c, apiKey)
+}
+
+// handleResumableUpload handles Gemini's resumable upload protocol
+func handleResumableUpload(c *gin.Context, apiKey string) {
+	// Build upstream URL
+	url := gemini.BuildGeminiFileURL("/upload/v1beta/files")
+
+	// Prepare headers - forward all X-Goog-Upload-* headers
+	headers := map[string]string{
+		"x-goog-api-key": apiKey,
+	}
+
+	// Forward all resumable upload headers
+	resumableHeaders := []string{
+		"X-Goog-Upload-Protocol",
+		"X-Goog-Upload-Command",
+		"X-Goog-Upload-Header-Content-Length",
+		"X-Goog-Upload-Header-Content-Type",
+		"X-Goog-Upload-Offset",
+		"Content-Type",
+		"Content-Length",
+	}
+
+	for _, header := range resumableHeaders {
+		if value := c.GetHeader(header); value != "" {
+			headers[header] = value
+		}
+	}
+
+	// Get request body (for start command, it's JSON metadata)
+	body := c.Request.Body
+
+	// Forward request to Gemini
+	err := gemini.ForwardGeminiFileRequest(c, http.MethodPost, url, body, headers)
+	if err != nil {
+		logger.LogError(c, fmt.Sprintf("failed to forward resumable upload request: %s", err.Error()))
+		return
+	}
+}
+
+// handleMultipartUpload handles standard multipart/form-data upload
+func handleMultipartUpload(c *gin.Context, apiKey string) {
 	// Parse multipart form
 	form, err := common.ParseMultipartFormReusable(c)
 	if err != nil {
@@ -30,20 +97,6 @@ func RelayGeminiFileUpload(c *gin.Context) {
 	}
 	defer form.RemoveAll()
 
-	// Get API key from channel context (set by setupGeminiFileChannel)
-	apiKey := common.GetContextKeyString(c, constant.ContextKeyChannelKey)
-	if apiKey == "" {
-		logger.LogError(c, "Failed to get Gemini channel API key")
-		c.JSON(http.StatusServiceUnavailable, gin.H{
-			"error": gin.H{
-				"message": "No available Gemini channel found",
-				"type":    "service_unavailable_error",
-				"code":    "no_available_channel",
-			},
-		})
-		return
-	}
-
 	// Rebuild multipart form for upstream request
 	body, contentType, err := gemini.RebuildMultipartForm(form)
 	if err != nil {
@@ -71,7 +124,6 @@ func RelayGeminiFileUpload(c *gin.Context) {
 	err = gemini.ForwardGeminiFileRequest(c, http.MethodPost, url, body, headers)
 	if err != nil {
 		logger.LogError(c, fmt.Sprintf("failed to forward file upload request: %s", err.Error()))
-		// Error response already sent by ForwardGeminiFileRequest
 		return
 	}
 }

+ 7 - 7
docker-compose.yml

@@ -5,18 +5,18 @@
 #   2. Access at http://localhost:3000
 #
 # Using MySQL instead of PostgreSQL:
-#   1. Comment out the postgres service and SQL_DSN line 15
-#   2. Uncomment the mysql service and SQL_DSN line 16
-#   3. Uncomment mysql in depends_on (line 28)
-#   4. Uncomment mysql_data in volumes section (line 64)
+#   1. Comment out the postgres service and SQL_DSN line 29
+#   2. Uncomment the mysql service and SQL_DSN line 30
+#   3. Uncomment mysql in depends_on (line 46)
+#   4. Uncomment mysql_data in volumes section (line 85)
 #
 # ⚠️  IMPORTANT: Change all default passwords before deploying to production!
 
-version: '3.4' # For compatibility with older Docker versions
-
 services:
   new-api:
-    image: calciumion/new-api:latest
+    build:
+      context: .
+      dockerfile: Dockerfile
     container_name: new-api
     restart: always
     command: --log-dir /app/logs