feat: add image generation for OpenAI DALL-E and Gemini Imagen
New `/v1/images/generations` endpoint proxies DALL-E 2/3 (OpenAI)
and Imagen 3 (Gemini). Same auth/logging as chat completions.
- Add ImageGenerationRequest/Response models
- Extend Provider interface with ImageGeneration()
- OpenAI: forward to /v1/images/generations
- Gemini: call /v1beta/models/{model}:predict, map OpenAI params
- Circuit breaker wraps image gen like chat completions
- Model routing: dall-e* -> openai, imagen*/gemini* -> gemini
- Unsupported providers (deepseek/moonshot/grok/ollama) return error
- Fix pre-existing CachedContentTokenCount bug in StreamGemini
This commit is contained in:
@@ -186,6 +186,7 @@ func (s *Server) setupRoutes() {
|
||||
v1.Use(middleware.AuthMiddleware(s.database, true))
|
||||
{
|
||||
v1.POST("/chat/completions", s.handleChatCompletions)
|
||||
v1.POST("/images/generations", s.handleImageGenerations)
|
||||
v1.GET("/models", s.handleListModels)
|
||||
v1.GET("/responses", s.handleListResponses)
|
||||
}
|
||||
@@ -501,6 +502,71 @@ func (s *Server) handleChatCompletions(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
|
||||
func (s *Server) handleImageGenerations(c *gin.Context) {
|
||||
startTime := time.Now()
|
||||
var req models.ImageGenerationRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// Determine provider based on model name
|
||||
providerName := "openai"
|
||||
modelLower := strings.ToLower(req.Model)
|
||||
switch {
|
||||
case strings.Contains(modelLower, "imagen"), strings.Contains(modelLower, "gemini"):
|
||||
providerName = "gemini"
|
||||
case strings.Contains(modelLower, "dall"), strings.HasPrefix(modelLower, "openai/"):
|
||||
providerName = "openai"
|
||||
}
|
||||
|
||||
// Default model for each provider if not specified
|
||||
if req.Model == "" {
|
||||
if providerName == "openai" {
|
||||
req.Model = "dall-e-3"
|
||||
} else {
|
||||
req.Model = "imagen-3.0-generate-001"
|
||||
}
|
||||
}
|
||||
|
||||
// Strip common prefixes
|
||||
prefixes := []string{"openai/", "gemini/", "google/"}
|
||||
for _, p := range prefixes {
|
||||
if strings.HasPrefix(req.Model, p) {
|
||||
req.Model = strings.TrimPrefix(req.Model, p)
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
provider, ok := s.providers[providerName]
|
||||
if !ok {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("Provider %s not enabled or supported", providerName)})
|
||||
return
|
||||
}
|
||||
|
||||
clientID := "default"
|
||||
if auth, ok := c.Get("auth"); ok {
|
||||
if authInfo, ok := auth.(models.AuthInfo); ok {
|
||||
clientID = authInfo.ClientID
|
||||
}
|
||||
}
|
||||
|
||||
resp, err := provider.ImageGeneration(c.Request.Context(), &req)
|
||||
if err != nil {
|
||||
s.logRequest(startTime, clientID, providerName, req.Model, nil, err, false)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
s.logRequest(startTime, clientID, providerName, req.Model, &models.Usage{
|
||||
PromptTokens: 1,
|
||||
CompletionTokens: uint32(len(resp.Data)),
|
||||
TotalTokens: 1 + uint32(len(resp.Data)),
|
||||
}, nil, false)
|
||||
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
|
||||
func (s *Server) logRequest(start time.Time, clientID, provider, model string, usage *models.Usage, err error, hasImages bool) {
|
||||
entry := RequestLog{
|
||||
Timestamp: start,
|
||||
|
||||
Reference in New Issue
Block a user