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
func RelayGeminiFileUpload(c *gin.Context) {
	// 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()

	// 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 {
		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()))
		// Error response already sent by ForwardGeminiFileRequest
		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
	}
}