Skip to content

feat: imagen for vertex channel #1614

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
creamlike1024 merged 1 commit into from
Aug 18, 2025
Merged

feat: imagen for vertex channel #1614

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 15 additions & 8 deletions relay/channel/gemini/adaptor.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,15 +59,22 @@ func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInf
return nil, errors.New("not supported model for image generation")
}

// convert size to aspect ratio
// convert size to aspect ratio but allow user to specify aspect ratio
aspectRatio := "1:1" // default aspect ratio
switch request.Size {
case "1024x1024":
aspectRatio = "1:1"
case "1024x1792":
aspectRatio = "9:16"
case "1792x1024":
aspectRatio = "16:9"
size := strings.TrimSpace(request.Size)
if size != "" {
if strings.Contains(size, ":") {
aspectRatio = size
} else {
switch size {
case "1024x1024":
aspectRatio = "1:1"
case "1024x1792":
aspectRatio = "9:16"
case "1792x1024":
aspectRatio = "16:9"
}
}
}

// build gemini imagen request
Expand Down
58 changes: 56 additions & 2 deletions relay/channel/vertex/adaptor.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,8 +66,8 @@ func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInf
}

func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
//TODO implement me
return nil, errors.New("not implemented")
geminiAdaptor := gemini.Adaptor{}
return geminiAdaptor.ConvertImageRequest(c, info, request)
}

func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
Expand Down Expand Up @@ -181,6 +181,60 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn
if request == nil {
return nil, errors.New("request is nil")
}
if a.RequestMode == RequestModeGemini && strings.HasPrefix(info.UpstreamModelName, "imagen") {
prompt := ""
for _, m := range request.Messages {
if m.Role == "user" {
prompt = m.StringContent()
if prompt != "" {
break
}
}
}
if prompt == "" {
if p, ok := request.Prompt.(string); ok {
prompt = p
}
}
if prompt == "" {
return nil, errors.New("prompt is required for image generation")
}

imgReq := dto.ImageRequest{
Model: request.Model,
Prompt: prompt,
N: 1,
Size: "1024x1024",
}
if request.N > 0 {
imgReq.N = uint(request.N)
}
if request.Size != "" {
imgReq.Size = request.Size
}
if len(request.ExtraBody) > 0 {
var extra map[string]any
if err := json.Unmarshal(request.ExtraBody, &extra); err == nil {
if n, ok := extra["n"].(float64); ok {
imgReq.N = uint(n)
}
if size, ok := extra["size"].(string); ok {
imgReq.Size = size
}
// accept aspectRatio in extra body (top-level or under parameters)
if ar, ok := extra["aspectRatio"].(string); ok && ar != "" {
imgReq.Size = ar
}
if params, ok := extra["parameters"].(map[string]any); ok {
if ar, ok := params["aspectRatio"].(string); ok && ar != "" {
imgReq.Size = ar
}
}
}
}
c.Set("request_model", request.Model)
return a.ConvertImageRequest(c, info, imgReq)
}
Comment on lines +184 to +237

@coderabbitai coderabbitai Bot Aug 18, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Imagen request translation: add guardrails for n, force non-streaming, and set request_model to upstream

The flow is solid. A few correctness/UX tweaks will make it safer:

  • Guard against negative/zero n in ExtraBody to avoid uint underflow or invalid counts.
  • Force non-streaming for image requests to prevent the streaming response path from mis-routing to text SSE handlers.
  • Set request_model to the actual upstream model name for accurate logging/metrics.

Apply this diff:

-	if a.RequestMode == RequestModeGemini && strings.HasPrefix(info.UpstreamModelName, "imagen") {
+	if a.RequestMode == RequestModeGemini && strings.HasPrefix(info.UpstreamModelName, "imagen") {
 		prompt := ""
 		for _, m := range request.Messages {
 			if m.Role == "user" {
 				prompt = m.StringContent()
 				if prompt != "" {
 					break
 				}
 			}
 		}
 		if prompt == "" {
 			if p, ok := request.Prompt.(string); ok {
 				prompt = p
 			}
 		}
 		if prompt == "" {
 			return nil, errors.New("prompt is required for image generation")
 		}

 		imgReq := dto.ImageRequest{
 			Model:  request.Model,
 			Prompt: prompt,
 			N:      1,
 			Size:   "1024x1024",
 		}
 		if request.N > 0 {
 			imgReq.N = uint(request.N)
 		}
 		if request.Size != "" {
 			imgReq.Size = request.Size
 		}
 		if len(request.ExtraBody) > 0 {
 			var extra map[string]any
 			if err := json.Unmarshal(request.ExtraBody, &extra); err == nil {
 				if n, ok := extra["n"].(float64); ok {
-					imgReq.N = uint(n)
+					if n > 0 {
+						imgReq.N = uint(n)
+					}
 				}
 				if size, ok := extra["size"].(string); ok {
 					imgReq.Size = size
 				}
 				// accept aspectRatio in extra body (top-level or under parameters)
 				if ar, ok := extra["aspectRatio"].(string); ok && ar != "" {
 					imgReq.Size = ar
 				}
 				if params, ok := extra["parameters"].(map[string]any); ok {
 					if ar, ok := params["aspectRatio"].(string); ok && ar != "" {
 						imgReq.Size = ar
 					}
 				}
 			}
 		}
-		c.Set("request_model", request.Model)
+		// Images are non-streaming; ensure downstream picks non-SSE handlers.
+		info.IsStream = false
+		// For logging/metrics, reflect the actual upstream model.
+		c.Set("request_model", info.UpstreamModelName)
 		return a.ConvertImageRequest(c, info, imgReq)
 	}

Notes:

  • If Imagen has a known max sampleCount (e.g., 1–8), consider clamping imgReq.N to that range.
  • If you intend to support personGeneration overrides, we can thread an optional field through dto.ImageRequest and forward it to the Gemini adaptor.

To ensure we don’t have duplicate Imagen translation blocks, run:

I can follow up with clamping logic once we confirm Imagen’s allowed sampleCount.

🏁 Script executed:

#!/bin/bash
set -euo pipefail
# Look for a second image-generation block in this file
rg -n -C2 -e 'prompt is required for image generation' -e 'aspectRatio' relay/channel/vertex/adaptor.go

Length of output: 640

Apply safety enhancements in the Gemini→Imagen translation block

We’ve confirmed there’s only one image‐generation block in relay/channel/vertex/adaptor.go. Please incorporate these critical fixes within the RequestModeGemini && strings.HasPrefix(info.UpstreamModelName, "imagen") branch (lines ~184–237):

  • Guard against negative or zero n from ExtraBody to prevent uint underflow.
  • Force non-streaming (info.IsStream = false) so image requests don’t hit text SSE handlers.
  • Record the actual upstream model name in the context for accurate logging/metrics.
--- a/relay/channel/vertex/adaptor.go
+++ b/relay/channel/vertex/adaptor.go
@@ -220,7 +220,9 @@ if len(request.ExtraBody) > 0 {
                     if err := json.Unmarshal(request.ExtraBody, &extra); err == nil {
-                        if n, ok := extra["n"].(float64); ok {
-                            imgReq.N = uint(n)
+                        if n, ok := extra["n"].(float64); ok && n > 0 {
+                            imgReq.N = uint(n)
+                        }
                     }
@@ -234,7 +236,11 @@ if len(request.ExtraBody) > 0 {
                     }
                 }
             }
-        c.Set("request_model", request.Model)
+        // Images are non-streaming.
+        info.IsStream = false
+        // Use the real upstream model for metrics.
+        c.Set("request_model", info.UpstreamModelName)
         return a.ConvertImageRequest(c, info, imgReq)
     }

Note: you may also clamp imgReq.N to Imagen’s supported range (e.g. 1–8) once confirmed.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if a.RequestMode == RequestModeGemini && strings.HasPrefix(info.UpstreamModelName, "imagen") {
prompt := ""
for _, m := range request.Messages {
if m.Role == "user" {
prompt = m.StringContent()
if prompt != "" {
break
}
}
}
if prompt == "" {
if p, ok := request.Prompt.(string); ok {
prompt = p
}
}
if prompt == "" {
return nil, errors.New("prompt is required for image generation")
}
imgReq := dto.ImageRequest{
Model: request.Model,
Prompt: prompt,
N: 1,
Size: "1024x1024",
}
if request.N > 0 {
imgReq.N = uint(request.N)
}
if request.Size != "" {
imgReq.Size = request.Size
}
if len(request.ExtraBody) > 0 {
var extra map[string]any
if err := json.Unmarshal(request.ExtraBody, &extra); err == nil {
if n, ok := extra["n"].(float64); ok {
imgReq.N = uint(n)
}
if size, ok := extra["size"].(string); ok {
imgReq.Size = size
}
// accept aspectRatio in extra body (top-level or under parameters)
if ar, ok := extra["aspectRatio"].(string); ok && ar != "" {
imgReq.Size = ar
}
if params, ok := extra["parameters"].(map[string]any); ok {
if ar, ok := params["aspectRatio"].(string); ok && ar != "" {
imgReq.Size = ar
}
}
}
}
c.Set("request_model", request.Model)
return a.ConvertImageRequest(c, info, imgReq)
}
if a.RequestMode == RequestModeGemini && strings.HasPrefix(info.UpstreamModelName, "imagen") {
prompt := ""
for _, m := range request.Messages {
if m.Role == "user" {
prompt = m.StringContent()
if prompt != "" {
break
}
}
}
if prompt == "" {
if p, ok := request.Prompt.(string); ok {
prompt = p
}
}
if prompt == "" {
return nil, errors.New("prompt is required for image generation")
}
imgReq := dto.ImageRequest{
Model: request.Model,
Prompt: prompt,
N: 1,
Size: "1024x1024",
}
if request.N > 0 {
imgReq.N = uint(request.N)
}
if request.Size != "" {
imgReq.Size = request.Size
}
if len(request.ExtraBody) > 0 {
var extra map[string]any
if err := json.Unmarshal(request.ExtraBody, &extra); err == nil {
if n, ok := extra["n"].(float64); ok && n > 0 {
imgReq.N = uint(n)
}
if size, ok := extra["size"].(string); ok {
imgReq.Size = size
}
// accept aspectRatio in extra body (top-level or under parameters)
if ar, ok := extra["aspectRatio"].(string); ok && ar != "" {
imgReq.Size = ar
}
if params, ok := extra["parameters"].(map[string]any); ok {
if ar, ok := params["aspectRatio"].(string); ok && ar != "" {
imgReq.Size = ar
}
}
}
}
// Images are non-streaming.
info.IsStream = false
// Use the real upstream model for metrics.
c.Set("request_model", info.UpstreamModelName)
return a.ConvertImageRequest(c, info, imgReq)
}
🤖 Prompt for AI Agents 
In relay/channel/vertex/adaptor.go around lines 184 to 237, inside the
RequestModeGemini && strings.HasPrefix(info.UpstreamModelName, "imagen") branch:
ensure safety by 1) when reading "n" from ExtraBody guard against zero or
negative values before converting to uint (ignore or default to 1 if n <= 0) to
avoid underflow and optionally clamp imgReq.N to a safe range (e.g. 1–8) if
desired; 2) force non-streaming by setting info.IsStream = false so image
requests do not flow into text SSE handlers; and 3) store the actual upstream
model name into the context (e.g. c.Set("upstream_model",
info.UpstreamModelName)) for correct logging/metrics before calling
ConvertImageRequest.

if a.RequestMode == RequestModeClaude {
claudeReq, err := claude.RequestOpenAI2ClaudeMessage(c, *request)
if err != nil {
Expand Down