From 5b0872d2d7e054a99df5e657249f702db00e4ab1 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Thu, 11 Dec 2025 06:55:09 -0800 Subject: [PATCH] [model-gateway] code clean up on oai router in responses (#14852) --- .../src/routers/openai/responses.rs | 398 ++++++++---------- 1 file changed, 179 insertions(+), 219 deletions(-) diff --git a/sgl-model-gateway/src/routers/openai/responses.rs b/sgl-model-gateway/src/routers/openai/responses.rs index c48134920..3b8cb1c76 100644 --- a/sgl-model-gateway/src/routers/openai/responses.rs +++ b/sgl-model-gateway/src/routers/openai/responses.rs @@ -1,8 +1,4 @@ -//! Response storage, patching, and extraction utilities - -use std::collections::HashMap; - -use serde_json::{json, Value}; +use serde_json::{json, Map, Value}; use tracing::warn; use crate::{ @@ -13,247 +9,154 @@ use crate::{ }, }; -// ============================================================================ -// Response Storage Operations -// ============================================================================ +/// Extract a string field from JSON, returning owned String +fn get_string(json: &Value, key: &str) -> Option { + json.get(key).and_then(|v| v.as_str()).map(String::from) +} + +/// Check if a JSON value is missing, null, or an empty string +fn is_missing_or_empty(value: Option<&Value>) -> bool { + match value { + None => true, + Some(v) => v.is_null() || v.as_str().is_some_and(|s| s.is_empty()), + } +} + +/// Insert a string value into a JSON object if the condition is met +fn insert_if(obj: &mut Map, key: &str, value: &str, condition: F) +where + F: FnOnce(&Map) -> bool, +{ + if condition(obj) { + obj.insert(key.to_string(), Value::String(value.to_string())); + } +} /// Build a StoredResponse from response JSON and original request pub(super) fn build_stored_response( response_json: &Value, original_body: &ResponsesRequest, ) -> StoredResponse { - let mut stored_response = StoredResponse::new(None); + let mut stored = StoredResponse::new(None); // Initialize empty arrays - will be populated by persist_items_with_storages - stored_response.input = Value::Array(vec![]); - stored_response.output = Value::Array(vec![]); + stored.input = Value::Array(vec![]); + stored.output = Value::Array(vec![]); - stored_response.instructions = response_json - .get("instructions") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()) - .or_else(|| original_body.instructions.clone()); + stored.instructions = + get_string(response_json, "instructions").or_else(|| original_body.instructions.clone()); - stored_response.model = response_json - .get("model") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()) - .or_else(|| Some(original_body.model.clone())); + stored.model = get_string(response_json, "model").or_else(|| Some(original_body.model.clone())); - if let Some(safety_identifier) = original_body.user.clone() { - stored_response.safety_identifier = Some(safety_identifier); - } + stored.safety_identifier = original_body.user.clone(); + stored.conversation_id = original_body.conversation.clone(); - // Set conversation id from request if provided - if let Some(conv_id) = original_body.conversation.clone() { - stored_response.conversation_id = Some(conv_id); - } - - stored_response.metadata = response_json + stored.metadata = response_json .get("metadata") .and_then(|v| v.as_object()) - .map(|m| { - m.iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect::>() - }) + .map(|m| m.iter().map(|(k, v)| (k.clone(), v.clone())).collect()) .unwrap_or_else(|| original_body.metadata.clone().unwrap_or_default()); - stored_response.previous_response_id = response_json - .get("previous_response_id") - .and_then(|v| v.as_str()) - .map(ResponseId::from) + stored.previous_response_id = get_string(response_json, "previous_response_id") + .map(|s| ResponseId::from(s.as_str())) .or_else(|| { original_body .previous_response_id - .as_ref() - .map(|id| ResponseId::from(id.as_str())) + .as_deref() + .map(ResponseId::from) }); - if let Some(id_str) = response_json.get("id").and_then(|v| v.as_str()) { - stored_response.id = ResponseId::from(id_str); + if let Some(id_str) = get_string(response_json, "id") { + stored.id = ResponseId::from(id_str.as_str()); } - stored_response.raw_response = response_json.clone(); - - stored_response + stored.raw_response = response_json.clone(); + stored } -// ============================================================================ -// Response JSON Patching -// ============================================================================ - /// Patch streaming response JSON with metadata from original request pub(super) fn patch_streaming_response_json( response_json: &mut Value, original_body: &ResponsesRequest, original_previous_response_id: Option<&str>, ) { - if let Some(obj) = response_json.as_object_mut() { - if let Some(prev_id) = original_previous_response_id { - let should_insert = obj - .get("previous_response_id") - .map(|v| v.is_null() || v.as_str().map(|s| s.is_empty()).unwrap_or(false)) - .unwrap_or(true); - if should_insert { - obj.insert( - "previous_response_id".to_string(), - Value::String(prev_id.to_string()), - ); - } + let Some(obj) = response_json.as_object_mut() else { + return; + }; + + // Set previous_response_id if missing/empty + if let Some(prev_id) = original_previous_response_id { + insert_if(obj, "previous_response_id", prev_id, |o| { + is_missing_or_empty(o.get("previous_response_id")) + }); + } + + // Set instructions if missing/null + if let Some(instructions) = &original_body.instructions { + insert_if(obj, "instructions", instructions, |o| { + is_missing_or_empty(o.get("instructions")) + }); + } + + // Set metadata if missing/null + if is_missing_or_empty(obj.get("metadata")) { + if let Some(metadata) = &original_body.metadata { + let metadata_map: Map = metadata + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + obj.insert("metadata".to_string(), Value::Object(metadata_map)); } + } - if !obj.contains_key("instructions") - || obj - .get("instructions") - .map(|v| v.is_null()) - .unwrap_or(false) - { - if let Some(instructions) = &original_body.instructions { - obj.insert( - "instructions".to_string(), - Value::String(instructions.clone()), - ); - } - } + // Always set store + obj.insert( + "store".to_string(), + Value::Bool(original_body.store.unwrap_or(false)), + ); - if !obj.contains_key("metadata") - || obj.get("metadata").map(|v| v.is_null()).unwrap_or(false) - { - if let Some(metadata) = &original_body.metadata { - let metadata_map: serde_json::Map = metadata - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - obj.insert("metadata".to_string(), Value::Object(metadata_map)); - } - } - - obj.insert( - "store".to_string(), - Value::Bool(original_body.store.unwrap_or(false)), - ); - - if obj - .get("model") - .and_then(|v| v.as_str()) - .map(|s| s.is_empty()) - .unwrap_or(true) - { - obj.insert( - "model".to_string(), - Value::String(original_body.model.clone()), - ); - } + // Set model if missing/empty + insert_if(obj, "model", &original_body.model, |o| { + is_missing_or_empty(o.get("model")) + }); + // Set safety_identifier if null (but key exists) + if let Some(user) = &original_body.user { if obj .get("safety_identifier") - .map(|v| v.is_null()) - .unwrap_or(false) + .is_some_and(|v: &Value| v.is_null()) { - if let Some(safety_identifier) = &original_body.user { - obj.insert( - "safety_identifier".to_string(), - Value::String(safety_identifier.clone()), - ); - } + obj.insert("safety_identifier".to_string(), Value::String(user.clone())); } + } - // Attach conversation id for client response if present (final aggregated JSON) - if let Some(conv_id) = original_body.conversation.clone() { - obj.insert("conversation".to_string(), json!({ "id": conv_id })); - } + // Attach conversation id for client response + if let Some(conv_id) = &original_body.conversation { + obj.insert("conversation".to_string(), json!({ "id": conv_id })); } } -/// Rewrite streaming SSE block to include metadata from original request -pub(super) fn rewrite_streaming_block( - block: &str, - original_body: &ResponsesRequest, - original_previous_response_id: Option<&str>, -) -> Option { - let trimmed = block.trim(); - if trimmed.is_empty() { - return None; - } - - let mut data_lines: Vec = Vec::new(); - - for line in trimmed.lines() { - if line.starts_with("data:") { - data_lines.push(line.trim_start_matches("data:").trim_start().to_string()); - } - } +/// Extract data payload from SSE block lines +fn extract_sse_data(block: &str) -> Option { + let data_lines: Vec<_> = block + .lines() + .filter(|line| line.starts_with("data:")) + .map(|line| line.trim_start_matches("data:").trim_start()) + .collect(); if data_lines.is_empty() { - return None; + None + } else { + Some(data_lines.join("\n")) } +} - let payload = data_lines.join("\n"); - let mut parsed: Value = match serde_json::from_str(&payload) { - Ok(value) => value, - Err(err) => { - warn!("Failed to parse streaming JSON payload: {}", err); - return None; - } - }; - - let event_type = parsed - .get("type") - .and_then(|v| v.as_str()) - .unwrap_or_default(); - - let should_patch = is_response_event(event_type); - - if !should_patch { - return None; - } - - let mut changed = false; - if let Some(response_obj) = parsed.get_mut("response").and_then(|v| v.as_object_mut()) { - let desired_store = Value::Bool(original_body.store.unwrap_or(false)); - if response_obj.get("store") != Some(&desired_store) { - response_obj.insert("store".to_string(), desired_store); - changed = true; - } - - if let Some(prev_id) = original_previous_response_id { - let needs_previous = response_obj - .get("previous_response_id") - .map(|v| v.is_null() || v.as_str().map(|s| s.is_empty()).unwrap_or(false)) - .unwrap_or(true); - - if needs_previous { - response_obj.insert( - "previous_response_id".to_string(), - Value::String(prev_id.to_string()), - ); - changed = true; - } - } - - // Attach conversation id into streaming event response content with ordering - if let Some(conv_id) = original_body.conversation.clone() { - response_obj.insert("conversation".to_string(), json!({ "id": conv_id })); - changed = true; - } - } - - if !changed { - return None; - } - - let new_payload = match serde_json::to_string(&parsed) { - Ok(json) => json, - Err(err) => { - warn!("Failed to serialize modified streaming payload: {}", err); - return None; - } - }; - +/// Rebuild SSE block with new data payload +fn rebuild_sse_block(block: &str, new_payload: &str) -> String { let mut rebuilt_lines = Vec::new(); let mut data_written = false; - for line in trimmed.lines() { + + for line in block.lines() { if line.starts_with("data:") { if !data_written { rebuilt_lines.push(format!("data: {}", new_payload)); @@ -268,7 +171,74 @@ pub(super) fn rewrite_streaming_block( rebuilt_lines.push(format!("data: {}", new_payload)); } - Some(rebuilt_lines.join("\n")) + rebuilt_lines.join("\n") +} + +/// Rewrite streaming SSE block to include metadata from original request +pub(super) fn rewrite_streaming_block( + block: &str, + original_body: &ResponsesRequest, + original_previous_response_id: Option<&str>, +) -> Option { + let trimmed = block.trim(); + if trimmed.is_empty() { + return None; + } + + let payload = extract_sse_data(trimmed)?; + let mut parsed: Value = serde_json::from_str(&payload) + .map_err(|e| warn!("Failed to parse streaming JSON payload: {}", e)) + .ok()?; + + let event_type = parsed + .get("type") + .and_then(|v| v.as_str()) + .unwrap_or_default(); + + if !is_response_event(event_type) { + return None; + } + + let response_obj = parsed.get_mut("response").and_then(|v| v.as_object_mut())?; + let mut changed = false; + + // Update store value if different + let desired_store = Value::Bool(original_body.store.unwrap_or(false)); + if response_obj.get("store") != Some(&desired_store) { + response_obj.insert("store".to_string(), desired_store); + changed = true; + } + + // Set previous_response_id if missing/empty + if let Some(prev_id) = original_previous_response_id { + if is_missing_or_empty(response_obj.get("previous_response_id")) { + response_obj.insert("previous_response_id".to_string(), json!(prev_id)); + changed = true; + } + } + + // Attach conversation id + if let Some(conv_id) = &original_body.conversation { + response_obj.insert("conversation".to_string(), json!({ "id": conv_id })); + changed = true; + } + + if !changed { + return None; + } + + let new_payload = serde_json::to_string(&parsed) + .map_err(|e| warn!("Failed to serialize modified streaming payload: {}", e)) + .ok()?; + + Some(rebuild_sse_block(trimmed, &new_payload)) +} + +/// Helper to insert an optional string field into a JSON map +fn insert_optional_string(map: &mut Map, key: &str, value: &Option) { + if let Some(v) = value { + map.insert(key.to_string(), Value::String(v.clone())); + } } /// Mask function tools as MCP tools in response for client @@ -278,37 +248,27 @@ pub(super) fn mask_tools_as_mcp(resp: &mut Value, original_body: &ResponsesReque .iter() .find(|t| matches!(t.r#type, ResponseToolType::Mcp) && t.server_url.is_some()) }); + let Some(t) = mcp_tool else { return; }; - let mut m = serde_json::Map::new(); - m.insert("type".to_string(), Value::String("mcp".to_string())); - if let Some(label) = &t.server_label { - m.insert("server_label".to_string(), Value::String(label.clone())); - } - if let Some(url) = &t.server_url { - m.insert("server_url".to_string(), Value::String(url.clone())); - } - if let Some(desc) = &t.server_description { - m.insert( - "server_description".to_string(), - Value::String(desc.clone()), - ); - } - if let Some(req) = &t.require_approval { - m.insert("require_approval".to_string(), Value::String(req.clone())); - } + let mut m = Map::new(); + m.insert("type".to_string(), json!("mcp")); + insert_optional_string(&mut m, "server_label", &t.server_label); + insert_optional_string(&mut m, "server_url", &t.server_url); + insert_optional_string(&mut m, "server_description", &t.server_description); + insert_optional_string(&mut m, "require_approval", &t.require_approval); + if let Some(allowed) = &t.allowed_tools { m.insert( "allowed_tools".to_string(), - Value::Array(allowed.iter().map(|s| Value::String(s.clone())).collect()), + Value::Array(allowed.iter().map(|s| json!(s)).collect()), ); } if let Some(obj) = resp.as_object_mut() { - obj.insert("tools".to_string(), Value::Array(vec![Value::Object(m)])); - obj.entry("tool_choice") - .or_insert(Value::String("auto".to_string())); + obj.insert("tools".to_string(), json!([Value::Object(m)])); + obj.entry("tool_choice").or_insert(json!("auto")); } }