[model-gateway] code clean up on oai router in responses (#14852)
This commit is contained in:
Simo Lin
@@ -1,8 +1,4 @@
|
|||||||
//! Response storage, patching, and extraction utilities
|
use serde_json::{json, Map, Value};
|
||||||
|
|
||||||
use std::collections::HashMap;
|
|
||||||
|
|
||||||
use serde_json::{json, Value};
|
|
||||||
use tracing::warn;
|
use tracing::warn;
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
@@ -13,247 +9,154 @@ use crate::{
|
|||||||
},
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
// ============================================================================
|
/// Extract a string field from JSON, returning owned String
|
||||||
// Response Storage Operations
|
fn get_string(json: &Value, key: &str) -> Option<String> {
|
||||||
// ============================================================================
|
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<F>(obj: &mut Map<String, Value>, key: &str, value: &str, condition: F)
|
||||||
|
where
|
||||||
|
F: FnOnce(&Map<String, Value>) -> bool,
|
||||||
|
{
|
||||||
|
if condition(obj) {
|
||||||
|
obj.insert(key.to_string(), Value::String(value.to_string()));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Build a StoredResponse from response JSON and original request
|
/// Build a StoredResponse from response JSON and original request
|
||||||
pub(super) fn build_stored_response(
|
pub(super) fn build_stored_response(
|
||||||
response_json: &Value,
|
response_json: &Value,
|
||||||
original_body: &ResponsesRequest,
|
original_body: &ResponsesRequest,
|
||||||
) -> StoredResponse {
|
) -> 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
|
// Initialize empty arrays - will be populated by persist_items_with_storages
|
||||||
stored_response.input = Value::Array(vec![]);
|
stored.input = Value::Array(vec![]);
|
||||||
stored_response.output = Value::Array(vec![]);
|
stored.output = Value::Array(vec![]);
|
||||||
|
|
||||||
stored_response.instructions = response_json
|
stored.instructions =
|
||||||
.get("instructions")
|
get_string(response_json, "instructions").or_else(|| original_body.instructions.clone());
|
||||||
.and_then(|v| v.as_str())
|
|
||||||
.map(|s| s.to_string())
|
|
||||||
.or_else(|| original_body.instructions.clone());
|
|
||||||
|
|
||||||
stored_response.model = response_json
|
stored.model = get_string(response_json, "model").or_else(|| Some(original_body.model.clone()));
|
||||||
.get("model")
|
|
||||||
.and_then(|v| v.as_str())
|
|
||||||
.map(|s| s.to_string())
|
|
||||||
.or_else(|| Some(original_body.model.clone()));
|
|
||||||
|
|
||||||
if let Some(safety_identifier) = original_body.user.clone() {
|
stored.safety_identifier = original_body.user.clone();
|
||||||
stored_response.safety_identifier = Some(safety_identifier);
|
stored.conversation_id = original_body.conversation.clone();
|
||||||
}
|
|
||||||
|
|
||||||
// Set conversation id from request if provided
|
stored.metadata = response_json
|
||||||
if let Some(conv_id) = original_body.conversation.clone() {
|
|
||||||
stored_response.conversation_id = Some(conv_id);
|
|
||||||
}
|
|
||||||
|
|
||||||
stored_response.metadata = response_json
|
|
||||||
.get("metadata")
|
.get("metadata")
|
||||||
.and_then(|v| v.as_object())
|
.and_then(|v| v.as_object())
|
||||||
.map(|m| {
|
.map(|m| m.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
|
||||||
m.iter()
|
|
||||||
.map(|(k, v)| (k.clone(), v.clone()))
|
|
||||||
.collect::<HashMap<_, _>>()
|
|
||||||
})
|
|
||||||
.unwrap_or_else(|| original_body.metadata.clone().unwrap_or_default());
|
.unwrap_or_else(|| original_body.metadata.clone().unwrap_or_default());
|
||||||
|
|
||||||
stored_response.previous_response_id = response_json
|
stored.previous_response_id = get_string(response_json, "previous_response_id")
|
||||||
.get("previous_response_id")
|
.map(|s| ResponseId::from(s.as_str()))
|
||||||
.and_then(|v| v.as_str())
|
|
||||||
.map(ResponseId::from)
|
|
||||||
.or_else(|| {
|
.or_else(|| {
|
||||||
original_body
|
original_body
|
||||||
.previous_response_id
|
.previous_response_id
|
||||||
.as_ref()
|
.as_deref()
|
||||||
.map(|id| ResponseId::from(id.as_str()))
|
.map(ResponseId::from)
|
||||||
});
|
});
|
||||||
|
|
||||||
if let Some(id_str) = response_json.get("id").and_then(|v| v.as_str()) {
|
if let Some(id_str) = get_string(response_json, "id") {
|
||||||
stored_response.id = ResponseId::from(id_str);
|
stored.id = ResponseId::from(id_str.as_str());
|
||||||
}
|
}
|
||||||
|
|
||||||
stored_response.raw_response = response_json.clone();
|
stored.raw_response = response_json.clone();
|
||||||
|
stored
|
||||||
stored_response
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// ============================================================================
|
|
||||||
// Response JSON Patching
|
|
||||||
// ============================================================================
|
|
||||||
|
|
||||||
/// Patch streaming response JSON with metadata from original request
|
/// Patch streaming response JSON with metadata from original request
|
||||||
pub(super) fn patch_streaming_response_json(
|
pub(super) fn patch_streaming_response_json(
|
||||||
response_json: &mut Value,
|
response_json: &mut Value,
|
||||||
original_body: &ResponsesRequest,
|
original_body: &ResponsesRequest,
|
||||||
original_previous_response_id: Option<&str>,
|
original_previous_response_id: Option<&str>,
|
||||||
) {
|
) {
|
||||||
if let Some(obj) = response_json.as_object_mut() {
|
let Some(obj) = response_json.as_object_mut() else {
|
||||||
if let Some(prev_id) = original_previous_response_id {
|
return;
|
||||||
let should_insert = obj
|
};
|
||||||
.get("previous_response_id")
|
|
||||||
.map(|v| v.is_null() || v.as_str().map(|s| s.is_empty()).unwrap_or(false))
|
// Set previous_response_id if missing/empty
|
||||||
.unwrap_or(true);
|
if let Some(prev_id) = original_previous_response_id {
|
||||||
if should_insert {
|
insert_if(obj, "previous_response_id", prev_id, |o| {
|
||||||
obj.insert(
|
is_missing_or_empty(o.get("previous_response_id"))
|
||||||
"previous_response_id".to_string(),
|
});
|
||||||
Value::String(prev_id.to_string()),
|
}
|
||||||
);
|
|
||||||
}
|
// 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<String, Value> = metadata
|
||||||
|
.iter()
|
||||||
|
.map(|(k, v)| (k.clone(), v.clone()))
|
||||||
|
.collect();
|
||||||
|
obj.insert("metadata".to_string(), Value::Object(metadata_map));
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
if !obj.contains_key("instructions")
|
// Always set store
|
||||||
|| obj
|
obj.insert(
|
||||||
.get("instructions")
|
"store".to_string(),
|
||||||
.map(|v| v.is_null())
|
Value::Bool(original_body.store.unwrap_or(false)),
|
||||||
.unwrap_or(false)
|
);
|
||||||
{
|
|
||||||
if let Some(instructions) = &original_body.instructions {
|
|
||||||
obj.insert(
|
|
||||||
"instructions".to_string(),
|
|
||||||
Value::String(instructions.clone()),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if !obj.contains_key("metadata")
|
// Set model if missing/empty
|
||||||
|| obj.get("metadata").map(|v| v.is_null()).unwrap_or(false)
|
insert_if(obj, "model", &original_body.model, |o| {
|
||||||
{
|
is_missing_or_empty(o.get("model"))
|
||||||
if let Some(metadata) = &original_body.metadata {
|
});
|
||||||
let metadata_map: serde_json::Map<String, Value> = 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 safety_identifier if null (but key exists)
|
||||||
|
if let Some(user) = &original_body.user {
|
||||||
if obj
|
if obj
|
||||||
.get("safety_identifier")
|
.get("safety_identifier")
|
||||||
.map(|v| v.is_null())
|
.is_some_and(|v: &Value| v.is_null())
|
||||||
.unwrap_or(false)
|
|
||||||
{
|
{
|
||||||
if let Some(safety_identifier) = &original_body.user {
|
obj.insert("safety_identifier".to_string(), Value::String(user.clone()));
|
||||||
obj.insert(
|
|
||||||
"safety_identifier".to_string(),
|
|
||||||
Value::String(safety_identifier.clone()),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Attach conversation id for client response if present (final aggregated JSON)
|
// Attach conversation id for client response
|
||||||
if let Some(conv_id) = original_body.conversation.clone() {
|
if let Some(conv_id) = &original_body.conversation {
|
||||||
obj.insert("conversation".to_string(), json!({ "id": conv_id }));
|
obj.insert("conversation".to_string(), json!({ "id": conv_id }));
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Rewrite streaming SSE block to include metadata from original request
|
/// Extract data payload from SSE block lines
|
||||||
pub(super) fn rewrite_streaming_block(
|
fn extract_sse_data(block: &str) -> Option<String> {
|
||||||
block: &str,
|
let data_lines: Vec<_> = block
|
||||||
original_body: &ResponsesRequest,
|
.lines()
|
||||||
original_previous_response_id: Option<&str>,
|
.filter(|line| line.starts_with("data:"))
|
||||||
) -> Option<String> {
|
.map(|line| line.trim_start_matches("data:").trim_start())
|
||||||
let trimmed = block.trim();
|
.collect();
|
||||||
if trimmed.is_empty() {
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut data_lines: Vec<String> = Vec::new();
|
|
||||||
|
|
||||||
for line in trimmed.lines() {
|
|
||||||
if line.starts_with("data:") {
|
|
||||||
data_lines.push(line.trim_start_matches("data:").trim_start().to_string());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if data_lines.is_empty() {
|
if data_lines.is_empty() {
|
||||||
return None;
|
None
|
||||||
|
} else {
|
||||||
|
Some(data_lines.join("\n"))
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
let payload = data_lines.join("\n");
|
/// Rebuild SSE block with new data payload
|
||||||
let mut parsed: Value = match serde_json::from_str(&payload) {
|
fn rebuild_sse_block(block: &str, new_payload: &str) -> String {
|
||||||
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;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let mut rebuilt_lines = Vec::new();
|
let mut rebuilt_lines = Vec::new();
|
||||||
let mut data_written = false;
|
let mut data_written = false;
|
||||||
for line in trimmed.lines() {
|
|
||||||
|
for line in block.lines() {
|
||||||
if line.starts_with("data:") {
|
if line.starts_with("data:") {
|
||||||
if !data_written {
|
if !data_written {
|
||||||
rebuilt_lines.push(format!("data: {}", new_payload));
|
rebuilt_lines.push(format!("data: {}", new_payload));
|
||||||
@@ -268,7 +171,74 @@ pub(super) fn rewrite_streaming_block(
|
|||||||
rebuilt_lines.push(format!("data: {}", new_payload));
|
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<String> {
|
||||||
|
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<String, Value>, key: &str, value: &Option<String>) {
|
||||||
|
if let Some(v) = value {
|
||||||
|
map.insert(key.to_string(), Value::String(v.clone()));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Mask function tools as MCP tools in response for client
|
/// 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()
|
.iter()
|
||||||
.find(|t| matches!(t.r#type, ResponseToolType::Mcp) && t.server_url.is_some())
|
.find(|t| matches!(t.r#type, ResponseToolType::Mcp) && t.server_url.is_some())
|
||||||
});
|
});
|
||||||
|
|
||||||
let Some(t) = mcp_tool else {
|
let Some(t) = mcp_tool else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut m = serde_json::Map::new();
|
let mut m = Map::new();
|
||||||
m.insert("type".to_string(), Value::String("mcp".to_string()));
|
m.insert("type".to_string(), json!("mcp"));
|
||||||
if let Some(label) = &t.server_label {
|
insert_optional_string(&mut m, "server_label", &t.server_label);
|
||||||
m.insert("server_label".to_string(), Value::String(label.clone()));
|
insert_optional_string(&mut m, "server_url", &t.server_url);
|
||||||
}
|
insert_optional_string(&mut m, "server_description", &t.server_description);
|
||||||
if let Some(url) = &t.server_url {
|
insert_optional_string(&mut m, "require_approval", &t.require_approval);
|
||||||
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()));
|
|
||||||
}
|
|
||||||
if let Some(allowed) = &t.allowed_tools {
|
if let Some(allowed) = &t.allowed_tools {
|
||||||
m.insert(
|
m.insert(
|
||||||
"allowed_tools".to_string(),
|
"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() {
|
if let Some(obj) = resp.as_object_mut() {
|
||||||
obj.insert("tools".to_string(), Value::Array(vec![Value::Object(m)]));
|
obj.insert("tools".to_string(), json!([Value::Object(m)]));
|
||||||
obj.entry("tool_choice")
|
obj.entry("tool_choice").or_insert(json!("auto"));
|
||||||
.or_insert(Value::String("auto".to_string()));
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user