Add Support for image generation tool

Author: zhoutt96

.oca/custom_code_review_guidelines.txt

@@ -0,0 +1,24 @@
+# Sample guideline, please follow similar structure for guideline with code samples
+# 1. Suggest using streams instead of simple loops for better readability.
+#          <example>
+#          *Comment:
+#          Category: Minor
+#          Issue: Use streams instead of a loop for better readability.
+#          Code Block:
+#
+#        ```java
+#            // Calculate squares of numbers
+#            List<Integer> squares = new ArrayList<>();
+#            for (int number : numbers) {
+#              squares.add(number * number);
+#            }
+#          ```
+#          Recommendation:
+#
+#        ```java
+#            // Calculate squares of numbers
+#            List<Integer> squares = Arrays.stream(numbers)
+#              .map(n -> n * n) // Map each number to its square
+#              .toList();
+#          ```
+#          </example>

crates/protocols/src/responses.rs

@@ -36,6 +36,10 @@ pub enum ResponseTool {
     #[serde(rename = "code_interpreter")]
     CodeInterpreter(CodeInterpreterTool),
 
+    /// Built-in tool.
+    #[serde(rename = "image_generation")]
+    ImageGeneration(ImageGenerationTool),
+
     /// MCP server tool.
     #[serde(rename = "mcp")]
     Mcp(McpTool),
@@ -80,6 +84,17 @@ pub struct CodeInterpreterTool {
     pub container: Option<Value>,
 }
 
+/// Built-in image generation tool.
+///
+/// The upstream API may add optional fields over time. We keep this tool shape
+/// open by preserving unknown key/value pairs.
+#[serde_with::skip_serializing_none]
+#[derive(Debug, Clone, Deserialize, Serialize, Default, schemars::JsonSchema)]
+pub struct ImageGenerationTool {
+    #[serde(flatten)]
+    pub options: HashMap<String, Value>,
+}
+
 /// `require_approval` values.
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, schemars::JsonSchema)]
 #[serde(rename_all = "snake_case")]
@@ -314,6 +329,13 @@ pub enum ResponseOutputItem {
         queries: Vec<String>,
         results: Option<Vec<FileSearchResult>>,
     },
+    #[serde(rename = "image_generation_call")]
+    ImageGenerationCall {
+        id: String,
+        status: ImageGenerationCallStatus,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        result: Option<String>,
+    },
 }
 
 // ============================================================================
@@ -389,6 +411,17 @@ pub enum FileSearchCallStatus {
     Failed,
 }
 
+/// Status for image generation tool calls.
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, schemars::JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum ImageGenerationCallStatus {
+    InProgress,
+    Generating,
+    Completed,
+    Incomplete,
+    Failed,
+}
+
 /// A result from file search.
 #[serde_with::skip_serializing_none]
 #[derive(Debug, Clone, Deserialize, Serialize, schemars::JsonSchema)]

model_gateway/src/routers/grpc/common/responses/utils.rs

@@ -44,7 +44,9 @@ pub(crate) async fn ensure_mcp_connection(
             t.iter().any(|tool| {
                 matches!(
                     tool,
-                    ResponseTool::WebSearchPreview(_) | ResponseTool::CodeInterpreter(_)
+                    ResponseTool::WebSearchPreview(_)
+                        | ResponseTool::CodeInterpreter(_)
+                        | ResponseTool::ImageGeneration(_)
                 )
             })
         })

model_gateway/src/routers/grpc/harmony/builder.rs

@@ -60,7 +60,12 @@ pub(crate) fn convert_harmony_logprobs(proto_logprobs: &ProtoOutputLogProbs) ->
 }
 
 /// Built-in tools that are added to the system message
-const BUILTIN_TOOLS: &[&str] = &["web_search_preview", "code_interpreter", "container"];
+const BUILTIN_TOOLS: &[&str] = &[
+    "web_search_preview",
+    "code_interpreter",
+    "image_generation",
+    "container",
+];
 
 /// Trait for tool-like objects that can be converted to Harmony ToolDescription
 trait ToolLike {
@@ -80,7 +85,7 @@ impl ToolLike for Tool {
     fn is_builtin(&self) -> bool {
         matches!(
             self.tool_type.as_str(),
-            "web_search_preview" | "code_interpreter" | "container"
+            "web_search_preview" | "code_interpreter" | "image_generation" | "container"
         )
     }
 
@@ -102,7 +107,9 @@ impl ToolLike for ResponseTool {
     fn is_builtin(&self) -> bool {
         matches!(
             self,
-            ResponseTool::WebSearchPreview(_) | ResponseTool::CodeInterpreter(_)
+            ResponseTool::WebSearchPreview(_)
+                | ResponseTool::CodeInterpreter(_)
+                | ResponseTool::ImageGeneration(_)
         )
     }
 
@@ -425,6 +432,7 @@ impl HarmonyBuilder {
                             ResponseTool::Function(_) => "function",
                             ResponseTool::WebSearchPreview(_) => "web_search_preview",
                             ResponseTool::CodeInterpreter(_) => "code_interpreter",
+                            ResponseTool::ImageGeneration(_) => "image_generation",
                             ResponseTool::Mcp(_) => "mcp",
                         })
                         .collect()

model_gateway/src/routers/openai/responses/utils.rs

@@ -230,6 +230,7 @@ pub(super) fn response_tool_to_value(tool: &ResponseTool) -> Option<Value> {
         }
         ResponseTool::WebSearchPreview(_) => serde_json::to_value(tool).ok(),
         ResponseTool::CodeInterpreter(_) => serde_json::to_value(tool).ok(),
+        ResponseTool::ImageGeneration(_) => serde_json::to_value(tool).ok(),
         ResponseTool::Function(_) => None,
     }
 }

oracle-schema-config.local.yaml

@@ -0,0 +1 @@
+auto_migrate: true