diff --git a/apps/api/internal/httpapi/gemini_compat.go b/apps/api/internal/httpapi/gemini_compat.go index e47e932..aac8b4b 100644 --- a/apps/api/internal/httpapi/gemini_compat.go +++ b/apps/api/internal/httpapi/gemini_compat.go @@ -38,6 +38,42 @@ type geminiUploadSession struct { CreatedAt time.Time } +var geminiGenerateContentRoutePrefixes = []string{ + "/v1beta/models/", + "/v1/models/", + "/models/", +} + +func (s *Server) registerGeminiGenerateContentRoutes(mux *http.ServeMux) { + handler := s.auth.Require(auth.PermissionBasic, http.HandlerFunc(s.geminiGenerateContent)) + for _, prefix := range geminiGenerateContentRoutePrefixes { + mux.Handle("POST "+prefix, geminiGenerateContentRouteHandler(prefix, handler)) + } +} + +func geminiGenerateContentRouteHandler(prefix string, next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + model, ok := geminiGenerateContentModelFromPath(prefix, r.URL.Path) + if !ok { + http.NotFound(w, r) + return + } + r.SetPathValue("model", model) + next.ServeHTTP(w, r) + }) +} + +func geminiGenerateContentModelFromPath(prefix string, requestPath string) (string, bool) { + if !strings.HasPrefix(requestPath, prefix) { + return "", false + } + model, ok := strings.CutSuffix(strings.TrimPrefix(requestPath, prefix), ":generateContent") + if !ok || strings.TrimSpace(model) == "" || strings.Contains(model, "/") { + return "", false + } + return model, true +} + func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) { user, ok := auth.UserFromContext(r.Context()) if !ok { diff --git a/apps/api/internal/httpapi/gemini_compat_test.go b/apps/api/internal/httpapi/gemini_compat_test.go index 753d43b..d7e5c79 100644 --- a/apps/api/internal/httpapi/gemini_compat_test.go +++ b/apps/api/internal/httpapi/gemini_compat_test.go @@ -2,6 +2,65 @@ package httpapi import "testing" +func TestGeminiGenerateContentModelFromPath(t *testing.T) { + tests := []struct { + name string + prefix string + requestPath string + wantModel string + wantOK bool + }{ + { + name: "v1beta model", + prefix: "/v1beta/models/", + requestPath: "/v1beta/models/gemini-2.5-flash:generateContent", + wantModel: "gemini-2.5-flash", + wantOK: true, + }, + { + name: "v1 model", + prefix: "/v1/models/", + requestPath: "/v1/models/gemini-compatible:generateContent", + wantModel: "gemini-compatible", + wantOK: true, + }, + { + name: "bare model path", + prefix: "/models/", + requestPath: "/models/gemini-image:generateContent", + wantModel: "gemini-image", + wantOK: true, + }, + { + name: "wrong suffix", + prefix: "/v1beta/models/", + requestPath: "/v1beta/models/gemini-2.5-flash:countTokens", + wantOK: false, + }, + { + name: "missing model", + prefix: "/v1beta/models/", + requestPath: "/v1beta/models/:generateContent", + wantOK: false, + }, + { + name: "extra path segment", + prefix: "/v1beta/models/", + requestPath: "/v1beta/models/group/gemini-2.5-flash:generateContent", + wantOK: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotModel, gotOK := geminiGenerateContentModelFromPath(tt.prefix, tt.requestPath) + if gotOK != tt.wantOK || gotModel != tt.wantModel { + t.Fatalf("geminiGenerateContentModelFromPath() = %q, %v; want %q, %v", gotModel, gotOK, tt.wantModel, tt.wantOK) + } + }) + } +} + func TestGeminiImageTaskBodyMapsTextOnlyToImageGenerate(t *testing.T) { mapping, err := geminiImageTaskBody("gemini-image", map[string]any{ "contents": []any{ diff --git a/apps/api/internal/httpapi/server.go b/apps/api/internal/httpapi/server.go index 33b6df2..124a571 100644 --- a/apps/api/internal/httpapi/server.go +++ b/apps/api/internal/httpapi/server.go @@ -149,9 +149,7 @@ 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))) + server.registerGeminiGenerateContentRoutes(mux) 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)))