Fix image generation using FLUX.2 via AI Gateway (#23)

* Fix image generation by using gpt-image-1 via ai gateway

* Update model to bfl/flux-2-pro

* Remove dalle 2 and 3 support

* Remove flash image and its preview

* Replace open api key with ai gateway api key

* Add example usage of experimental_generateImage for 1024x1024 image

---------

Co-authored-by: Chris Tate <ctate@users.noreply.github.com>

Author: ctate

app/api/ai/generate/route.ts

@@ -183,7 +183,7 @@ Action types:
 - Database Query: {"actionType": "Database Query", "dbQuery": "SELECT * FROM table", "dbTable": "table"}
 - HTTP Request: {"actionType": "HTTP Request", "httpMethod": "POST", "endpoint": "https://api.example.com", "httpHeaders": "{}", "httpBody": "{}"}
 - Generate Text: {"actionType": "Generate Text", "aiModel": "gpt-5", "aiFormat": "text", "aiPrompt": "Your prompt here"}
-- Generate Image: {"actionType": "Generate Image", "imageModel": "openai/dall-e-3", "imagePrompt": "Image description"}
+- Generate Image: {"actionType": "Generate Image", "imageModel": "bfl/flux-2-pro", "imagePrompt": "Image description"}
 - Condition: {"actionType": "Condition", "condition": "{{@nodeId:Label.field}} === 'value'"}
 
 CRITICAL ABOUT CONDITION NODES:

components/workflow/config/action-config.tsx

@@ -501,20 +501,13 @@ function GenerateImageFields({
         <Select
           disabled={disabled}
           onValueChange={(value) => onUpdateConfig("imageModel", value)}
-          value={(config?.imageModel as string) || "openai/dall-e-3"}
+          value={(config?.imageModel as string) || "bfl/flux-2-pro"}
         >
           <SelectTrigger className="w-full" id="imageModel">
             <SelectValue placeholder="Select model" />
           </SelectTrigger>
           <SelectContent>
-            <SelectItem value="openai/dall-e-3">OpenAI DALL-E 3</SelectItem>
-            <SelectItem value="openai/dall-e-2">OpenAI DALL-E 2</SelectItem>
-            <SelectItem value="google/gemini-2.5-flash-image">
-              Google Gemini 2.5 Flash Image
-            </SelectItem>
-            <SelectItem value="google/gemini-2.5-flash-image-preview">
-              Google Gemini 2.5 Flash Image Preview
-            </SelectItem>
+            <SelectItem value="bfl/flux-2-pro">FLUX.2 Pro</SelectItem>
           </SelectContent>
         </Select>
       </div>

lib/codegen-templates/generate-image.ts

@@ -2,22 +2,24 @@
  * Code template for Generate Image action step
  * This is a string template used for code generation - keep as string export
  */
-export default `import OpenAI from 'openai';
+export default `import { experimental_generateImage as generateImage } from 'ai';
 
 export async function generateImageStep(input: {
   model: string;
   prompt: string;
 }) {
   "use step";
   
-  const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
-  
-  const response = await openai.images.generate({
-    model: input.model,
+  const result = await generateImage({
+    model: input.model as any,
     prompt: input.prompt,
-    n: 1,
-    response_format: 'b64_json',
+    size: '1024x1024',
+    providerOptions: {
+      openai: {
+        apiKey: process.env.AI_GATEWAY_API_KEY,
+      },
+    },
   });
   
-  return { base64: response.data[0].b64_json };
+  return { base64: result.image.toString() };
 }`;

lib/credential-fetcher.ts

@@ -24,7 +24,6 @@ export type WorkflowCredentials = {
   LINEAR_TEAM_ID?: string;
   SLACK_API_KEY?: string;
   AI_GATEWAY_API_KEY?: string;
-  OPENAI_API_KEY?: string;
   DATABASE_URL?: string;
 };
 
@@ -71,9 +70,6 @@ function mapAiGatewayConfig(config: IntegrationConfig): WorkflowCredentials {
   if (config.apiKey) {
     creds.AI_GATEWAY_API_KEY = config.apiKey;
   }
-  if (config.openaiApiKey) {
-    creds.OPENAI_API_KEY = config.openaiApiKey;
-  }
   return creds;
 }
 

lib/steps/generate-image.ts

@@ -6,7 +6,7 @@
  */
 import "server-only";
 
-import OpenAI from "openai";
+import { experimental_generateImage as generateImage } from "ai";
 import { fetchCredentials } from "../credential-fetcher";
 
 export async function generateImageStep(input: {
@@ -20,26 +20,32 @@ export async function generateImageStep(input: {
     ? await fetchCredentials(input.integrationId)
     : {};
 
-  const apiKey = credentials.OPENAI_API_KEY || credentials.AI_GATEWAY_API_KEY;
+  const apiKey = credentials.AI_GATEWAY_API_KEY;
 
   if (!apiKey) {
     throw new Error(
-      "OPENAI_API_KEY or AI_GATEWAY_API_KEY is not configured. Please add it in Project Integrations."
+      "AI_GATEWAY_API_KEY is not configured. Please add it in Project Integrations."
     );
   }
 
-  const openai = new OpenAI({ apiKey });
-
-  const response = await openai.images.generate({
-    model: input.model,
+  const result = await generateImage({
+    // biome-ignore lint/suspicious/noExplicitAny: model string needs type coercion for ai package
+    model: input.model as any,
     prompt: input.prompt,
-    n: 1,
-    response_format: "b64_json",
+    size: "1024x1024",
+    providerOptions: {
+      openai: {
+        apiKey,
+      },
+    },
   });
 
-  if (!response.data?.[0]) {
+  if (!result.image) {
     throw new Error("Failed to generate image");
   }
 
-  return { base64: response.data[0].b64_json };
+  // Convert the GeneratedFile to base64 string
+  const base64 = result.image.toString();
+
+  return { base64 };
 }

lib/steps/generate-text.ts

@@ -63,11 +63,11 @@ export async function generateTextStep(input: {
     ? await fetchCredentials(input.integrationId)
     : {};
 
-  const apiKey = credentials.AI_GATEWAY_API_KEY || credentials.OPENAI_API_KEY;
+  const apiKey = credentials.AI_GATEWAY_API_KEY;
 
   if (!apiKey) {
     throw new Error(
-      "AI_GATEWAY_API_KEY or OPENAI_API_KEY is not configured. Please add it in Project Integrations."
+      "AI_GATEWAY_API_KEY is not configured. Please add it in Project Integrations."
     );
   }
 

lib/workflow-codegen-sdk.ts

@@ -304,14 +304,14 @@ function _generateGenerateImageStepBody(
   const imagePrompt = (config.imagePrompt as string) || "";
   const convertedImagePrompt = convertTemplateToJS(imagePrompt);
 
-  return `  const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
+  return `  const openai = new OpenAI({ apiKey: process.env.AI_GATEWAY_API_KEY });
   
   // Use template literal with dynamic values from outputs
   const imagePrompt = \`${escapeForTemplateLiteral(convertedImagePrompt)}\`;
   const finalPrompt = (input.imagePrompt as string) || imagePrompt;
   
   const response = await openai.images.generate({
-    model: '${config.imageModel || "dall-e-3"}',
+    model: '${config.imageModel || "bfl/flux-2-pro"}',
     prompt: finalPrompt,
     n: 1,
     response_format: 'b64_json',
@@ -503,12 +503,15 @@ export function generateWorkflowSDKCode(
   }
 
   function buildAIImageParams(config: Record<string, unknown>): string[] {
-    imports.add("import OpenAI from 'openai';");
-    const imageModel = (config.imageModel as string) || "dall-e-3";
+    imports.add(
+      "import { experimental_generateImage as generateImage } from 'ai';"
+    );
+    const imageModel = (config.imageModel as string) || "bfl/flux-2-pro";
     return [
       `model: "${imageModel}"`,
       `prompt: \`${convertTemplateToJS((config.imagePrompt as string) || "")}\``,
-      "apiKey: process.env.OPENAI_API_KEY!",
+      'size: "1024x1024"',
+      "providerOptions: { openai: { apiKey: process.env.AI_GATEWAY_API_KEY! } }",
     ];
   }
 

lib/workflow-codegen.ts

@@ -438,17 +438,20 @@ export function generateWorkflowCode(
     indent: string,
     varName: string
   ): string[] {
-    imports.add("import { generateImage } from './integrations/ai';");
+    imports.add(
+      "import { experimental_generateImage as generateImage } from 'ai';"
+    );
     const imagePrompt =
       (node.data.config?.imagePrompt as string) || "A beautiful landscape";
     const imageModel =
-      (node.data.config?.imageModel as string) || "openai/dall-e-3";
+      (node.data.config?.imageModel as string) || "bfl/flux-2-pro";
 
     return [
       `${indent}// Generate image using AI`,
       `${indent}const ${varName} = await generateImage({`,
       `${indent}  model: "${imageModel}",`,
       `${indent}  prompt: \`${imagePrompt}\`,`,
+      `${indent}  size: "1024x1024",`,
       `${indent}});`,
     ];
   }