diff --git a/apps/api/internal/auth/auth.go b/apps/api/internal/auth/auth.go index db90018..09877f2 100644 --- a/apps/api/internal/auth/auth.go +++ b/apps/api/internal/auth/auth.go @@ -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 } diff --git a/apps/api/internal/clients/clients_test.go b/apps/api/internal/clients/clients_test.go index 5429115..74456f1 100644 --- a/apps/api/internal/clients/clients_test.go +++ b/apps/api/internal/clients/clients_test.go @@ -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) { diff --git a/apps/api/internal/clients/gemini.go b/apps/api/internal/clients/gemini.go index 292d29b..3c70d12 100644 --- a/apps/api/internal/clients/gemini.go +++ b/apps/api/internal/clients/gemini.go @@ -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)), + }) } } } diff --git a/apps/api/internal/httpapi/gemini_compat.go b/apps/api/internal/httpapi/gemini_compat.go new file mode 100644 index 0000000..e47e932 --- /dev/null +++ b/apps/api/internal/httpapi/gemini_compat.go @@ -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 +} diff --git a/apps/api/internal/httpapi/gemini_compat_test.go b/apps/api/internal/httpapi/gemini_compat_test.go new file mode 100644 index 0000000..753d43b --- /dev/null +++ b/apps/api/internal/httpapi/gemini_compat_test.go @@ -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) + } +} diff --git a/apps/api/internal/httpapi/server.go b/apps/api/internal/httpapi/server.go index 88df9da..33b6df2 100644 --- a/apps/api/internal/httpapi/server.go +++ b/apps/api/internal/httpapi/server.go @@ -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 {