feat: add Gemini image compatibility
This commit is contained in:
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user