Image generation (Imagen + Gemini 2.5) + top-bar quick pickers - #9
Conversation
- Imagen integration: new `generate_image` tool that calls
`models/{imagen}:predict` with the API key and the imagen model
selected in Settings (default imagen-3.0-generate-002; picker lists
3 built-in variants + custom ID). Decoded PNGs are persisted under
<filesDir>/attachments/imagen-*.png and returned both as text
(for the model turn) and as `attachmentPaths` on the ToolCallResult,
so the tool-result bubble renders them as thumbnails.
- Gemini 2.5 native multimodal output: when the selected chat model's
name contains "-image" (heuristic, auto-detects
`gemini-2.5-flash-image-preview`), the request gets
`generationConfig.responseModalities = [TEXT, IMAGE]`. Response
parser now handles `inlineData` parts, persists the bytes, and
attaches the paths to the live model bubble alongside streamed text.
- Model + tool-result bubbles now render `AttachmentThumbnails`, same
120-dp thumbnails as the user side.
- Shared persistence: `RestGeminiCore.persistAttachmentBytes(id,bytes,mime)`
is the single helper used by the response parser and by the Imagen
tool.
- Top bar quick pickers: tapping the model name opens a DropdownMenu
listing the discovered models (+ "More models…" shortcut to full
settings). Tapping the workspace folder name opens a menu with
"Open folder" (best-effort ACTION_VIEW on the tree/file URI,
falls back to a toast) and "Change folder" (launches
OpenDocumentTree).
Reviewer's GuideAdds end-to-end image generation support via a new Imagen-based tool and native Gemini 2.5 image outputs, including shared attachment persistence and thumbnail rendering, and introduces top-bar quick pickers for switching models and workspace folders directly from the chat UI. Sequence diagram for generate_image tool-based image generationsequenceDiagram
actor User
participant ChatScreen
participant ChatViewModel
participant RestGeminiCore
participant GeminiModel
participant GenerateImageTool
participant ImagenAPI
participant Storage as AttachmentsStorage
User->>ChatScreen: Type prompt "draw a red fox"
ChatScreen->>ChatViewModel: sendPrompt(prompt, currentModel)
ChatViewModel->>RestGeminiCore: startChatTurn(prompt)
RestGeminiCore->>GeminiModel: POST contents (non image model)
GeminiModel-->>RestGeminiCore: Response with tool call generate_image
RestGeminiCore->>GenerateImageTool: execute(ToolCall)
GenerateImageTool->>GenerateImageTool: read apiKey via getApiKey()
GenerateImageTool->>GenerateImageTool: read model via getModel()
GenerateImageTool->>ImagenAPI: POST models/{imagenModel}:predict
ImagenAPI-->>GenerateImageTool: predictions[].bytesBase64Encoded
GenerateImageTool->>GenerateImageTool: decode Base64 to bytes
loop for each image
GenerateImageTool->>RestGeminiCore: persist(id, bytes, mime)
RestGeminiCore->>AttachmentsStorage: persistAttachmentBytes(id, bytes, mime)
AttachmentsStorage-->>RestGeminiCore: imagePath
RestGeminiCore-->>GenerateImageTool: imagePath
end
GenerateImageTool-->>RestGeminiCore: ToolCallResult(attachmentPaths)
RestGeminiCore->>RestGeminiCore: create GeminiMessage(role TOOL, attachmentPaths)
RestGeminiCore-->>ChatViewModel: GeminiEvent.ToolCallCompleted
ChatViewModel-->>ChatScreen: updated messages with tool-result message
ChatScreen->>ChatScreen: MessageBubble(toolMessage)
ChatScreen->>ChatScreen: AttachmentThumbnails(message.attachmentPaths)
ChatScreen-->>User: Tool-result bubble with image thumbnails
Sequence diagram for inline Gemini 2.5 image output handlingsequenceDiagram
actor User
participant ChatScreen
participant ChatViewModel
participant RestGeminiCore
participant GeminiModel as Gemini2_5_ImageModel
participant AttachmentsStorage
User->>ChatScreen: Select gemini-2.5-flash-image-preview
ChatScreen->>ChatViewModel: setModel("gemini-2.5-flash-image-preview")
ChatViewModel->>RestGeminiCore: setModel(name)
User->>ChatScreen: Ask for "a landscape oil painting"
ChatScreen->>ChatViewModel: sendPrompt(prompt)
ChatViewModel->>RestGeminiCore: startChatTurn(prompt)
RestGeminiCore->>RestGeminiCore: buildRequestBody()
RestGeminiCore->>RestGeminiCore: modelEmitsImages(model) == true
RestGeminiCore->>RestGeminiCore: add generationConfig.responseModalities [TEXT, IMAGE]
RestGeminiCore->>GeminiModel: POST chat request with responseModalities
GeminiModel-->>RestGeminiCore: Streaming response parts
loop for each part
alt part has inlineData
RestGeminiCore->>RestGeminiCore: read inlineData.mimeType, data
RestGeminiCore->>RestGeminiCore: Base64 decode bytes
RestGeminiCore->>AttachmentsStorage: persistAttachmentBytes(id, bytes, mime)
AttachmentsStorage-->>RestGeminiCore: imagePath
alt first image and liveMessage is null
RestGeminiCore->>RestGeminiCore: create liveMessage(role MODEL, attachmentPaths=[imagePath])
RestGeminiCore-->>ChatViewModel: addUiMessage(liveMessage)
else subsequent image
RestGeminiCore->>RestGeminiCore: liveMessage.copy(attachmentPaths + imagePath)
RestGeminiCore-->>ChatViewModel: replaceUiMessage(liveMessage)
end
else text part
RestGeminiCore->>RestGeminiCore: append to accumulated text
end
end
ChatViewModel-->>ChatScreen: updated model message with attachmentPaths
ChatScreen->>ChatScreen: MessageBubble(modelMessage)
ChatScreen->>ChatScreen: AttachmentThumbnails(message.attachmentPaths)
ChatScreen-->>User: Model bubble showing text and image thumbnails
Class diagram for Imagen tool, core, and UI integrationclassDiagram
class RestGeminiCore {
- Context appContext
- SecurePrefs prefs
- ToolRegistry registry
+ String imagenModel()
+ void setImagenModel(name String)
+ String~nullable~ persistAttachmentBytes(id String, bytes ByteArray, mime String)
+ String buildRequestBody()
+ void setModel(name String)
+ List~String~ listModels()
+ String currentModel()
+ Boolean isAutoCompressEnabled()
+ void setAutoCompressEnabled(enabled Boolean)
+ Float autoCompressThreshold()
+ void setAutoCompressThreshold(fraction Float)
+ void handleToolResult(result ToolCallResult)
+ void parseInlineData(part JSONObject)
<<companion>>
+ String DEFAULT_MODEL
+ String DEFAULT_IMAGEN_MODEL
+ List~String~ AVAILABLE_MODELS
+ List~String~ AVAILABLE_IMAGEN_MODELS
+ Boolean modelEmitsImages(modelName String)
}
class GenerateImageTool {
- () -> String getApiKey
- () -> String getModel
- (String, ByteArray, String) -> String~nullable~ persist
+ ToolSpec spec
+ ToolCallResult execute(call ToolCall)
}
class SecurePrefs {
- SharedPreferences plain
+ String~nullable~ model
+ String~nullable~ imagenModel
+ String~nullable~ workspaceUri
+ Boolean autoApprove
+ Boolean autoCompressEnabled
+ Float autoCompressThreshold
}
class Workspace {
- DocumentFile~nullable~ root
- String label
+ void init()
+ String rootLabel()
+ Uri~nullable~ rootUri()
+ void setFile(file File)
+ void setTree(uri Uri)
+ String absolutePath()
+ String~nullable~ unreachableReason()
}
class ChatViewModel {
- RestGeminiCore core
- MutableStateFlow~String~ _model
- MutableStateFlow~List~String~~ _availableModels
- MutableStateFlow~String~ _imagenModel
- MutableStateFlow~String~ _workspaceLabel
- MutableStateFlow~String~ _workspacePath
- MutableStateFlow~String~ _workspaceReason
- MutableStateFlow~String~ _workspaceUri
+ StateFlow~String~ model
+ StateFlow~List~String~~ availableModels
+ StateFlow~String~ imagenModel
+ StateFlow~String~ workspaceLabel
+ StateFlow~String~ workspaceUri
+ void setModel(name String)
+ void setImagenModel(name String)
+ void setProjectFolder(uri String)
+ void refreshWorkspace()
}
class ToolCallResult {
+ String callId
+ Boolean ok
+ String output
+ Boolean truncated
+ List~String~ attachmentPaths
}
class GeminiMessage {
+ String id
+ String text
+ Boolean isUser
+ Long timestamp
+ MessageRole role
+ ToolCallResult~nullable~ toolResult
+ List~String~ attachmentPaths
}
class ChatScreen {
+ void ChatScreen(viewModel ChatViewModel, onEvent Function)
+ void openWorkspaceFolder(context Context, workspaceUri String~nullable~)
+ void AttachmentThumbnails(paths List~String~)
+ void MessageBubble(message GeminiMessage)
}
class SettingsSheet {
+ void SettingsSheet(viewModel ChatViewModel)
}
class ToolRegistry {
+ void register(tool Tool)
}
class Tool {
<<interface>>
+ ToolSpec spec
+ ToolCallResult execute(call ToolCall)
}
RestGeminiCore --> SecurePrefs : uses
RestGeminiCore --> Workspace : uses
RestGeminiCore --> ToolRegistry : owns
RestGeminiCore ..> GenerateImageTool : registers
RestGeminiCore --> ToolCallResult : creates
RestGeminiCore --> GeminiMessage : creates
RestGeminiCore "1" --> "*" Tool : invokes
GenerateImageTool ..|> Tool
GenerateImageTool --> ToolCallResult : returns
ChatViewModel --> RestGeminiCore : delegates
ChatViewModel --> Workspace : reads rootUri
ChatScreen --> ChatViewModel : observes state
ChatScreen --> GeminiMessage : renders
SettingsSheet --> ChatViewModel : reads and sets imagenModel
SecurePrefs --> Workspace : stores workspaceUri
GeminiMessage --> ToolCallResult : toolResult
ToolRegistry --> Tool : holds
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
|
Warning Rate limit exceeded
Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 54 minutes and 38 seconds. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThis PR implements end-to-end image generation support by introducing the Changes
Sequence DiagramsequenceDiagram
participant User
participant ChatScreen
participant GenerateImageTool
participant GeminiAPI as Gemini API
participant Persistence
participant ChatViewModel
User->>ChatScreen: Sends message triggering image generation
ChatScreen->>ChatViewModel: Captures message with generate_image intent
ChatViewModel->>GenerateImageTool: Executes tool via RestGeminiCore
GenerateImageTool->>GeminiAPI: POST /v1beta/models/$model:predict<br/>(prompt, aspectRatio, sampleCount)
GeminiAPI-->>GenerateImageTool: Returns predictions with<br/>bytesBase64Encoded & mimeType
GenerateImageTool->>Persistence: persistAttachmentBytes(id, bytes, mime)
Persistence-->>GenerateImageTool: Returns file path(s)
GenerateImageTool-->>ChatViewModel: Returns ToolCallResult<br/>with attachmentPaths
ChatViewModel->>ChatScreen: Updates message with<br/>attachmentPaths
ChatScreen->>ChatScreen: Renders AttachmentThumbnails
ChatScreen-->>User: Displays generated images
Estimated Code Review Effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly Related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
There was a problem hiding this comment.
Hey - I've found 2 issues, and left some high level feedback:
- In
persistAttachmentBytes, the mime-type checks use case-sensitivecontainswhich may miss valid types likeimage/PNG; consider usingcontains("png", ignoreCase = true)(and similarly for the other types) or parsing the subtype more robustly before mapping to an extension. - In
openWorkspaceFolder, the fallbackIntent(Intent.ACTION_VIEW, target)does not carryFLAG_GRANT_READ_URI_PERMISSION, so opening acontent://URI may fail where the first attempt would have worked; propagate the read permission flag to the fallback intent as well. - The new top-bar model and folder pickers rely on clickable
Textwith a small padded area; for better accessibility/usability, consider wrapping them in a larger hit target (e.g.Modifier.sizeIn(minHeight = ...)or anIconButton-style container) so they are easier to tap.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `persistAttachmentBytes`, the mime-type checks use case-sensitive `contains` which may miss valid types like `image/PNG`; consider using `contains("png", ignoreCase = true)` (and similarly for the other types) or parsing the subtype more robustly before mapping to an extension.
- In `openWorkspaceFolder`, the fallback `Intent(Intent.ACTION_VIEW, target)` does not carry `FLAG_GRANT_READ_URI_PERMISSION`, so opening a `content://` URI may fail where the first attempt would have worked; propagate the read permission flag to the fallback intent as well.
- The new top-bar model and folder pickers rely on clickable `Text` with a small padded area; for better accessibility/usability, consider wrapping them in a larger hit target (e.g. `Modifier.sizeIn(minHeight = ...)` or an `IconButton`-style container) so they are easier to tap.
## Individual Comments
### Comment 1
<location path="app/src/main/kotlin/com/gemini/app/ui/chat/ChatScreen.kt" line_range="1218-1220" />
<code_context>
+ } else uri
+ intent.setDataAndType(target, "vnd.android.document/directory")
+ runCatching { context.startActivity(intent) }.onFailure {
+ // Retry without mime type — some Files apps only match on the URI.
+ val fallback = Intent(Intent.ACTION_VIEW, target)
+ .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK)
+ runCatching { context.startActivity(fallback) }.onFailure {
+ Toast.makeText(
</code_context>
<issue_to_address>
**issue (bug_risk):** Fallback folder intent drops the read-permission flag, which can break `content://` URIs.
In `openWorkspaceFolder`, the fallback `Intent` created on retry omits `FLAG_GRANT_READ_URI_PERMISSION`, unlike the primary intent. For SAF `content://` URIs this means the target app may not be able to read the folder on the fallback path. Please mirror the permission flags from the primary intent on the fallback (at least `FLAG_GRANT_READ_URI_PERMISSION`) so both code paths behave consistently.
</issue_to_address>
### Comment 2
<location path="core-bridge/src/main/kotlin/com/gemini/bridge/RestGeminiCore.kt" line_range="193-206" />
<code_context>
+ * by the GenerateImageTool (for Imagen outputs). Returns the absolute path
+ * that can be stored in `GeminiMessage.attachmentPaths`.
+ */
+ fun persistAttachmentBytes(id: String, bytes: ByteArray, mime: String): String? = runCatching {
+ val ext = when {
+ mime.contains("png") -> "png"
+ mime.contains("webp") -> "webp"
+ mime.contains("gif") -> "gif"
+ mime.contains("heic") -> "heic"
+ mime.contains("heif") -> "heif"
+ else -> "jpg"
+ }
+ val dir = java.io.File(appContext.filesDir, "attachments").also { it.mkdirs() }
+ val file = java.io.File(dir, "$id.$ext")
+ file.writeBytes(bytes)
+ file.absolutePath
+ }.getOrNull()
</code_context>
<issue_to_address>
**suggestion (bug_risk):** MIME-type checks are case-sensitive and may misclassify some valid image types.
In `persistAttachmentBytes`, `mime.contains("png")` / `"webp"` / etc. are case-sensitive, so values like `image/PNG` or `IMAGE/WEBP` will incorrectly fall back to `jpg`. Normalize once (e.g., `val m = mime.lowercase()`) and run the `contains` checks on `m` to avoid this misclassification.
```suggestion
fun persistAttachmentBytes(id: String, bytes: ByteArray, mime: String): String? = runCatching {
val m = mime.lowercase()
val ext = when {
m.contains("png") -> "png"
m.contains("webp") -> "webp"
m.contains("gif") -> "gif"
m.contains("heic") -> "heic"
m.contains("heif") -> "heif"
else -> "jpg"
}
val dir = java.io.File(appContext.filesDir, "attachments").also { it.mkdirs() }
val file = java.io.File(dir, "$id.$ext")
file.writeBytes(bytes)
file.absolutePath
}.getOrNull()
```
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| // Retry without mime type — some Files apps only match on the URI. | ||
| val fallback = Intent(Intent.ACTION_VIEW, target) | ||
| .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK) |
There was a problem hiding this comment.
issue (bug_risk): Fallback folder intent drops the read-permission flag, which can break content:// URIs.
In openWorkspaceFolder, the fallback Intent created on retry omits FLAG_GRANT_READ_URI_PERMISSION, unlike the primary intent. For SAF content:// URIs this means the target app may not be able to read the folder on the fallback path. Please mirror the permission flags from the primary intent on the fallback (at least FLAG_GRANT_READ_URI_PERMISSION) so both code paths behave consistently.
| fun persistAttachmentBytes(id: String, bytes: ByteArray, mime: String): String? = runCatching { | ||
| val ext = when { | ||
| mime.contains("png") -> "png" | ||
| mime.contains("webp") -> "webp" | ||
| mime.contains("gif") -> "gif" | ||
| mime.contains("heic") -> "heic" | ||
| mime.contains("heif") -> "heif" | ||
| else -> "jpg" | ||
| } | ||
| val dir = java.io.File(appContext.filesDir, "attachments").also { it.mkdirs() } | ||
| val file = java.io.File(dir, "$id.$ext") | ||
| file.writeBytes(bytes) | ||
| file.absolutePath | ||
| }.getOrNull() |
There was a problem hiding this comment.
suggestion (bug_risk): MIME-type checks are case-sensitive and may misclassify some valid image types.
In persistAttachmentBytes, mime.contains("png") / "webp" / etc. are case-sensitive, so values like image/PNG or IMAGE/WEBP will incorrectly fall back to jpg. Normalize once (e.g., val m = mime.lowercase()) and run the contains checks on m to avoid this misclassification.
| fun persistAttachmentBytes(id: String, bytes: ByteArray, mime: String): String? = runCatching { | |
| val ext = when { | |
| mime.contains("png") -> "png" | |
| mime.contains("webp") -> "webp" | |
| mime.contains("gif") -> "gif" | |
| mime.contains("heic") -> "heic" | |
| mime.contains("heif") -> "heif" | |
| else -> "jpg" | |
| } | |
| val dir = java.io.File(appContext.filesDir, "attachments").also { it.mkdirs() } | |
| val file = java.io.File(dir, "$id.$ext") | |
| file.writeBytes(bytes) | |
| file.absolutePath | |
| }.getOrNull() | |
| fun persistAttachmentBytes(id: String, bytes: ByteArray, mime: String): String? = runCatching { | |
| val m = mime.lowercase() | |
| val ext = when { | |
| m.contains("png") -> "png" | |
| m.contains("webp") -> "webp" | |
| m.contains("gif") -> "gif" | |
| m.contains("heic") -> "heic" | |
| m.contains("heif") -> "heif" | |
| else -> "jpg" | |
| } | |
| val dir = java.io.File(appContext.filesDir, "attachments").also { it.mkdirs() } | |
| val file = java.io.File(dir, "$id.$ext") | |
| file.writeBytes(bytes) | |
| file.absolutePath | |
| }.getOrNull() |
Summary
Image generation (both paths, auto-enabled)
generate_imagetool that callsmodels/{imagen}:predictwith the user's API key. The chat model invokes it automatically when asked to draw/illustrate. Bytes are persisted under<filesDir>/attachments/imagen-*.pngand attached to the tool-result bubble as thumbnails.-image(e.g.gemini-2.5-flash-image-preview), the request payload getsgenerationConfig.responseModalities = ["TEXT","IMAGE"]. The response parser now handlesinlineDataparts, decodes them, persists to disk, and attaches the paths to the live model bubble. Toggle is automatic per the user's choice — picking an image-capable model is enough.imagen-3.0-generate-002(default),imagen-3.0-fast-generate-001,imagen-4.0-generate-preview-06-06, plus a custom-ID field. Stored in plain prefs underimagen_model.RestGeminiCore.persistAttachmentBytes(id, bytes, mime)used by both the response parser and the Imagen tool.AttachmentThumbnails(120-dp squares, sampled BitmapFactory), same as the user side.Top-bar quick pickers
DropdownMenulisting the discovered models with the active one highlighted. Last entry is "More models…" which opens the full Settings sheet (so the user can still set a custom ID / refresh the list).ACTION_VIEWwithvnd.android.document/directoryon the tree's document URI (falls back to a plainACTION_VIEWon the URI, and to a toast when no Files app handles it).OpenDocumentTreedirectly from the chat screen (no need to go through Settings).Files touched
core-bridge/.../RestGeminiCore.kt— tool registration,persistAttachmentBytes,inlineDataparser branch,generationConfig.responseModalitiesgated onmodelEmitsImages(),imagenModel()/setImagenModel(), newAVAILABLE_IMAGEN_MODELS+DEFAULT_IMAGEN_MODELconstants.core-bridge/.../tools/GenerateImageTool.kt— new; POSTs to:predict, decodespredictions[].bytesBase64Encoded, persists + returnsToolCallResult(attachmentPaths=…).core-bridge/.../workspace/Workspace.kt— exposerootUri(): Uri?for "Open folder".core-bridge/.../storage/SecurePrefs.kt— newimagenModelplain pref.domain/.../ToolTypes.kt—ToolCallResult.attachmentPathsfield.app/.../ui/chat/ChatScreen.kt— top-bar clickable model/folder withDropdownMenus, folder-picker launcher,openWorkspaceFolderhelper, model bubble now rendersAttachmentThumbnails.app/.../ui/chat/ChatViewModel.kt—imagenModelstate + setter,workspaceUristate.app/.../ui/settings/SettingsSheet.kt— "Image generation model" picker subsection inside Model accordion (dropdown + custom ID input).Test plan
generate_image, image appears as thumbnail in the tool-result bubble.imagen-3.0-fast-generate-001, repeat — new model name appears in the tool-result header.gemini-2.5-flash-image-preview, ask for "a landscape oil painting" — image appears directly in the model bubble alongside any text.https://claude.ai/code/session_015dVKN2jG34HKP9SeBpSvS5
Summary by Sourcery
Add automatic image generation support via Imagen and Gemini image-capable models, and introduce quick-access model and workspace folder pickers in the chat top bar.
New Features:
Enhancements:
Summary by CodeRabbit