feat(grok): add runtime image surfaces

This commit is contained in:
mayrain
2026-05-16 21:15:38 +08:00
parent 936e1ae37b
commit 5bf236957e
43 changed files with 7044 additions and 400 deletions
@@ -1,4 +1,4 @@
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
use serde_json::{json, Map, Value};
@@ -50,6 +50,7 @@ pub struct OpenAIResponsesProviderState {
tool_calls: BTreeMap<usize, OpenAIResponsesProviderToolState>,
tool_results: BTreeMap<usize, OpenAIResponsesProviderToolResultState>,
tool_index_by_key: BTreeMap<String, usize>,
image_item_keys: BTreeSet<String>,
last_tool_index: Option<usize>,
}
@@ -718,6 +719,55 @@ impl OpenAIResponsesProviderState {
}
}
fn emit_image_generation_item(
&mut self,
report_context: &Value,
out: &mut Vec<CanonicalStreamFrame>,
item: &Map<String, Value>,
output_index: Option<usize>,
final_item: bool,
) {
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
return;
}
if !final_item
&& !item
.get("status")
.and_then(Value::as_str)
.is_some_and(|value| value.eq_ignore_ascii_case("completed"))
{
return;
}
let has_image_payload = item
.get("result")
.or_else(|| item.get("url"))
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty());
if !has_image_payload {
return;
}
let index = output_index.unwrap_or(self.image_item_keys.len());
let key = item
.get("id")
.and_then(Value::as_str)
.map(ToOwned::to_owned)
.unwrap_or_else(|| format!("image_generation_call:{index}"));
if !self.image_item_keys.insert(key) {
return;
}
self.ensure_started(report_context, out);
let (id, model) = self.identity(report_context);
out.push(CanonicalStreamFrame {
id,
model,
event: CanonicalStreamEvent::ImageGenerationCall {
index,
item: Value::Object(item.clone()),
},
});
}
pub fn push_line(
&mut self,
report_context: &Value,
@@ -894,6 +944,15 @@ impl OpenAIResponsesProviderState {
"reasoning" => {
self.ensure_started(report_context, &mut out);
}
"image_generation_call" => {
self.emit_image_generation_item(
report_context,
&mut out,
item,
output_index,
false,
);
}
_ => {
out.push(self.unknown_frame(report_context, Value::Object(item.clone())));
}
@@ -1107,6 +1166,15 @@ impl OpenAIResponsesProviderState {
"reasoning" => {
self.emit_reasoning_item(report_context, &mut out, item);
}
"image_generation_call" => {
self.emit_image_generation_item(
report_context,
&mut out,
item,
output_index,
true,
);
}
_ => {
out.push(self.unknown_frame(report_context, Value::Object(item.clone())));
}
@@ -1152,6 +1220,15 @@ impl OpenAIResponsesProviderState {
"reasoning" => {
self.emit_reasoning_item(report_context, &mut out, item);
}
"image_generation_call" => {
self.emit_image_generation_item(
report_context,
&mut out,
item,
Some(output_index),
true,
);
}
_ => {
out.push(
self.unknown_frame(report_context, Value::Object(item.clone())),
@@ -1255,6 +1332,7 @@ pub struct OpenAIResponsesClientEmitter {
reasoning_summary_parts: Vec<String>,
tool_calls: BTreeMap<usize, OpenAIResponsesClientToolState>,
tool_results: BTreeMap<usize, OpenAIResponsesClientToolResultState>,
image_generation_items: BTreeMap<usize, Value>,
}
impl OpenAIChatClientEmitter {
@@ -1361,6 +1439,26 @@ impl OpenAIChatClientEmitter {
)?);
Ok(out)
}
CanonicalStreamEvent::ImageGenerationCall { item, .. } => {
let Some(part) = content_part_from_openai_image_generation_item(&item) else {
return Ok(Vec::new());
};
let placeholder = openai_stream_placeholder_for_content_part(&part);
let mut out = self.ensure_started()?;
out.extend(encode_json_sse(
None,
&build_openai_chat_chunk(
self.response_id
.as_deref()
.unwrap_or("chatcmpl-local-stream"),
self.model.as_deref().unwrap_or("unknown"),
placeholder,
None,
None,
),
)?);
Ok(out)
}
CanonicalStreamEvent::ToolCallStart {
index,
call_id,
@@ -1653,6 +1751,11 @@ impl OpenAIResponsesClientEmitter {
output_index
}
fn ensure_image_generation_output_index(&mut self, index: usize) -> usize {
self.next_output_index = self.next_output_index.max(index.saturating_add(1));
index
}
fn ensure_reasoning_item_started(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = self.ensure_started()?;
let output_index = self.ensure_reasoning_output_index();
@@ -1956,6 +2059,42 @@ impl OpenAIResponsesClientEmitter {
Ok(out)
}
fn emit_image_generation_call_item(
&mut self,
index: usize,
item: Value,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut out = self.ensure_started()?;
let output_index = self.ensure_image_generation_output_index(index);
let mut item = item.as_object().cloned().unwrap_or_default();
item.insert(
"type".to_string(),
Value::String("image_generation_call".to_string()),
);
if !item.contains_key("id") {
item.insert(
"id".to_string(),
Value::String(format!("{}_ig_{}", self.response_id(), output_index)),
);
}
if !item.contains_key("status") {
item.insert("status".to_string(), Value::String("completed".to_string()));
}
let item = Value::Object(item);
self.image_generation_items
.insert(output_index, item.clone());
out.extend(self.encode_response_event(
"response.output_item.done",
json!({
"type": "response.output_item.done",
"response_id": self.response_id(),
"output_index": output_index,
"item": item,
}),
)?);
Ok(out)
}
fn completed_response(&self, usage: CanonicalUsage) -> Value {
let mut ordered_output = Vec::new();
let summary = if self.reasoning_summary_parts.is_empty() {
@@ -2058,6 +2197,9 @@ impl OpenAIResponsesClientEmitter {
ordered_output.push((output_index, Value::Object(item)));
}
}
for (output_index, item) in &self.image_generation_items {
ordered_output.push((*output_index, item.clone()));
}
ordered_output.sort_by_key(|(output_index, _)| *output_index);
let mut usage_payload = Map::new();
@@ -2183,6 +2325,9 @@ impl OpenAIResponsesClientEmitter {
)?);
Ok(out)
}
CanonicalStreamEvent::ImageGenerationCall { index, item } => {
self.emit_image_generation_call_item(index, item)
}
CanonicalStreamEvent::ToolCallStart {
index,
call_id,
@@ -2876,6 +3021,134 @@ mod tests {
)));
}
#[test]
fn openai_responses_provider_state_preserves_image_generation_calls() {
let mut state = OpenAIResponsesProviderState::default();
let report_context = json!({});
let frames = state
.push_line(
&report_context,
data_line(json!({
"type": "response.completed",
"response": {
"id": "resp_img_123",
"model": "gpt-image-2",
"output": [{
"id": "ig_123",
"type": "image_generation_call",
"status": "completed",
"output_format": "png",
"result": "aGVsbG8="
}],
"usage": {"input_tokens": 1, "output_tokens": 2, "total_tokens": 3}
}
})),
)
.expect("completed event should parse");
assert!(frames.iter().any(|frame| matches!(
frame.event,
CanonicalStreamEvent::ImageGenerationCall {
index: 0,
ref item,
} if item["type"] == json!("image_generation_call")
&& item["result"] == json!("aGVsbG8=")
)));
}
#[test]
fn openai_responses_provider_state_waits_for_final_image_generation_item() {
let mut state = OpenAIResponsesProviderState::default();
let report_context = json!({});
let added_frames = state
.push_line(
&report_context,
data_line(json!({
"type": "response.output_item.added",
"output_index": 0,
"item": {
"id": "ig_123",
"type": "image_generation_call",
"status": "generating",
"output_format": "png",
"result": "early"
}
})),
)
.expect("added event should parse");
assert!(!added_frames.iter().any(|frame| matches!(
frame.event,
CanonicalStreamEvent::ImageGenerationCall { .. }
)));
let done_frames = state
.push_line(
&report_context,
data_line(json!({
"type": "response.output_item.done",
"output_index": 0,
"item": {
"id": "ig_123",
"type": "image_generation_call",
"status": "completed",
"output_format": "png",
"result": "final"
}
})),
)
.expect("done event should parse");
assert!(done_frames.iter().any(|frame| matches!(
frame.event,
CanonicalStreamEvent::ImageGenerationCall {
index: 0,
ref item,
} if item["status"] == json!("completed") && item["result"] == json!("final")
)));
}
#[test]
fn openai_responses_client_emitter_emits_image_generation_call_events() {
let mut emitter = OpenAIResponsesClientEmitter::default();
let mut bytes = emitter
.emit(CanonicalStreamFrame {
id: "resp_img_123".to_string(),
model: "gpt-image-2".to_string(),
event: CanonicalStreamEvent::ImageGenerationCall {
index: 0,
item: json!({
"id": "ig_123",
"type": "image_generation_call",
"status": "completed",
"output_format": "png",
"result": "aGVsbG8="
}),
},
})
.expect("image event should encode");
bytes.extend(
emitter
.emit(CanonicalStreamFrame {
id: "resp_img_123".to_string(),
model: "gpt-image-2".to_string(),
event: CanonicalStreamEvent::Finish {
finish_reason: Some("stop".to_string()),
usage: None,
},
})
.expect("finish should encode"),
);
let sse = String::from_utf8(bytes).expect("sse should be utf8");
assert!(sse.contains("event: response.output_item.done\n"));
assert!(sse.contains("\"type\":\"image_generation_call\""));
assert!(sse.contains("\"result\":\"aGVsbG8=\""));
assert!(sse.contains("\"output\":["));
assert!(sse.contains("\"id\":\"ig_123\""));
}
#[test]
fn openai_responses_client_emitter_emits_function_call_output_events() {
let mut emitter = OpenAIResponsesClientEmitter::default();