feat: add Gemini image compatibility

This commit is contained in:
wangbo 2026-06-24 00:42:49 +08:00
parent 7f32446466
commit c15842a94d
6 changed files with 966 additions and 15 deletions

View File

@ -97,6 +97,12 @@ func (a *Authenticator) Authenticate(r *http.Request) (*User, error) {
if token == "" {
token = strings.TrimSpace(r.Header.Get("x-comfy-api-key"))
}
if token == "" {
token = strings.TrimSpace(r.Header.Get("x-goog-api-key"))
}
if token == "" {
token = strings.TrimSpace(r.URL.Query().Get("key"))
}
if token == "" {
return nil, ErrUnauthorized
}

View File

@ -1026,6 +1026,130 @@ func TestGeminiClientChatConvertsMediaContentParts(t *testing.T) {
}
}
func TestGeminiClientImageGenerateBuildsNativeImageBody(t *testing.T) {
var captured map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := json.NewDecoder(r.Body).Decode(&captured); err != nil {
t.Fatalf("decode request: %v", err)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"candidates": []any{map[string]any{
"content": map[string]any{
"parts": []any{map[string]any{
"inlineData": map[string]any{
"mimeType": "image/png",
"data": "aW1hZ2U=",
},
}},
},
}},
})
}))
defer server.Close()
response, err := (GeminiClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "images.generations",
ModelType: "image_generate",
Model: "gemini-image",
Body: map[string]any{
"model": "gemini-image",
"prompt": "draw a cat",
"generationConfig": map[string]any{
"temperature": 0.5,
"imageConfig": map[string]any{
"aspectRatio": "16:9",
},
},
"resolution": "2K",
"n": 2,
},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
ProviderModelName: "gemini-image",
ModelType: "image_generate",
Credentials: map[string]any{"apiKey": "gemini-key"},
},
})
if err != nil {
t.Fatalf("run gemini image generate: %v", err)
}
config := captured["generationConfig"].(map[string]any)
modalities := config["responseModalities"].([]any)
imageConfig := config["imageConfig"].(map[string]any)
if modalities[0] != "IMAGE" || config["temperature"] != 0.5 || numericValue(config["candidateCount"], 0) != 2 || imageConfig["aspectRatio"] != "16:9" || imageConfig["imageSize"] != "2K" {
t.Fatalf("unexpected generationConfig: %+v", config)
}
contents := captured["contents"].([]any)
parts := contents[0].(map[string]any)["parts"].([]any)
if parts[0].(map[string]any)["text"] != "draw a cat" {
t.Fatalf("unexpected image contents: %+v", captured)
}
data := response.Result["data"].([]any)
if data[0].(map[string]any)["b64_json"] != "aW1hZ2U=" {
t.Fatalf("unexpected image response: %+v", response.Result)
}
}
func TestGeminiClientImageEditPreservesNativeContentsAndFileData(t *testing.T) {
var captured map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := json.NewDecoder(r.Body).Decode(&captured); err != nil {
t.Fatalf("decode request: %v", err)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"candidates": []any{map[string]any{
"content": map[string]any{
"parts": []any{map[string]any{
"fileData": map[string]any{
"mimeType": "image/webp",
"fileUri": "https://cdn.example/out.webp",
},
}},
},
}},
})
}))
defer server.Close()
_, err := (GeminiClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "images.edits",
ModelType: "image_edit",
Model: "gemini-image",
Body: map[string]any{
"model": "gemini-image",
"contents": []any{map[string]any{
"parts": []any{
map[string]any{"fileData": map[string]any{
"mimeType": "image/png",
"fileUri": "https://cdn.example/input.png",
}},
map[string]any{"text": "edit it"},
},
}},
"generationConfig": map[string]any{"temperature": 0.2},
},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
ProviderModelName: "gemini-image",
ModelType: "image_edit",
Credentials: map[string]any{"apiKey": "gemini-key"},
},
})
if err != nil {
t.Fatalf("run gemini image edit: %v", err)
}
config := captured["generationConfig"].(map[string]any)
if config["temperature"] != 0.2 {
t.Fatalf("generationConfig should preserve caller settings: %+v", config)
}
contents := captured["contents"].([]any)
parts := contents[0].(map[string]any)["parts"].([]any)
fileData := parts[0].(map[string]any)["fileData"].(map[string]any)
if fileData["fileUri"] != "https://cdn.example/input.png" {
t.Fatalf("native fileData should be preserved: %+v", captured)
}
}
func TestGeminiClientChatRestoresToolContext(t *testing.T) {
var captured map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

View File

@ -69,17 +69,20 @@ func geminiURL(baseURL string, model string, apiKey string) string {
}
func geminiBody(request Request) map[string]any {
if geminiImageModelType(request.ModelType) {
return geminiImageBody(request)
}
if contents, ok := request.Body["contents"]; ok {
return map[string]any{"contents": contents}
return geminiApplyRequestOptions(map[string]any{"contents": contents}, request, false)
}
prompt := firstNonEmptyPrompt(request.Body, "")
if prompt != "" {
return map[string]any{
return geminiApplyRequestOptions(map[string]any{
"contents": []any{map[string]any{
"role": "user",
"parts": []any{map[string]any{"text": prompt}},
}},
}
}, request, false)
}
body := map[string]any{"contents": geminiContentsFromMessages(request.Body)}
if tools := geminiToolsFromOpenAITools(request.Body["tools"]); len(tools) > 0 {
@ -87,12 +90,156 @@ func geminiBody(request Request) map[string]any {
}
contents, _ := body["contents"].([]any)
if len(contents) > 0 {
return body
return geminiApplyRequestOptions(body, request, false)
}
return map[string]any{"contents": []any{map[string]any{
return geminiApplyRequestOptions(map[string]any{"contents": []any{map[string]any{
"role": "user",
"parts": []any{map[string]any{"text": textFromMessages(request.Body)}},
}}}
}}}, request, false)
}
func geminiImageBody(request Request) map[string]any {
body := map[string]any{}
if contents, ok := request.Body["contents"]; ok {
body["contents"] = contents
} else {
parts := make([]any, 0)
for _, image := range geminiImageInputValues(request.Body) {
if media := geminiMediaURLPart(image.URI, image.MimeType, "image"); media != nil {
parts = append(parts, media)
}
}
if prompt := firstNonEmptyPrompt(request.Body, ""); prompt != "" {
parts = append(parts, map[string]any{"text": prompt})
}
body["contents"] = []any{map[string]any{
"role": "user",
"parts": parts,
}}
}
return geminiApplyRequestOptions(body, request, true)
}
type geminiImageInputValue struct {
URI string
MimeType string
}
func geminiImageInputValues(body map[string]any) []geminiImageInputValue {
values := make([]geminiImageInputValue, 0)
for _, key := range []string{"image", "images", "image_url", "imageUrl", "image_urls", "imageUrls"} {
values = append(values, geminiImageInputValuesFromAny(body[key])...)
}
return values
}
func geminiImageInputValuesFromAny(value any) []geminiImageInputValue {
switch typed := value.(type) {
case string:
if uri := strings.TrimSpace(typed); uri != "" {
return []geminiImageInputValue{{URI: uri}}
}
case []any:
out := make([]geminiImageInputValue, 0, len(typed))
for _, item := range typed {
out = append(out, geminiImageInputValuesFromAny(item)...)
}
return out
case []string:
out := make([]geminiImageInputValue, 0, len(typed))
for _, item := range typed {
if uri := strings.TrimSpace(item); uri != "" {
out = append(out, geminiImageInputValue{URI: uri})
}
}
return out
case map[string]any:
if uri := firstNonEmptyString(typed["url"], typed["fileUri"], typed["file_uri"]); uri != "" {
return []geminiImageInputValue{{
URI: uri,
MimeType: firstNonEmptyString(typed["mime_type"], typed["mimeType"]),
}}
}
}
return nil
}
func geminiApplyRequestOptions(body map[string]any, request Request, imageResponse bool) map[string]any {
if generationConfig := geminiGenerationConfig(request.Body, imageResponse); len(generationConfig) > 0 {
body["generationConfig"] = generationConfig
}
for _, item := range []struct {
from []string
to string
}{
{from: []string{"systemInstruction", "system_instruction"}, to: "systemInstruction"},
{from: []string{"safetySettings", "safety_settings"}, to: "safetySettings"},
{from: []string{"toolConfig", "tool_config"}, to: "toolConfig"},
{from: []string{"cachedContent", "cached_content"}, to: "cachedContent"},
} {
for _, key := range item.from {
if value, ok := request.Body[key]; ok {
body[item.to] = value
break
}
}
}
if _, exists := body["tools"]; !exists {
if tools, ok := request.Body["tools"]; ok {
body["tools"] = tools
}
}
return body
}
func geminiGenerationConfig(body map[string]any, imageResponse bool) map[string]any {
source := mapFromAny(firstPresent(body["generationConfig"], body["generation_config"]))
out := cloneMapAny(source)
if aspectRatio := firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"]); aspectRatio != "" {
imageConfig := cloneMapAny(mapFromAny(firstPresent(out["imageConfig"], out["image_config"])))
imageConfig["aspectRatio"] = aspectRatio
out["imageConfig"] = imageConfig
delete(out, "image_config")
}
if imageSize := firstNonEmptyString(body["resolution"], body["imageSize"], body["image_size"], body["size"]); imageSize != "" {
imageConfig := cloneMapAny(mapFromAny(firstPresent(out["imageConfig"], out["image_config"])))
imageConfig["imageSize"] = imageSize
out["imageConfig"] = imageConfig
delete(out, "image_config")
}
if count := intFromAny(firstPresent(body["n"], body["candidateCount"], body["candidate_count"])); count > 0 {
out["candidateCount"] = count
}
if imageResponse {
out["responseModalities"] = ensureGeminiResponseModalityImage(out["responseModalities"])
}
return out
}
func ensureGeminiResponseModalityImage(value any) []any {
items := []any{}
switch typed := value.(type) {
case []any:
items = append(items, typed...)
case []string:
for _, item := range typed {
items = append(items, item)
}
case string:
if strings.TrimSpace(typed) != "" {
items = append(items, typed)
}
}
for _, item := range items {
if strings.EqualFold(strings.TrimSpace(stringFromAny(item)), "IMAGE") {
return items
}
}
return append(items, "IMAGE")
}
func geminiImageModelType(modelType string) bool {
return modelType == "image_generate" || modelType == "image_edit"
}
func geminiContentsFromMessages(body map[string]any) []any {
@ -387,7 +534,7 @@ func geminiToolsFromOpenAITools(value any) []any {
}
func geminiResult(request Request, raw map[string]any) map[string]any {
if request.ModelType == "image" {
if geminiImageModelType(request.ModelType) {
data := geminiImageData(raw)
if len(data) == 0 {
data = []any{map[string]any{"url": "/static/provider/gemini-image-placeholder.png"}}
@ -507,7 +654,21 @@ func geminiImageData(raw map[string]any) []any {
inline, _ = partMap["inline_data"].(map[string]any)
}
if data, ok := inline["data"].(string); ok && data != "" {
out = append(out, map[string]any{"b64_json": data, "mime_type": inline["mimeType"]})
out = append(out, map[string]any{
"b64_json": data,
"mime_type": firstNonEmptyString(inline["mimeType"], inline["mime_type"]),
})
continue
}
fileData, _ := partMap["fileData"].(map[string]any)
if fileData == nil {
fileData, _ = partMap["file_data"].(map[string]any)
}
if uri := firstNonEmptyString(fileData["fileUri"], fileData["file_uri"], fileData["uri"]); uri != "" {
out = append(out, map[string]any{
"url": uri,
"mime_type": firstNonEmptyString(fileData["mimeType"], fileData["mime_type"], mimeFromURI(uri)),
})
}
}
}

View File

@ -0,0 +1,566 @@
package httpapi
import (
"crypto/rand"
"encoding/hex"
"encoding/json"
"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
}
func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok {
writeError(w, http.StatusUnauthorized, "unauthorized")
return
}
var native map[string]any
if err := json.NewDecoder(r.Body).Decode(&native); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body", "invalid_json_body")
return
}
mapping, err := geminiImageTaskBody(r.PathValue("model"), native)
if err != nil {
writeError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
return
}
if !apiKeyScopeAllowed(user, mapping.Kind) {
writeError(w, http.StatusForbidden, "api key scope does not allow this capability")
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
}
writeError(w, status, err.Error(), clients.ErrorCode(err))
return
}
task, err := s.store.CreateTask(r.Context(), store.CreateTaskInput{
Kind: mapping.Kind,
Model: mapping.Model,
RunMode: runModeFromRequest(prepared.Body),
Async: false,
Request: prepared.Body,
ConversationID: prepared.ConversationID,
NewMessageCount: prepared.NewMessageCount,
MessageRefs: prepared.MessageRefs,
}, user)
if err != nil {
s.logger.Error("create gemini task failed", "kind", mapping.Kind, "error", err)
writeError(w, http.StatusInternalServerError, "create task failed")
return
}
runCtx, cancelRun := s.requestExecutionContext(r)
defer cancelRun()
result, runErr := s.runner.Execute(runCtx, task, user)
if runErr != nil {
if !requestStillConnected(r) {
return
}
writeErrorWithDetails(w, statusFromRunError(runErr), runErrorMessage(runErr), runErrorDetails(runErr), runErrorCode(runErr))
return
}
if !requestStillConnected(r) {
return
}
writeJSON(w, http.StatusOK, geminiGenerateContentResponse(result.Output, mapping.Model))
}
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
}
func (s *Server) geminiFilesUpload(w http.ResponseWriter, r *http.Request) {
if err := validateGeminiFilesVersion(r.PathValue("version")); err != nil {
writeError(w, http.StatusBadRequest, err.Error(), 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(),
})
}
func (s *Server) geminiFilesUploadFinalize(w http.ResponseWriter, r *http.Request) {
if err := validateGeminiFilesVersion(r.PathValue("version")); err != nil {
writeError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
return
}
uploadID := strings.TrimSpace(r.PathValue("uploadID"))
value, ok := s.geminiUploadSessions.LoadAndDelete(uploadID)
if !ok {
writeError(w, http.StatusBadRequest, "Gemini upload session not found", "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, "/upload/"+session.Version+"/files/"+uploadID))
w.Header().Set("X-Goog-Upload-Status", "active")
writeJSON(w, http.StatusOK, map[string]any{})
}
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 {
writeError(w, http.StatusBadRequest, "file payload is required", "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
}
writeError(w, status, err.Error(), 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
}

View File

@ -0,0 +1,87 @@
package httpapi
import "testing"
func TestGeminiImageTaskBodyMapsTextOnlyToImageGenerate(t *testing.T) {
mapping, err := geminiImageTaskBody("gemini-image", map[string]any{
"contents": []any{
map[string]any{"parts": []any{
map[string]any{"text": "draw"},
map[string]any{"text": "a cat"},
}},
},
"generationConfig": map[string]any{
"candidateCount": float64(2),
"imageConfig": map[string]any{
"aspectRatio": "16:9",
"imageSize": "2K",
},
},
})
if err != nil {
t.Fatalf("map Gemini request: %v", err)
}
if mapping.Kind != "images.generations" || mapping.Body["modelType"] != "image_generate" {
t.Fatalf("unexpected mapping kind/body: %+v", mapping)
}
if mapping.Body["prompt"] != "draw\na cat" || mapping.Body["n"] != 2 || mapping.Body["aspect_ratio"] != "16:9" || mapping.Body["resolution"] != "2K" {
t.Fatalf("unexpected mapped body: %+v", mapping.Body)
}
}
func TestGeminiImageTaskBodyMapsInlineDataToImageEdit(t *testing.T) {
mapping, err := geminiImageTaskBody("gemini-image", map[string]any{
"contents": []any{
map[string]any{"parts": []any{
map[string]any{"text": "make it brighter"},
map[string]any{"inlineData": map[string]any{
"mimeType": "image/png",
"data": "aW1hZ2U=",
}},
}},
},
})
if err != nil {
t.Fatalf("map Gemini request: %v", err)
}
if mapping.Kind != "images.edits" || mapping.Body["modelType"] != "image_edit" {
t.Fatalf("unexpected mapping kind/body: %+v", mapping)
}
if mapping.Body["prompt"] != "make it brighter" || mapping.Body["image"] != "data:image/png;base64,aW1hZ2U=" {
t.Fatalf("unexpected mapped body: %+v", mapping.Body)
}
}
func TestGeminiGenerateContentResponseMapsURLAndUsage(t *testing.T) {
response := geminiGenerateContentResponse(map[string]any{
"data": []any{map[string]any{
"type": "image",
"url": "https://cdn.example/out.webp",
}},
"usage": map[string]any{
"prompt_tokens": 3,
"completion_tokens": 4,
"total_tokens": 7,
},
}, "gemini-image")
candidates := response["candidates"].([]any)
content := candidates[0].(map[string]any)["content"].(map[string]any)
parts := content["parts"].([]any)
fileData := parts[0].(map[string]any)["fileData"].(map[string]any)
if fileData["fileUri"] != "https://cdn.example/out.webp" || fileData["mimeType"] != "image/webp" {
t.Fatalf("unexpected fileData part: %+v", fileData)
}
usage := response["usageMetadata"].(map[string]any)
if usage["promptTokenCount"] != 3 || usage["candidatesTokenCount"] != 4 || usage["totalTokenCount"] != 7 {
t.Fatalf("unexpected usage metadata: %+v", usage)
}
}
func TestGeminiFileResponseShape(t *testing.T) {
response := geminiFileResponse("files/test", "https://cdn.example/input.png", "image/png", 12)
file := response["file"].(map[string]any)
if file["name"] != "files/test" || file["uri"] != "https://cdn.example/input.png" || file["mimeType"] != "image/png" || file["sizeBytes"] != "12" || file["state"] != "ACTIVE" {
t.Fatalf("unexpected file response: %+v", file)
}
}

View File

@ -5,6 +5,7 @@ import (
"log/slog"
"net/http"
"strings"
"sync"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
@ -13,12 +14,13 @@ import (
)
type Server struct {
ctx context.Context
cfg config.Config
store *store.Store
auth *auth.Authenticator
runner *runner.Service
logger *slog.Logger
ctx context.Context
cfg config.Config
store *store.Store
auth *auth.Authenticator
runner *runner.Service
logger *slog.Logger
geminiUploadSessions sync.Map
}
func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Handler {
@ -147,6 +149,11 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
mux.Handle("GET /api/v1/voice_clone/voices", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.listClonedVoices)))
mux.Handle("DELETE /api/v1/voice_clone/voices/{voiceID}", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.deleteClonedVoice)))
mux.Handle("POST /api/v1/files/upload", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.uploadFile)))
mux.Handle("POST /v1beta/models/{model}:generateContent", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.geminiGenerateContent)))
mux.Handle("POST /v1/models/{model}:generateContent", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.geminiGenerateContent)))
mux.Handle("POST /models/{model}:generateContent", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.geminiGenerateContent)))
mux.Handle("POST /upload/{version}/files", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.geminiFilesUpload)))
mux.Handle("POST /upload/{version}/files/{uploadID}", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.geminiFilesUploadFinalize)))
mux.Handle("GET /api/v1/tasks", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.listTasks)))
mux.Handle("GET /api/v1/tasks/{taskID}", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.getTask)))
mux.Handle("POST /api/v1/tasks/{taskID}/cancel", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.cancelTask)))
@ -205,7 +212,7 @@ func (s *Server) cors(next http.Handler) http.Handler {
w.Header().Set("Access-Control-Allow-Origin", origin)
w.Header().Set("Vary", "Origin")
w.Header().Set("Access-Control-Allow-Credentials", "true")
w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type, X-Comfy-Api-Key, X-Async, X-EasyAI-Conversation-ID")
w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type, X-Comfy-Api-Key, X-Goog-Api-Key, X-Goog-Upload-Protocol, X-Goog-Upload-Command, X-Goog-Upload-Header-Content-Length, X-Goog-Upload-Header-Content-Type, X-Goog-Upload-Offset, X-Async, X-EasyAI-Conversation-ID")
w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE, OPTIONS")
}
if r.Method == http.MethodOptions {