new-api/controller/relay_gemini_file.go

package controller

import (
	"fmt"
	"net/http"
	"strings"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/logger"
	"github.com/QuantumNous/new-api/relay/channel/gemini"

	"github.com/gin-gonic/gin"
)

// 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 {
		logger.LogError(c, fmt.Sprintf("failed to parse multipart form: %s", err.Error()))
		c.JSON(http.StatusBadRequest, gin.H{
			"error": gin.H{
				"message": fmt.Sprintf("failed to parse multipart form: %s", err.Error()),
				"type":    "invalid_request_error",
				"code":    "invalid_multipart_form",
			},
		})
		return
	}
	defer form.RemoveAll()

	// Rebuild multipart form for upstream request
	body, contentType, err := gemini.RebuildMultipartForm(form)
	if err != nil {
		logger.LogError(c, fmt.Sprintf("failed to rebuild multipart form: %s", err.Error()))
		c.JSON(http.StatusInternalServerError, gin.H{
			"error": gin.H{
				"message": fmt.Sprintf("failed to rebuild multipart form: %s", err.Error()),
				"type":    "internal_error",
				"code":    "form_rebuild_error",
			},
		})
		return
	}

	// Build upstream URL
	url := gemini.BuildGeminiFileURL("/upload/v1beta/files")

	// Prepare headers
	headers := map[string]string{
		"Content-Type":   contentType,
		"x-goog-api-key": apiKey,
	}

	// Forward request to Gemini
	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()))
		return
	}
}

// RelayGeminiFileList lists files from Gemini File API
func RelayGeminiFileList(c *gin.Context) {
	// Get API key from channel context
	apiKey := common.GetContextKeyString(c, constant.ContextKeyChannelKey)
	if apiKey == "" {
		logger.LogError(c, "API key not found in context")
		c.JSON(http.StatusUnauthorized, gin.H{
			"error": gin.H{
				"message": "API key not found",
				"type":    "authentication_error",
				"code":    "invalid_api_key",
			},
		})
		return
	}

	// Build upstream URL with query parameters
	url := gemini.BuildGeminiFileURL("/v1beta/files")

	// Add query parameters if present
	queryParams := []string{}
	if pageSize := c.Query("pageSize"); pageSize != "" {
		queryParams = append(queryParams, fmt.Sprintf("pageSize=%s", pageSize))
	}
	if pageToken := c.Query("pageToken"); pageToken != "" {
		queryParams = append(queryParams, fmt.Sprintf("pageToken=%s", pageToken))
	}

	if len(queryParams) > 0 {
		url = fmt.Sprintf("%s?%s", url, strings.Join(queryParams, "&"))
	}

	// Prepare headers
	headers := map[string]string{
		"x-goog-api-key": apiKey,
	}

	// Forward request to Gemini
	err := gemini.ForwardGeminiFileRequest(c, http.MethodGet, url, nil, headers)
	if err != nil {
		logger.LogError(c, fmt.Sprintf("failed to forward file list request: %s", err.Error()))
		return
	}
}