feat: imagen for vertex channel - #1614
Merged
Merged
Conversation
Contributor
WalkthroughAdds explicit aspect-ratio handling in Gemini image conversion and wires Vertex adaptor to route OpenAI image-generation requests (including overrides) to Gemini’s image conversion. Normalizes Vertex model names for thinking variants and enforces imagen “predict” endpoint selection. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant VertexAdaptor
participant GeminiAdaptor
participant VertexAPI
Client->>VertexAdaptor: OpenAI-style request (messages/prompt, size/n/extra)
alt Gemini mode AND model startsWith "imagen"
VertexAdaptor->>VertexAdaptor: Extract prompt and overrides
VertexAdaptor->>GeminiAdaptor: ConvertImageRequest(dto.ImageRequest)
GeminiAdaptor->>GeminiAdaptor: Map size to aspectRatio (e.g., "16:9" or presets)
GeminiAdaptor-->>VertexAdaptor: Converted request payload
VertexAdaptor->>VertexAPI: POST /.../imagen:predict
VertexAPI-->>VertexAdaptor: Response
else Other models/paths
VertexAdaptor->>VertexAPI: Route per existing logic
end
VertexAdaptor-->>Client: Translated response
sequenceDiagram
participant Caller
participant VertexAdaptor
Caller->>VertexAdaptor: GetRequestURL(model, thinking flags)
VertexAdaptor->>VertexAdaptor: Normalize -thinking/-nothinking/-budget suffix
alt model startsWith "imagen"
VertexAdaptor-->>Caller: .../predict URL
else
VertexAdaptor-->>Caller: .../generateContent or streaming URL
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
relay/channel/gemini/adaptor.go (1)
62-78: Aspect-ratio parsing is headed the right direction; normalize and sanitize input to avoid edge-case mismatchesGood call allowing explicit aspect ratios. Two small robustness tweaks:
- Normalize the size string to lowercase and strip spaces to handle variants like "1024X1792" or "1 : 1".
- When the input is already a ratio, sanitize embedded spaces.
Apply this diff:
- size := strings.TrimSpace(request.Size) + size := strings.ToLower(strings.TrimSpace(request.Size)) + // normalize common variants + size = strings.ReplaceAll(size, "X", "x") if size != "" { - if strings.Contains(size, ":") { - aspectRatio = size + if strings.Contains(size, ":") { + aspectRatio = strings.ReplaceAll(size, " ", "") } else { switch size { case "1024x1024": aspectRatio = "1:1" case "1024x1792": aspectRatio = "9:16" case "1792x1024": aspectRatio = "16:9" } } }Optional: If Imagen only accepts a fixed set (1:1, 9:16, 16:9), consider validating and rejecting unsupported ratios early with a clear error.
Would you like me to wire a small whitelist check and return a 400-like error for unsupported ratios?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (2)
relay/channel/gemini/adaptor.go(1 hunks)relay/channel/vertex/adaptor.go(2 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-08-08T17:12:43.157Z
Learnt from: RedwindA
PR: QuantumNous/new-api#1537
File: relay/gemini_handler.go:330-342
Timestamp: 2025-08-08T17:12:43.157Z
Learning: In the new-api repository, the `GeminiEmbeddingHandler` function in `relay/gemini_handler.go` is designed specifically for native Gemini embedding requests and therefore does not require the `ConvertGeminiRequest` step that is used in the chat handler. The embedding requests are already in the native Gemini format and don't need conversion.
Applied to files:
relay/channel/vertex/adaptor.go
🧬 Code Graph Analysis (1)
relay/channel/vertex/adaptor.go (4)
relay/channel/gemini/adaptor.go (1)
Adaptor(20-21)relay/channel/claude/adaptor.go (1)
Adaptor(23-25)relay/channel/aws/adaptor.go (1)
Adaptor(21-23)relay/channel/openai/adaptor.go (1)
Adaptor(32-35)
🔇 Additional comments (1)
relay/channel/vertex/adaptor.go (1)
69-71: Delegating image conversion to Gemini adaptor is a good consolidationCentralizing Imagen payload shaping in one place prevents drift between Vertex and Gemini paths. LGTM.
Comment on lines
+184
to
+237
| 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) | ||
| } |
Contributor
There was a problem hiding this comment.
💡 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.goLength 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
nfromExtraBodyto 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.
10 tasks
11 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
该 pr 包含以下更改：
Summary by CodeRabbit
New Features
Bug Fixes