feat: [ENG-3733] Image output support on AI Gateway (#5409)

* added image pro preview pricing

* image output support in one turn

* add top k

* update sdk and add image gen config

---------

Co-authored-by: Justin <justintorre75@gmail.com>

Author: connortbot

packages/__tests__/cost/__snapshots__/registrySnapshots.test.ts.snap

@@ -8751,8 +8751,8 @@ exports[`Registry Snapshots verify registry state 1`] = `
     "claude-3.5-haiku:anthropic:*",
   ],
   "totalArchivedConfigs": 0,
-  "totalEndpoints": 255,
-  "totalModelProviderConfigs": 255,
+  "totalEndpoints": 256,
+  "totalModelProviderConfigs": 256,
   "totalModelsWithPtb": 95,
   "totalProviders": 21,
 }

packages/llm-mapper/transform/providers/google/response/toOpenai.ts

@@ -6,6 +6,7 @@ import {
   OpenAIToolCall,
   OpenAIUsage,
   OpenAIReasoningDetail,
+  ChatCompletionContentPartImage,
 } from "../../../types/openai";
 import {
   GoogleCandidate,
@@ -56,13 +57,15 @@ function mapCandidates(candidates: GoogleCandidate[]): OpenAIChoice[] {
     let tool_calls: OpenAIToolCall[] = [];
     let reasoning: string | undefined;
     let reasoning_details: OpenAIReasoningDetail[] | undefined;
+    let images: ChatCompletionContentPartImage[] | undefined;
 
     if (candidate.content) {
       const extracted = extractContent(candidate.content);
       content = extracted.content;
       tool_calls = extracted.tool_calls;
       reasoning = extracted.reasoning;
       reasoning_details = extracted.reasoning_details;
+      images = extracted.images;
     }
 
     return {
@@ -73,6 +76,7 @@ function mapCandidates(candidates: GoogleCandidate[]): OpenAIChoice[] {
         ...(reasoning && { reasoning }),
         ...(reasoning_details && reasoning_details.length > 0 && { reasoning_details }),
         ...(tool_calls.length > 0 && { tool_calls }),
+        ...(images && images.length > 0 && { images }),
       } as OpenAIResponseMessage,
       finish_reason: mapGoogleFinishReason(candidate.finishReason),
       logprobs: null,
@@ -88,6 +92,7 @@ interface ExtractedContent {
   tool_calls: OpenAIToolCall[];
   reasoning?: string;
   reasoning_details?: OpenAIReasoningDetail[];
+  images?: ChatCompletionContentPartImage[];
 }
 
 /**
@@ -106,6 +111,7 @@ function extractContent(
   const textParts: string[] = [];
   const thinkingTexts: string[] = [];
   const toolCalls: OpenAIToolCall[] = [];
+  const imageParts: ChatCompletionContentPartImage[] = [];
   let collectedSignature: string | undefined;
 
   for (const block of contents) {
@@ -127,6 +133,15 @@ function extractContent(
 
       if (part.functionCall) {
         toolCalls.push(mapToolCall(part.functionCall, toolCalls.length));
+      } else if (part.inlineData) {
+        const mimeType = part.inlineData.mimeType || "image/png";
+        const dataUri = `data:${mimeType};base64,${part.inlineData.data}`;
+        imageParts.push({
+          type: "image_url",
+          image_url: {
+            url: dataUri,
+          },
+        });
       } else if (part.text) {
         // Check if this is a thinking part (Google uses thought: true)
         if (part.thought === true) {
@@ -155,6 +170,11 @@ function extractContent(
     }));
   }
 
+  // Add images if image parts were found
+  if (imageParts.length > 0) {
+    result.images = imageParts;
+  }
+
   return result;
 }
 

packages/llm-mapper/transform/providers/google/streamedResponse/toOpenai.ts

@@ -136,6 +136,31 @@ export class GoogleToOpenAIStreamConverter {
           );
 
           this.toolCallIndex += 1;
+        } else if (part.inlineData) {
+          // Handle image output from Google's image generation models
+          const mimeType = part.inlineData.mimeType || "image/png";
+          const dataUri = `data:${mimeType};base64,${part.inlineData.data}`;
+          chunks.push(
+            this.createChunk({
+              choices: [
+                {
+                  index: candidate.index ?? 0,
+                  delta: {
+                    images: [
+                      {
+                        type: "image_url",
+                        image_url: {
+                          url: dataUri,
+                        },
+                      },
+                    ],
+                  },
+                  logprobs: null,
+                  finish_reason: null,
+                },
+              ],
+            })
+          );
         } else if (part.text) {
           // Check if this is a thinking part (Google uses thought: true)
           if (part.thought === true) {

packages/llm-mapper/transform/providers/openai/request/toAnthropic.ts

@@ -379,6 +379,11 @@ function mapMessages(
     };
 
     if (message.role === "assistant" || message.role === "user") {
+      // Check if assistant message has images - not supported by Anthropic
+      if (message.role === "assistant" && message.images && message.images.length > 0) {
+        throw new Error("Image outputs in assistant messages are not supported by Anthropic");
+      }
+
       const processedToolCallContent: AnthropicContentBlock[] = [];
 
       if (message.role === "assistant" && message.tool_calls) {

packages/llm-mapper/transform/providers/openai/request/toGoogle.ts

@@ -1,10 +1,14 @@
-import { ExtendedHeliconeChatCreateParams, GeminiContent, GeminiGenerateContentRequest, GeminiGenerationConfig, GeminiPart, GeminiThinkingConfig, GeminiTool, GeminiToolConfig, GoogleReasoningOptions, ChatCompletionMessage } from "../../../types/google";
+import { GeminiContent, GeminiGenerateContentRequest, GeminiGenerationConfig, GeminiImageConfig, GeminiPart, GeminiThinkingConfig, GeminiTool, GeminiToolConfig, GoogleReasoningOptions } from "../../../types/google";
 import {
   HeliconeChatCompletionContentPart,
   HeliconeChatCreateParams,
+  HeliconeChatCompletionMessageParam,
+  HeliconeImageGenerationConfig,
 } from "@helicone-package/prompts/types";
 import { ChatCompletionTool } from "openai/resources/chat/completions";
 
+type ChatCompletionMessage = NonNullable<HeliconeChatCreateParams["messages"]>[number];
+
 export function toGoogle(
   openAIBody: HeliconeChatCreateParams
 ): GeminiGenerateContentRequest {
@@ -63,6 +67,9 @@ export function toGoogle(
     if (message.role === "assistant") {
       const toolCallParts = mapToolCallsToParts(message);
       parts.push(...toolCallParts);
+
+      const imageParts = mapImagesToParts(message);
+      parts.push(...imageParts);
     }
 
     if (parts.length === 0) {
@@ -107,16 +114,13 @@ export function toGoogle(
 function buildGenerationConfig(
   body: HeliconeChatCreateParams
 ): GeminiGenerationConfig | undefined {
-  const bodyWithExtensions = body as ExtendedHeliconeChatCreateParams;
-
   const getNumberOrUndefined = (
     value?: number | null
   ): number | undefined => {
     return typeof value === "number" ? value : undefined;
   };
 
   const maxOutputTokens =
-    getNumberOrUndefined(bodyWithExtensions.max_output_tokens) ??
     getNumberOrUndefined(body.max_completion_tokens) ??
     getNumberOrUndefined(body.max_tokens) ??
     undefined;
@@ -137,7 +141,7 @@ function buildGenerationConfig(
   if (topP !== undefined) {
     config.topP = topP;
   }
-  const topK = getNumberOrUndefined(bodyWithExtensions.top_k);
+  const topK = getNumberOrUndefined(body.top_k);
   if (topK !== undefined) {
     config.topK = topK;
   }
@@ -166,6 +170,11 @@ function buildGenerationConfig(
     config.thinkingConfig = thinkingConfig;
   }
 
+  const imageConfig = buildImageConfig(body);
+  if (imageConfig) {
+    config.imageConfig = imageConfig;
+  }
+
   return Object.keys(config).length > 0 ? config : undefined;
 }
 
@@ -264,6 +273,20 @@ function buildThinkingConfig(
   return thinkingConfig;
 }
 
+function buildImageConfig(body: HeliconeChatCreateParams): GeminiImageConfig | undefined {
+  const heliconeImageConfig = body.image_generation;
+
+  if (heliconeImageConfig === undefined) {
+    return undefined;
+  }
+
+  const imageConfig: GeminiImageConfig = {
+    aspectRatio: heliconeImageConfig.aspect_ratio,
+    imageSize: heliconeImageConfig.image_size
+  };
+  return imageConfig;
+}
+
 function buildTools(body: HeliconeChatCreateParams): GeminiTool[] | undefined {
   if (!body.tools || body.tools.length === 0) {
     return undefined;
@@ -499,3 +522,33 @@ function parseArguments(
     return { raw: value };
   }
 }
+
+/**
+ * Maps images from assistant messages to Gemini inlineData parts.
+ * This handles image outputs that were generated by previous model responses.
+ */
+function mapImagesToParts(message: HeliconeChatCompletionMessageParam): GeminiPart[] {
+  const parts: GeminiPart[] = [];
+
+  if (!message.images) {
+    return parts;
+  }
+
+  for (const image of message.images) {
+    if (image.type === "image_url" && image.image_url?.url) {
+      const dataUri = image.image_url.url;
+      if (dataUri.startsWith("data:")) {
+        const [meta, data] = dataUri.split(",");
+        let mimeType = meta.split(";")[0].replace("data:", "");
+        if (!mimeType || mimeType.trim() === "") {
+          mimeType = "image/png";
+        }
+        parts.push({
+          inlineData: { mimeType, data },
+        });
+      }
+    }
+  }
+
+  return parts;
+}

packages/llm-mapper/transform/providers/responses/openai/response/toResponses.ts

@@ -1,9 +1,10 @@
-import { OpenAIResponseBody } from "../../../../types/openai";
+import { OpenAIResponseBody, ChatCompletionContentPartImage } from "../../../../types/openai";
 import {
   ResponsesResponseBody,
   ResponsesMessageOutputItem,
   ResponsesUsage,
   ResponsesReasoningItem,
+  ResponsesOutputContentPart,
 } from "../../../../types/responses";
 
 export function toResponses(body: OpenAIResponseBody): ResponsesResponseBody {
@@ -31,13 +32,36 @@ export function toResponses(body: OpenAIResponseBody): ResponsesResponseBody {
     output.push(reasoningItem);
   }
 
+  // Build message content from text and images
+  const messageContent: ResponsesOutputContentPart[] = [];
+
   if (message?.content) {
+    messageContent.push({
+      type: "output_text",
+      text: message.content,
+      annotations: message.annotations ?? [],
+    });
+  }
+
+  // Map images to output_image parts
+  if (message?.images && message.images.length > 0) {
+    for (const img of message.images) {
+      messageContent.push({
+        type: "output_image",
+        image_url: img.image_url?.url || "",
+        detail: img.image_url?.detail,
+      });
+    }
+  }
+
+  // Only create a message output item if we have content
+  if (messageContent.length > 0) {
     const msg: ResponsesMessageOutputItem = {
       id: `msg_${Math.random().toString(36).slice(2, 10)}`,
       type: "message",
       status: "completed",
       role: "assistant",
-      content: [{ type: "output_text", text: message.content, annotations: message.annotations ?? [] }],
+      content: messageContent,
     };
     output.push(msg);
   }

packages/llm-mapper/transform/providers/responses/request/toChatCompletions.ts

@@ -29,13 +29,20 @@ function convertContentParts(
         return { type: "text", text: p.text };
       case "input_image": {
         if (p.image_url) {
-          return { type: "image_url", image_url: { url: p.image_url } };
+          return { type: "image_url", image_url: { url: p.image_url, detail: p.detail } };
         }
         // Chat Completions does not support file_id for images directly
         throw new Error(
           "input_image with file_id is not supported by Chat Completions"
         );
       }
+      // Handle output_image when responses output is fed back as input
+      case "output_image": {
+        if (p.image_url) {
+          return { type: "image_url", image_url: { url: p.image_url, detail: p.detail } };
+        }
+        throw new Error("output_image missing image_url");
+      }
       case "input_file":
         // Chat Completions API does not support arbitrary files as message parts
         throw new Error("input_file is not supported by Chat Completions");
@@ -215,6 +222,7 @@ export function toChatCompletions(
     max_tokens: body.max_output_tokens,
     temperature: body.temperature,
     top_p: body.top_p,
+    top_k: body.top_k,
     n: body.n,
     stream: body.stream,
     tools,
@@ -234,6 +242,7 @@ export function toChatCompletions(
     stream_options: body.stream_options,
     // Context editing passthrough (only supported by Anthropic - will be stripped for other providers)
     context_editing: body.context_editing,
+    image_generation: body.image_generation,
     // Deprecated passthroughs (supported by Chat Completions clients)
     function_call: body.function_call,
     functions: body.functions,