-
Couldn't load subscription status.
- Fork 2.4k
feat: add image generation tool with OpenRouter integration #7474
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.
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
Changes from 5 commits
9ed3d53
ea8ff25
f65866a
73a42a3
229f608
7a7b12d
e85409f
0199992
2af9108
7a5db06
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -26,6 +26,33 @@ import { DEFAULT_HEADERS } from "./constants" | |||||
| import { BaseProvider } from "./base-provider" | ||||||
| import type { SingleCompletionHandler } from "../index" | ||||||
|
|
||||||
| // Image generation types | ||||||
| interface ImageGenerationResponse { | ||||||
| choices?: Array<{ | ||||||
| message?: { | ||||||
| content?: string | ||||||
| images?: Array<{ | ||||||
| type?: string | ||||||
| image_url?: { | ||||||
| url?: string | ||||||
| } | ||||||
| }> | ||||||
| } | ||||||
| }> | ||||||
| error?: { | ||||||
| message?: string | ||||||
| type?: string | ||||||
| code?: string | ||||||
| } | ||||||
| } | ||||||
|
|
||||||
| export interface ImageGenerationResult { | ||||||
| success: boolean | ||||||
| imageData?: string | ||||||
| imageFormat?: string | ||||||
| error?: string | ||||||
| } | ||||||
|
|
||||||
| // Add custom interface for OpenRouter params. | ||||||
| type OpenRouterChatCompletionParams = OpenAI.Chat.ChatCompletionCreateParams & { | ||||||
| transforms?: string[] | ||||||
|
|
@@ -242,4 +269,105 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH | |||||
| const completion = response as OpenAI.Chat.ChatCompletion | ||||||
| return completion.choices[0]?.message?.content || "" | ||||||
| } | ||||||
|
|
||||||
| /** | ||||||
| * Generate an image using OpenRouter's image generation API | ||||||
| * @param prompt The text prompt for image generation | ||||||
| * @param model The model to use for generation | ||||||
| * @param apiKey The OpenRouter API key (must be explicitly provided) | ||||||
| * @returns The generated image data and format, or an error | ||||||
| */ | ||||||
| async generateImage(prompt: string, model: string, apiKey: string): Promise<ImageGenerationResult> { | ||||||
| if (!apiKey) { | ||||||
| return { | ||||||
| success: false, | ||||||
| error: "OpenRouter API key is required for image generation", | ||||||
| } | ||||||
| } | ||||||
|
|
||||||
| try { | ||||||
| const response = await fetch("https://openrouter.ai/api/v1/chat/completions", { | ||||||
| method: "POST", | ||||||
| headers: { | ||||||
| Authorization: `Bearer ${apiKey}`, | ||||||
| "Content-Type": "application/json", | ||||||
| "HTTP-Referer": "https://github.com/RooVetGit/Roo-Code", | ||||||
| "X-Title": "Roo Code", | ||||||
| }, | ||||||
| body: JSON.stringify({ | ||||||
| model, | ||||||
| messages: [ | ||||||
| { | ||||||
| role: "user", | ||||||
| content: prompt, | ||||||
| }, | ||||||
| ], | ||||||
| modalities: ["image", "text"], | ||||||
| }), | ||||||
| }) | ||||||
|
|
||||||
| if (!response.ok) { | ||||||
| const errorText = await response.text() | ||||||
| let errorMessage = `Failed to generate image: ${response.status} ${response.statusText}` | ||||||
| try { | ||||||
| const errorJson = JSON.parse(errorText) | ||||||
| if (errorJson.error?.message) { | ||||||
| errorMessage = `Failed to generate image: ${errorJson.error.message}` | ||||||
| } | ||||||
| } catch { | ||||||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In the JSON.parse try/catch block (lines 312-319), consider logging the caught error (e.g. using console.error) to aid debugging of parsing failures.
Suggested change
This comment was generated because it violated a code review rule: irule_PTI8rjtnhwrWq6jS. |
||||||
| // Use default error message | ||||||
| } | ||||||
| return { | ||||||
| success: false, | ||||||
| error: errorMessage, | ||||||
| } | ||||||
| } | ||||||
|
|
||||||
| const result: ImageGenerationResponse = await response.json() | ||||||
|
|
||||||
| if (result.error) { | ||||||
| return { | ||||||
| success: false, | ||||||
| error: `Failed to generate image: ${result.error.message}`, | ||||||
| } | ||||||
| } | ||||||
|
|
||||||
| // Extract the generated image from the response | ||||||
| const images = result.choices?.[0]?.message?.images | ||||||
| if (!images || images.length === 0) { | ||||||
| return { | ||||||
| success: false, | ||||||
| error: "No image was generated in the response", | ||||||
| } | ||||||
| } | ||||||
|
|
||||||
| const imageData = images[0]?.image_url?.url | ||||||
| if (!imageData) { | ||||||
| return { | ||||||
| success: false, | ||||||
| error: "Invalid image data in response", | ||||||
| } | ||||||
| } | ||||||
|
|
||||||
| // Extract base64 data from data URL | ||||||
| const base64Match = imageData.match(/^data:image\/(png|jpeg|jpg);base64,(.+)$/) | ||||||
| if (!base64Match) { | ||||||
| return { | ||||||
| success: false, | ||||||
| error: "Invalid image format received", | ||||||
| } | ||||||
| } | ||||||
|
|
||||||
| return { | ||||||
| success: true, | ||||||
| imageData: imageData, | ||||||
| imageFormat: base64Match[1], | ||||||
| } | ||||||
| } catch (error) { | ||||||
| return { | ||||||
| success: false, | ||||||
| error: error instanceof Error ? error.message : "Unknown error occurred", | ||||||
| } | ||||||
| } | ||||||
| } | ||||||
| } | ||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| import { ToolArgs } from "./types" | ||
|
|
||
| export function getGenerateImageDescription(args: ToolArgs): string { | ||
| return `## generate_image | ||
| Description: Request to generate an image using AI models through OpenRouter API. This tool creates images from text prompts and saves them to the specified path. | ||
| Parameters: | ||
| - prompt: (required) The text prompt describing the image to generate | ||
| - path: (required) The file path where the generated image should be saved (relative to the current workspace directory ${args.cwd}). The tool will automatically add the appropriate image extension if not provided. | ||
| Usage: | ||
| <generate_image> | ||
| <prompt>Your image description here</prompt> | ||
| <path>path/to/save/image.png</path> | ||
| </generate_image> | ||
|
|
||
| Example: Requesting to generate a sunset image | ||
| <generate_image> | ||
| <prompt>A beautiful sunset over mountains with vibrant orange and purple colors</prompt> | ||
| <path>images/sunset.png</path> | ||
| </generate_image>` | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I wonder if we can remove this from the base provider schema