wangbo
4fa44f4c41
媒体请求从认证、任务创建到异步准入登记共用可配置的请求并发令牌，完成登记后立即释放。这样仍会尽早丢弃 Base64 正文，但不会让高并发请求各占一个 PostgreSQL 连接等待 advisory lock。\n\n扩展 Gemini 双角色压力测试，使用生产同档 API/Worker 连接池并记录数据库连接与 advisory lock 峰值。\n\n验证：go test ./... -count=1；go vet ./...；pnpm openapi；迁移安全检查；256 个 256 KiB 异步 Gemini 请求全部通过，连接峰值 45、锁等待峰值 14。
844 lines
28 KiB
Go
844 lines
28 KiB
Go
package httpapi
|
||
|
||
import (
|
||
"context"
|
||
"crypto/rand"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"mime"
|
||
"net/http"
|
||
"path"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
|
||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||
)
|
||
|
||
type geminiImageTaskMapping struct {
|
||
Kind string
|
||
Body map[string]any
|
||
Model string
|
||
}
|
||
|
||
type geminiImageInput struct {
|
||
URI string
|
||
MimeType string
|
||
}
|
||
|
||
type geminiUploadSession struct {
|
||
Version string
|
||
DisplayName string
|
||
MimeType string
|
||
SizeBytes int64
|
||
CreatedAt time.Time
|
||
}
|
||
|
||
var geminiGenerateContentRoutePrefixes = []string{
|
||
"/api/v1/models/",
|
||
"/v1beta/models/",
|
||
"/v1/models/",
|
||
}
|
||
|
||
func (s *Server) registerGeminiGenerateContentRoutes(mux *http.ServeMux) {
|
||
handler := s.requireProtocolUser(clients.ProtocolGeminiGenerateContent, http.HandlerFunc(s.geminiGenerateContent))
|
||
for _, prefix := range geminiGenerateContentRoutePrefixes {
|
||
mux.Handle("POST "+prefix, geminiGenerateContentRouteHandler(prefix, s.withMediaRequestBodySlot(handler)))
|
||
}
|
||
}
|
||
|
||
func geminiGenerateContentRouteHandler(prefix string, next http.Handler) http.Handler {
|
||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
model, stream, ok := geminiContentModelFromPath(prefix, r.URL.Path)
|
||
if !ok {
|
||
http.NotFound(w, r)
|
||
return
|
||
}
|
||
r.SetPathValue("model", model)
|
||
r.SetPathValue("stream", strconv.FormatBool(stream))
|
||
next.ServeHTTP(w, r)
|
||
})
|
||
}
|
||
|
||
func geminiContentModelFromPath(prefix string, requestPath string) (string, bool, bool) {
|
||
if !strings.HasPrefix(requestPath, prefix) {
|
||
return "", false, false
|
||
}
|
||
relative := strings.TrimPrefix(requestPath, prefix)
|
||
for _, action := range []struct {
|
||
suffix string
|
||
stream bool
|
||
}{{":generateContent", false}, {":streamGenerateContent", true}} {
|
||
model, ok := strings.CutSuffix(relative, action.suffix)
|
||
if ok && strings.TrimSpace(model) != "" && !strings.Contains(model, "/") {
|
||
return model, action.stream, true
|
||
}
|
||
}
|
||
return "", false, false
|
||
}
|
||
|
||
func geminiGenerateContentModelFromPath(prefix string, requestPath string) (string, bool) {
|
||
model, stream, ok := geminiContentModelFromPath(prefix, requestPath)
|
||
if !ok || stream {
|
||
return "", false
|
||
}
|
||
return model, true
|
||
}
|
||
|
||
// geminiGenerateContent godoc
|
||
// @Summary Gemini generateContent 兼容接口
|
||
// @Description 使用统一 /api/v1 前缀接收 Gemini generateContent 请求；旧 /v1 和 /v1beta 路径保留兼容。
|
||
// @Tags gemini-compatible
|
||
// @Accept json
|
||
// @Produce json
|
||
// @Produce text/event-stream
|
||
// @Security BearerAuth
|
||
// @Param model path string true "模型名称"
|
||
// @Param input body map[string]interface{} true "Gemini generateContent 请求"
|
||
// @Success 200 {object} map[string]interface{}
|
||
// @Failure 400 {object} GeminiErrorEnvelope
|
||
// @Failure 401 {object} GeminiErrorEnvelope
|
||
// @Failure 403 {object} GeminiErrorEnvelope
|
||
// @Failure 502 {object} GeminiErrorEnvelope
|
||
// @Router /api/v1/models/{model}:generateContent [post]
|
||
// @Router /api/v1/models/{model}:streamGenerateContent [post]
|
||
func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
|
||
writeGeminiTaskError := func(status int, message string, details map[string]any, code string) {
|
||
writeProtocolError(w, clients.ProtocolGeminiGenerateContent, status, message, details, code)
|
||
}
|
||
user, ok := auth.UserFromContext(r.Context())
|
||
if !ok {
|
||
writeGeminiTaskError(http.StatusUnauthorized, "unauthorized", nil, "unauthorized")
|
||
return
|
||
}
|
||
admission, err := s.authorizeTaskTraffic(r, user)
|
||
if err != nil {
|
||
writeTaskTrafficError(w, err, clients.ProtocolGeminiGenerateContent)
|
||
return
|
||
}
|
||
releaseRequestBody, err := s.requestMediaBodyRelease(r.Context())
|
||
if err != nil {
|
||
return
|
||
}
|
||
defer func() {
|
||
if releaseRequestBody != nil {
|
||
releaseRequestBody()
|
||
}
|
||
}()
|
||
var native map[string]any
|
||
if err := json.NewDecoder(r.Body).Decode(&native); err != nil {
|
||
writeGeminiTaskError(http.StatusBadRequest, "invalid json body", nil, "invalid_json_body")
|
||
return
|
||
}
|
||
mapping, err := geminiImageTaskBody(r.PathValue("model"), native)
|
||
if err != nil {
|
||
writeGeminiTaskError(http.StatusBadRequest, err.Error(), nil, clients.ErrorCode(err))
|
||
return
|
||
}
|
||
streamMode := r.PathValue("stream") == "true"
|
||
if streamMode {
|
||
mapping.Body["stream"] = true
|
||
}
|
||
if !apiKeyScopeAllowed(user, mapping.Kind) {
|
||
writeGeminiTaskError(http.StatusForbidden, "api key scope does not allow this capability", nil, "permission_denied")
|
||
return
|
||
}
|
||
idempotencyKey, hasIdempotencyKey, err := optionalTaskIdempotencyKey(r)
|
||
if err != nil {
|
||
writeGeminiTaskError(http.StatusBadRequest, "Idempotency-Key must contain one non-empty value", nil, "invalid_idempotency_key")
|
||
return
|
||
}
|
||
prepared, err := s.prepareTaskRequest(r.Context(), r, user, mapping.Body)
|
||
if err != nil {
|
||
s.logger.Warn("prepare gemini task request failed", "kind", mapping.Kind, "error", err)
|
||
status := http.StatusBadRequest
|
||
if code := clients.ErrorCode(err); strings.HasPrefix(code, "upload_") || code == "request_asset_upload_failed" {
|
||
status = http.StatusBadGateway
|
||
}
|
||
writeGeminiTaskError(status, err.Error(), nil, clients.ErrorCode(err))
|
||
return
|
||
}
|
||
// The queued task only needs the materialized assetRef. Drop the native
|
||
// Base64 request body now, but retain the lightweight request slot until
|
||
// distributed admission is durably registered. This prevents a burst from
|
||
// turning advisory-lock waiters into one PostgreSQL connection per request.
|
||
native = nil
|
||
mapping.Body = nil
|
||
runMode, err := s.admittedTaskRunMode(admission, prepared.Body)
|
||
if err != nil {
|
||
writeTaskTrafficError(w, err, clients.ProtocolGeminiGenerateContent)
|
||
return
|
||
}
|
||
createInput := store.CreateTaskInput{
|
||
Kind: mapping.Kind,
|
||
Model: mapping.Model,
|
||
RunMode: runMode,
|
||
AcceptanceRunID: admission.AcceptanceRunID,
|
||
Async: !streamMode,
|
||
Request: prepared.Body,
|
||
ConversationID: prepared.ConversationID,
|
||
NewMessageCount: prepared.NewMessageCount,
|
||
MessageRefs: prepared.MessageRefs,
|
||
}
|
||
if hasIdempotencyKey {
|
||
createInput.IdempotencyKeyHash = taskIdempotencyKeyHash(idempotencyKey)
|
||
createInput.IdempotencyRequestHash = taskIdempotencyRequestHash(mapping.Kind, false, streamMode, prepared.Body)
|
||
}
|
||
created, err := s.store.CreateTaskIdempotent(r.Context(), createInput, user)
|
||
if err != nil {
|
||
if errors.Is(err, store.ErrIdempotencyKeyReused) {
|
||
writeGeminiTaskError(http.StatusConflict, "Idempotency-Key was reused for a different request", nil, "idempotency_key_reused")
|
||
return
|
||
}
|
||
if errors.Is(err, store.ErrTaskRequestBinaryNotMaterialized) {
|
||
writeGeminiTaskError(http.StatusBadRequest, err.Error(), nil, clients.ErrorCode(err))
|
||
return
|
||
}
|
||
s.logger.Error("create gemini task failed", "kind", mapping.Kind, "error_category", "task_create_failed")
|
||
writeGeminiTaskError(http.StatusInternalServerError, "create task failed", nil, "task_create_failed")
|
||
return
|
||
}
|
||
task := created.Task
|
||
if created.Replayed {
|
||
if s.billingMetrics != nil {
|
||
s.billingMetrics.ObserveBillingEvent("idempotent_replay")
|
||
}
|
||
if streamMode {
|
||
writeGeminiTaskError(http.StatusConflict, "streaming idempotent replay is not supported", nil, "idempotency_stream_replay_unsupported")
|
||
return
|
||
}
|
||
if task.Status == "succeeded" {
|
||
task, err = s.hydrateTaskResult(r.Context(), task)
|
||
if err != nil {
|
||
writeGeminiTaskError(statusFromRunError(err), err.Error(), nil, clients.ErrorCode(err))
|
||
return
|
||
}
|
||
if _, native := task.Result["candidates"]; native {
|
||
writeJSON(w, http.StatusOK, task.Result)
|
||
return
|
||
}
|
||
writeJSON(w, http.StatusOK, geminiGenerateContentResponse(task.Result, mapping.Model))
|
||
return
|
||
}
|
||
status := http.StatusConflict
|
||
if task.Status == "failed" || task.Status == "cancelled" {
|
||
status = storedTaskErrorStatus(task.ErrorCode)
|
||
}
|
||
writeGeminiTaskError(status, firstNonEmpty(task.ErrorMessage, task.Error, task.Message, "task is not complete"), nil, firstNonEmpty(task.ErrorCode, task.Status))
|
||
return
|
||
}
|
||
runCtx, cancelRun := s.requestExecutionContext(r)
|
||
defer cancelRun()
|
||
if streamMode {
|
||
releaseRequestBody()
|
||
releaseRequestBody = nil
|
||
s.writeGeminiGenerateContentStream(runCtx, w, r, task, user, mapping.Model, writeGeminiTaskError)
|
||
return
|
||
}
|
||
if err := s.runner.SubmitAsyncTask(runCtx, task); err != nil {
|
||
if !requestStillConnected(r) {
|
||
return
|
||
}
|
||
applyRunErrorHeaders(w, err)
|
||
writeGeminiTaskError(statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), runErrorCode(err))
|
||
return
|
||
}
|
||
releaseRequestBody()
|
||
releaseRequestBody = nil
|
||
task, err = s.runner.WaitForTaskCompletion(runCtx, task.ID)
|
||
if err != nil || !requestStillConnected(r) {
|
||
return
|
||
}
|
||
if task.Status != "succeeded" {
|
||
status := storedTaskErrorStatus(task.ErrorCode)
|
||
writeGeminiTaskError(status, firstNonEmpty(task.ErrorMessage, task.Error, task.Message, "task failed"), nil, firstNonEmpty(task.ErrorCode, task.Status))
|
||
return
|
||
}
|
||
task, err = s.hydrateTaskResult(runCtx, task)
|
||
if err != nil {
|
||
writeGeminiTaskError(statusFromRunError(err), err.Error(), nil, clients.ErrorCode(err))
|
||
return
|
||
}
|
||
if !requestStillConnected(r) {
|
||
return
|
||
}
|
||
if _, native := task.Result["candidates"]; native {
|
||
writeJSON(w, http.StatusOK, task.Result)
|
||
return
|
||
}
|
||
writeJSON(w, http.StatusOK, geminiGenerateContentResponse(task.Result, mapping.Model))
|
||
}
|
||
|
||
func (s *Server) writeGeminiGenerateContentStream(runCtx context.Context, w http.ResponseWriter, r *http.Request, task store.GatewayTask, user *auth.User, model string, writeGeminiTaskError func(int, string, map[string]any, string)) {
|
||
w.Header().Set("Content-Type", "text/event-stream")
|
||
w.Header().Set("Cache-Control", "no-cache")
|
||
flusher, _ := w.(http.Flusher)
|
||
nativePassthrough := false
|
||
convertedFrames := false
|
||
result, runErr := s.runner.ExecuteStream(runCtx, task, user, func(delta clients.StreamDeltaEvent) error {
|
||
if !requestStillConnected(r) {
|
||
return nil
|
||
}
|
||
if delta.Event != nil && delta.WireProtocol == clients.ProtocolGeminiGenerateContent && !delta.WireConverted {
|
||
for name, values := range delta.WireHeaders {
|
||
for _, value := range values {
|
||
if strings.EqualFold(name, "Content-Type") {
|
||
w.Header().Set(name, value)
|
||
} else {
|
||
w.Header().Add(name, value)
|
||
}
|
||
}
|
||
}
|
||
nativePassthrough = true
|
||
writeGeminiSSEFrame(w, delta.Event)
|
||
} else if delta.Text != "" {
|
||
convertedFrames = true
|
||
writeGeminiSSEFrame(w, geminiTextStreamFrame(delta.Text))
|
||
}
|
||
if flusher != nil {
|
||
flusher.Flush()
|
||
}
|
||
return nil
|
||
})
|
||
if runErr != nil {
|
||
if !requestStillConnected(r) {
|
||
return
|
||
}
|
||
applyRunErrorHeaders(w, runErr)
|
||
if !nativePassthrough && !convertedFrames {
|
||
if wire := clients.ErrorWireResponse(runErr); wireResponseMatches(wire, clients.ProtocolGeminiGenerateContent) {
|
||
writeWireResponse(w, wire)
|
||
return
|
||
}
|
||
writeGeminiTaskError(statusFromRunError(runErr), runErrorMessage(runErr), runErrorDetails(runErr), runErrorCode(runErr))
|
||
} else if !nativePassthrough {
|
||
status := statusFromRunError(runErr)
|
||
writeGeminiSSEFrame(w, geminiErrorEnvelope(status, runErrorMessage(runErr), runErrorDetails(runErr), runErrorCode(runErr)))
|
||
if flusher != nil {
|
||
flusher.Flush()
|
||
}
|
||
}
|
||
return
|
||
}
|
||
if !nativePassthrough && !convertedFrames && requestStillConnected(r) {
|
||
writeGeminiSSEFrame(w, geminiGenerateContentResponse(result.Output, model))
|
||
if flusher != nil {
|
||
flusher.Flush()
|
||
}
|
||
}
|
||
}
|
||
|
||
func writeGeminiSSEFrame(w http.ResponseWriter, payload map[string]any) {
|
||
raw, _ := json.Marshal(payload)
|
||
_, _ = fmt.Fprintf(w, "data: %s\n\n", raw)
|
||
}
|
||
|
||
func geminiTextStreamFrame(text string) map[string]any {
|
||
return map[string]any{"candidates": []any{map[string]any{
|
||
"content": map[string]any{"role": "model", "parts": []any{map[string]any{"text": text}}},
|
||
}}}
|
||
}
|
||
|
||
func geminiImageTaskBody(model string, native map[string]any) (geminiImageTaskMapping, error) {
|
||
model = strings.TrimSpace(model)
|
||
if model == "" {
|
||
return geminiImageTaskMapping{}, &clients.ClientError{Code: "bad_request", Message: "model is required", Retryable: false}
|
||
}
|
||
if native == nil {
|
||
native = map[string]any{}
|
||
}
|
||
prompt := geminiPromptFromNative(native)
|
||
imageInputs := geminiImageInputsFromNative(native)
|
||
body := cloneMap(native)
|
||
body["model"] = model
|
||
body["prompt"] = prompt
|
||
if count := geminiCandidateCount(native); count > 0 {
|
||
body["n"] = count
|
||
}
|
||
if aspectRatio, imageSize := geminiImageConfig(native); aspectRatio != "" {
|
||
body["aspect_ratio"] = aspectRatio
|
||
} else if imageSize != "" {
|
||
body["resolution"] = imageSize
|
||
}
|
||
if _, imageSize := geminiImageConfig(native); imageSize != "" {
|
||
body["resolution"] = imageSize
|
||
}
|
||
if len(imageInputs) == 0 {
|
||
body["modelType"] = "image_generate"
|
||
return geminiImageTaskMapping{Kind: "images.generations", Body: body, Model: model}, nil
|
||
}
|
||
images := make([]any, 0, len(imageInputs))
|
||
for _, input := range imageInputs {
|
||
images = append(images, input.URI)
|
||
}
|
||
if len(images) == 1 {
|
||
body["image"] = images[0]
|
||
} else {
|
||
body["image"] = images
|
||
}
|
||
body["modelType"] = "image_edit"
|
||
return geminiImageTaskMapping{Kind: "images.edits", Body: body, Model: model}, nil
|
||
}
|
||
|
||
func geminiPromptFromNative(body map[string]any) string {
|
||
parts := []string{}
|
||
for _, part := range geminiOrderedNativeParts(body) {
|
||
if text := strings.TrimSpace(stringFromRequestAny(part["text"])); text != "" {
|
||
parts = append(parts, text)
|
||
}
|
||
}
|
||
if len(parts) > 0 {
|
||
return strings.Join(parts, "\n")
|
||
}
|
||
return firstNonEmptyRequestString(body, "prompt", "positive", "input")
|
||
}
|
||
|
||
func geminiImageInputsFromNative(body map[string]any) []geminiImageInput {
|
||
inputs := []geminiImageInput{}
|
||
for _, part := range geminiOrderedNativeParts(body) {
|
||
inline := mapFromGeminiAny(firstPresentGemini(part["inlineData"], part["inline_data"]))
|
||
if data := strings.TrimSpace(stringFromRequestAny(inline["data"])); data != "" {
|
||
mimeType := firstNonEmptyGeminiString(inline["mimeType"], inline["mime_type"])
|
||
if mimeType == "" {
|
||
mimeType = "image/png"
|
||
}
|
||
if !strings.HasPrefix(strings.ToLower(data), "data:") {
|
||
data = "data:" + mimeType + ";base64," + data
|
||
}
|
||
inputs = append(inputs, geminiImageInput{URI: data, MimeType: mimeType})
|
||
continue
|
||
}
|
||
fileData := mapFromGeminiAny(firstPresentGemini(part["fileData"], part["file_data"]))
|
||
if fileURI := firstNonEmptyGeminiString(fileData["fileUri"], fileData["file_uri"], fileData["uri"]); fileURI != "" {
|
||
inputs = append(inputs, geminiImageInput{
|
||
URI: fileURI,
|
||
MimeType: firstNonEmptyGeminiString(fileData["mimeType"], fileData["mime_type"]),
|
||
})
|
||
}
|
||
}
|
||
return inputs
|
||
}
|
||
|
||
func geminiOrderedNativeParts(body map[string]any) []map[string]any {
|
||
contents := normalizeGeminiContents(body["contents"])
|
||
parts := []map[string]any{}
|
||
for _, content := range contents {
|
||
contentMap := mapFromGeminiAny(content)
|
||
rawParts, _ := contentMap["parts"].([]any)
|
||
if len(rawParts) == 0 {
|
||
if text := stringFromRequestAny(contentMap["text"]); text != "" {
|
||
rawParts = []any{map[string]any{"text": text}}
|
||
}
|
||
}
|
||
for _, rawPart := range rawParts {
|
||
switch typed := rawPart.(type) {
|
||
case string:
|
||
parts = append(parts, map[string]any{"text": typed})
|
||
case map[string]any:
|
||
parts = append(parts, typed)
|
||
}
|
||
}
|
||
}
|
||
return parts
|
||
}
|
||
|
||
func normalizeGeminiContents(value any) []any {
|
||
switch typed := value.(type) {
|
||
case []any:
|
||
out := make([]any, 0, len(typed))
|
||
for _, item := range typed {
|
||
if text, ok := item.(string); ok {
|
||
out = append(out, map[string]any{"parts": []any{map[string]any{"text": text}}})
|
||
} else {
|
||
out = append(out, item)
|
||
}
|
||
}
|
||
return out
|
||
case string:
|
||
return []any{map[string]any{"parts": []any{map[string]any{"text": typed}}}}
|
||
case map[string]any:
|
||
return []any{typed}
|
||
default:
|
||
return nil
|
||
}
|
||
}
|
||
|
||
func geminiCandidateCount(body map[string]any) int {
|
||
generationConfig := mapFromGeminiAny(firstPresentGemini(body["generationConfig"], body["generation_config"]))
|
||
for _, value := range []any{body["n"], body["candidateCount"], body["candidate_count"], generationConfig["n"], generationConfig["candidateCount"], generationConfig["candidate_count"]} {
|
||
if count := positiveIntFromGeminiAny(value); count > 0 {
|
||
return count
|
||
}
|
||
}
|
||
return 0
|
||
}
|
||
|
||
func geminiImageConfig(body map[string]any) (string, string) {
|
||
generationConfig := mapFromGeminiAny(firstPresentGemini(body["generationConfig"], body["generation_config"]))
|
||
imageConfig := mapFromGeminiAny(firstPresentGemini(generationConfig["imageConfig"], generationConfig["image_config"]))
|
||
return firstNonEmptyGeminiString(imageConfig["aspectRatio"], imageConfig["aspect_ratio"]),
|
||
firstNonEmptyGeminiString(imageConfig["imageSize"], imageConfig["image_size"])
|
||
}
|
||
|
||
func geminiGenerateContentResponse(output map[string]any, model string) map[string]any {
|
||
parts := geminiPartsFromOutput(output)
|
||
response := map[string]any{
|
||
"candidates": []any{map[string]any{
|
||
"index": 0,
|
||
"content": map[string]any{
|
||
"role": "model",
|
||
"parts": parts,
|
||
},
|
||
"finishReason": "STOP",
|
||
}},
|
||
"modelVersion": model,
|
||
}
|
||
if usage := geminiUsageMetadataFromOutput(output); len(usage) > 0 {
|
||
response["usageMetadata"] = usage
|
||
}
|
||
return response
|
||
}
|
||
|
||
func geminiPartsFromOutput(output map[string]any) []any {
|
||
parts := []any{}
|
||
data, _ := output["data"].([]any)
|
||
for _, rawItem := range data {
|
||
item, _ := rawItem.(map[string]any)
|
||
if item == nil {
|
||
continue
|
||
}
|
||
if text := firstNonEmptyGeminiString(item["text"], item["content"]); text != "" {
|
||
parts = append(parts, map[string]any{"text": text})
|
||
continue
|
||
}
|
||
mimeType := firstNonEmptyGeminiString(item["mime_type"], item["mimeType"])
|
||
if b64 := strings.TrimSpace(stringFromRequestAny(item["b64_json"])); b64 != "" {
|
||
if parsed := parseGeminiDataURL(b64); parsed != nil {
|
||
mimeType = firstNonEmptyGeminiString(mimeType, parsed.MimeType)
|
||
b64 = parsed.Data
|
||
}
|
||
parts = append(parts, map[string]any{"inlineData": map[string]any{
|
||
"mimeType": firstNonEmptyGeminiString(mimeType, "image/png"),
|
||
"data": b64,
|
||
}})
|
||
continue
|
||
}
|
||
upload := mapFromGeminiAny(item["upload"])
|
||
if uri := firstNonEmptyGeminiString(item["url"], item["image_url"], upload["url"]); uri != "" {
|
||
parts = append(parts, map[string]any{"fileData": map[string]any{
|
||
"mimeType": firstNonEmptyGeminiString(mimeType, mimeTypeFromGeminiURI(uri)),
|
||
"fileUri": uri,
|
||
}})
|
||
}
|
||
}
|
||
if len(parts) == 0 {
|
||
if outputs, ok := output["output"].([]any); ok {
|
||
for _, raw := range outputs {
|
||
if uri := stringFromRequestAny(raw); uri != "" {
|
||
parts = append(parts, map[string]any{"fileData": map[string]any{
|
||
"mimeType": mimeTypeFromGeminiURI(uri),
|
||
"fileUri": uri,
|
||
}})
|
||
}
|
||
}
|
||
}
|
||
}
|
||
if len(parts) == 0 {
|
||
parts = append(parts, map[string]any{"text": firstNonEmptyGeminiString(output["message"])})
|
||
}
|
||
return parts
|
||
}
|
||
|
||
func geminiUsageMetadataFromOutput(output map[string]any) map[string]any {
|
||
usage := mapFromGeminiAny(output["usage"])
|
||
prompt := positiveIntFromGeminiAny(firstPresentGemini(usage["prompt_tokens"], usage["input_tokens"]))
|
||
candidates := positiveIntFromGeminiAny(firstPresentGemini(usage["completion_tokens"], usage["output_tokens"]))
|
||
total := positiveIntFromGeminiAny(usage["total_tokens"])
|
||
if total == 0 && (prompt > 0 || candidates > 0) {
|
||
total = prompt + candidates
|
||
}
|
||
meta := map[string]any{}
|
||
if prompt > 0 {
|
||
meta["promptTokenCount"] = prompt
|
||
}
|
||
if candidates > 0 {
|
||
meta["candidatesTokenCount"] = candidates
|
||
}
|
||
if total > 0 {
|
||
meta["totalTokenCount"] = total
|
||
}
|
||
return meta
|
||
}
|
||
|
||
// geminiFilesUpload godoc
|
||
// @Summary Gemini Files 上传接口
|
||
// @Description 使用统一 /api/v1 前缀启动或直接完成 Gemini Files 上传。
|
||
// @Tags gemini-compatible
|
||
// @Accept octet-stream
|
||
// @Produce json
|
||
// @Security BearerAuth
|
||
// @Param version path string true "Gemini 版本（v1 或 v1beta）"
|
||
// @Success 200 {object} map[string]interface{}
|
||
// @Failure 400 {object} ErrorEnvelope
|
||
// @Failure 401 {object} ErrorEnvelope
|
||
// @Router /api/v1/gemini/upload/{version}/files [post]
|
||
func (s *Server) geminiFilesUpload(w http.ResponseWriter, r *http.Request) {
|
||
if err := validateGeminiFilesVersion(r.PathValue("version")); err != nil {
|
||
writeProtocolError(w, clients.ProtocolGeminiGenerateContent, http.StatusBadRequest, err.Error(), nil, clients.ErrorCode(err))
|
||
return
|
||
}
|
||
if geminiUploadCommandHas(r, "start") {
|
||
s.startGeminiFilesUpload(w, r)
|
||
return
|
||
}
|
||
s.finalizeGeminiFilesUpload(w, r, "", geminiUploadSession{
|
||
Version: r.PathValue("version"),
|
||
DisplayName: "gemini-upload-" + newGeminiUploadID(),
|
||
MimeType: r.Header.Get("Content-Type"),
|
||
CreatedAt: time.Now(),
|
||
})
|
||
}
|
||
|
||
// geminiFilesUploadFinalize godoc
|
||
// @Summary 完成 Gemini Files 分段上传
|
||
// @Tags gemini-compatible
|
||
// @Accept octet-stream
|
||
// @Produce json
|
||
// @Security BearerAuth
|
||
// @Param version path string true "Gemini 版本（v1 或 v1beta）"
|
||
// @Param uploadID path string true "上传会话 ID"
|
||
// @Success 200 {object} map[string]interface{}
|
||
// @Failure 400 {object} ErrorEnvelope
|
||
// @Failure 401 {object} ErrorEnvelope
|
||
// @Router /api/v1/gemini/upload/{version}/files/{uploadID} [post]
|
||
func (s *Server) geminiFilesUploadFinalize(w http.ResponseWriter, r *http.Request) {
|
||
if err := validateGeminiFilesVersion(r.PathValue("version")); err != nil {
|
||
writeProtocolError(w, clients.ProtocolGeminiGenerateContent, http.StatusBadRequest, err.Error(), nil, clients.ErrorCode(err))
|
||
return
|
||
}
|
||
uploadID := strings.TrimSpace(r.PathValue("uploadID"))
|
||
value, ok := s.geminiUploadSessions.LoadAndDelete(uploadID)
|
||
if !ok {
|
||
writeProtocolError(w, clients.ProtocolGeminiGenerateContent, http.StatusBadRequest, "Gemini upload session not found", nil, "bad_request")
|
||
return
|
||
}
|
||
session, _ := value.(geminiUploadSession)
|
||
s.finalizeGeminiFilesUpload(w, r, uploadID, session)
|
||
}
|
||
|
||
func (s *Server) startGeminiFilesUpload(w http.ResponseWriter, r *http.Request) {
|
||
var body map[string]any
|
||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||
fileMeta := mapFromGeminiAny(body["file"])
|
||
uploadID := newGeminiUploadID()
|
||
session := geminiUploadSession{
|
||
Version: r.PathValue("version"),
|
||
DisplayName: firstNonEmptyGeminiString(fileMeta["display_name"], fileMeta["displayName"], "gemini-upload-"+uploadID),
|
||
MimeType: firstNonEmptyGeminiString(r.Header.Get("X-Goog-Upload-Header-Content-Type"), "application/octet-stream"),
|
||
SizeBytes: int64(positiveIntFromGeminiAny(r.Header.Get("X-Goog-Upload-Header-Content-Length"))),
|
||
CreatedAt: time.Now(),
|
||
}
|
||
s.geminiUploadSessions.Store(uploadID, session)
|
||
w.Header().Set("X-Goog-Upload-URL", absoluteRequestURL(r, geminiUploadPath(r, session.Version, uploadID)))
|
||
w.Header().Set("X-Goog-Upload-Status", "active")
|
||
writeJSON(w, http.StatusOK, map[string]any{})
|
||
}
|
||
|
||
func geminiUploadPath(r *http.Request, version string, uploadID string) string {
|
||
if strings.HasPrefix(r.URL.Path, "/api/v1/gemini/upload/") {
|
||
return "/api/v1/gemini/upload/" + version + "/files/" + uploadID
|
||
}
|
||
return "/upload/" + version + "/files/" + uploadID
|
||
}
|
||
|
||
func (s *Server) finalizeGeminiFilesUpload(w http.ResponseWriter, r *http.Request, uploadID string, session geminiUploadSession) {
|
||
if uploadID == "" {
|
||
uploadID = newGeminiUploadID()
|
||
}
|
||
r.Body = http.MaxBytesReader(w, r.Body, maxGatewayUploadBytes)
|
||
payload, err := io.ReadAll(r.Body)
|
||
if err != nil || len(payload) == 0 {
|
||
writeProtocolError(w, clients.ProtocolGeminiGenerateContent, http.StatusBadRequest, "file payload is required", nil, "bad_request")
|
||
return
|
||
}
|
||
mimeType := firstNonEmptyGeminiString(session.MimeType, http.DetectContentType(payload), "application/octet-stream")
|
||
upload, err := s.runner.UploadFile(r.Context(), runner.FileUploadPayload{
|
||
Bytes: payload,
|
||
ContentType: mimeType,
|
||
FileName: geminiUploadFileName(uploadID, session.DisplayName, mimeType),
|
||
Source: "gemini-files-api",
|
||
Scene: store.FileStorageSceneUpload,
|
||
})
|
||
if err != nil {
|
||
s.logger.Error("Gemini files upload failed", "error", err)
|
||
status := http.StatusBadGateway
|
||
if clients.ErrorCode(err) == "upload_no_channel" {
|
||
status = http.StatusServiceUnavailable
|
||
}
|
||
writeProtocolError(w, clients.ProtocolGeminiGenerateContent, status, err.Error(), nil, clients.ErrorCode(err))
|
||
return
|
||
}
|
||
sizeBytes := session.SizeBytes
|
||
if sizeBytes <= 0 {
|
||
sizeBytes = int64(len(payload))
|
||
}
|
||
writeJSON(w, http.StatusOK, geminiFileResponse("files/"+uploadID, stringFromRequestAny(upload["url"]), mimeType, sizeBytes))
|
||
}
|
||
|
||
func geminiFileResponse(name string, uri string, mimeType string, sizeBytes int64) map[string]any {
|
||
return map[string]any{"file": map[string]any{
|
||
"name": name,
|
||
"uri": uri,
|
||
"mimeType": firstNonEmptyGeminiString(mimeType, "application/octet-stream"),
|
||
"sizeBytes": strconv.FormatInt(sizeBytes, 10),
|
||
"state": "ACTIVE",
|
||
}}
|
||
}
|
||
|
||
func validateGeminiFilesVersion(version string) error {
|
||
if version == "v1" || version == "v1beta" {
|
||
return nil
|
||
}
|
||
return &clients.ClientError{Code: "bad_request", Message: "unsupported Gemini Files API version", Retryable: false}
|
||
}
|
||
|
||
func geminiUploadCommandHas(r *http.Request, command string) bool {
|
||
for _, item := range strings.Split(r.Header.Get("X-Goog-Upload-Command"), ",") {
|
||
if strings.EqualFold(strings.TrimSpace(item), command) {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
func newGeminiUploadID() string {
|
||
var payload [12]byte
|
||
if _, err := rand.Read(payload[:]); err != nil {
|
||
return strconv.FormatInt(time.Now().UnixNano(), 36)
|
||
}
|
||
return hex.EncodeToString(payload[:])
|
||
}
|
||
|
||
func geminiUploadFileName(uploadID string, displayName string, mimeType string) string {
|
||
name := strings.TrimSpace(strings.ReplaceAll(displayName, " ", ""))
|
||
name = path.Base(name)
|
||
if name == "." || name == "/" {
|
||
name = ""
|
||
}
|
||
if strings.Contains(name, ".") {
|
||
return name
|
||
}
|
||
if name == "" {
|
||
name = "gemini-upload-" + uploadID
|
||
}
|
||
return name + extensionFromGeminiMime(mimeType)
|
||
}
|
||
|
||
func extensionFromGeminiMime(mimeType string) string {
|
||
extensions, err := mime.ExtensionsByType(strings.ToLower(strings.TrimSpace(mimeType)))
|
||
if err == nil && len(extensions) > 0 {
|
||
return extensions[0]
|
||
}
|
||
return ".bin"
|
||
}
|
||
|
||
func absoluteRequestURL(r *http.Request, targetPath string) string {
|
||
proto := strings.TrimSpace(r.Header.Get("X-Forwarded-Proto"))
|
||
if proto == "" {
|
||
proto = "http"
|
||
}
|
||
host := strings.TrimSpace(r.Header.Get("X-Forwarded-Host"))
|
||
if host == "" {
|
||
host = r.Host
|
||
}
|
||
return fmt.Sprintf("%s://%s%s", proto, host, targetPath)
|
||
}
|
||
|
||
type parsedGeminiDataURL struct {
|
||
MimeType string
|
||
Data string
|
||
}
|
||
|
||
func parseGeminiDataURL(value string) *parsedGeminiDataURL {
|
||
if !strings.HasPrefix(strings.ToLower(value), "data:") {
|
||
return nil
|
||
}
|
||
prefix, data, ok := strings.Cut(value, ",")
|
||
if !ok || !strings.Contains(strings.ToLower(prefix), ";base64") {
|
||
return nil
|
||
}
|
||
mimeType := strings.TrimPrefix(strings.Split(prefix, ";")[0], "data:")
|
||
if mimeType == "" {
|
||
mimeType = "image/png"
|
||
}
|
||
return &parsedGeminiDataURL{MimeType: mimeType, Data: data}
|
||
}
|
||
|
||
func mimeTypeFromGeminiURI(value string) string {
|
||
extension := strings.ToLower(path.Ext(strings.Split(value, "?")[0]))
|
||
if extension == "" {
|
||
return "image/png"
|
||
}
|
||
if mimeType := mime.TypeByExtension(extension); mimeType != "" {
|
||
return mimeType
|
||
}
|
||
return "image/png"
|
||
}
|
||
|
||
func mapFromGeminiAny(value any) map[string]any {
|
||
if typed, ok := value.(map[string]any); ok {
|
||
return typed
|
||
}
|
||
return map[string]any{}
|
||
}
|
||
|
||
func firstPresentGemini(values ...any) any {
|
||
for _, value := range values {
|
||
if value != nil {
|
||
return value
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func firstNonEmptyGeminiString(values ...any) string {
|
||
for _, value := range values {
|
||
if text := strings.TrimSpace(stringFromRequestAny(value)); text != "" {
|
||
return text
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
func positiveIntFromGeminiAny(value any) int {
|
||
switch typed := value.(type) {
|
||
case int:
|
||
if typed > 0 {
|
||
return typed
|
||
}
|
||
case int64:
|
||
if typed > 0 {
|
||
return int(typed)
|
||
}
|
||
case float64:
|
||
if typed > 0 {
|
||
return int(typed)
|
||
}
|
||
case json.Number:
|
||
if parsed, err := typed.Int64(); err == nil && parsed > 0 {
|
||
return int(parsed)
|
||
}
|
||
case string:
|
||
if parsed, err := strconv.Atoi(strings.TrimSpace(typed)); err == nil && parsed > 0 {
|
||
return parsed
|
||
}
|
||
}
|
||
return 0
|
||
}
|