[model-gateway] optimize Vec and HashMap allocations in responses api (#16406)

This commit is contained in:
Simo Lin
2026-01-04 10:07:55 -08:00
committed by GitHub
parent f16606d6f6
commit cf6800f6cc
2 changed files with 16 additions and 13 deletions
+13 -11
View File
@@ -279,8 +279,7 @@ pub fn generate_tool_constraints(
/// Build JSON schema for required tool calls (array with minItems: 1) /// Build JSON schema for required tool calls (array with minItems: 1)
/// Includes $defs consolidation from all tools (matching Python's behavior) /// Includes $defs consolidation from all tools (matching Python's behavior)
fn build_required_array_schema(tools: &[Tool]) -> Result<String, String> { fn build_required_array_schema(tools: &[Tool]) -> Result<String, String> {
// Build anyOf schemas for each tool let mut any_of_schemas = Vec::with_capacity(tools.len());
let mut any_of_schemas = Vec::new();
for tool in tools { for tool in tools {
let tool_schema = json!({ let tool_schema = json!({
"properties": { "properties": {
@@ -296,7 +295,7 @@ fn build_required_array_schema(tools: &[Tool]) -> Result<String, String> {
} }
// Consolidate $defs from all tools (matching Python's _get_tool_schema_defs) // Consolidate $defs from all tools (matching Python's _get_tool_schema_defs)
let mut all_defs: HashMap<String, Value> = HashMap::new(); let mut all_defs: Map<String, Value> = Map::new();
for tool in tools { for tool in tools {
if let Value::Object(params) = &tool.function.parameters { if let Value::Object(params) = &tool.function.parameters {
if let Some(Value::Object(defs)) = params.get("$defs") { if let Some(Value::Object(defs)) = params.get("$defs") {
@@ -332,8 +331,7 @@ fn build_required_array_schema(tools: &[Tool]) -> Result<String, String> {
// Add $defs if any were found (matching Python's behavior) // Add $defs if any were found (matching Python's behavior)
if !all_defs.is_empty() { if !all_defs.is_empty() {
if let Value::Object(ref mut schema_obj) = array_schema { if let Value::Object(ref mut schema_obj) = array_schema {
let defs_value = Value::Object(all_defs.into_iter().collect::<Map<String, Value>>()); schema_obj.insert("$defs".to_string(), Value::Object(all_defs));
schema_obj.insert("$defs".to_string(), defs_value);
} }
} }
@@ -439,8 +437,8 @@ pub fn process_chat_messages(
.transpose() .transpose()
.map_err(|e| format!("Failed to serialize tools: {}", e))?; .map_err(|e| format!("Failed to serialize tools: {}", e))?;
// Build template kwargs, merging reasoning_effort if present let kwargs_capacity = 1 + request.chat_template_kwargs.as_ref().map_or(0, |k| k.len());
let mut combined_template_kwargs = HashMap::new(); let mut combined_template_kwargs = HashMap::with_capacity(kwargs_capacity);
// Add reasoning_effort if present (like Python does) // Add reasoning_effort if present (like Python does)
if let Some(reasoning_effort) = &request.reasoning_effort { if let Some(reasoning_effort) = &request.reasoning_effort {
@@ -878,7 +876,7 @@ pub fn convert_proto_to_openai_logprobs(
proto_logprobs: &OutputLogProbs, proto_logprobs: &OutputLogProbs,
tokenizer: &Arc<dyn Tokenizer>, tokenizer: &Arc<dyn Tokenizer>,
) -> Result<ChatLogProbs, String> { ) -> Result<ChatLogProbs, String> {
let mut content_items = Vec::new(); let mut content_items = Vec::with_capacity(proto_logprobs.token_logprobs.len());
// Decode token IDs to text (always with skip_special_tokens=false for logprobs) // Decode token IDs to text (always with skip_special_tokens=false for logprobs)
let token_texts: Vec<String> = proto_logprobs let token_texts: Vec<String> = proto_logprobs
@@ -901,8 +899,9 @@ pub fn convert_proto_to_openai_logprobs(
let bytes = Some(token_text.as_bytes().to_vec()); let bytes = Some(token_text.as_bytes().to_vec());
// Build top_logprobs for this position // Build top_logprobs for this position
let mut top_logprobs = Vec::new(); let top_logprobs = if let Some(top_logprobs_entry) = proto_logprobs.top_logprobs.get(i) {
if let Some(top_logprobs_entry) = proto_logprobs.top_logprobs.get(i) { let mut top_logprobs = Vec::with_capacity(top_logprobs_entry.values.len());
// Decode top token IDs (always with skip_special_tokens=false) // Decode top token IDs (always with skip_special_tokens=false)
let top_token_texts: Vec<String> = top_logprobs_entry let top_token_texts: Vec<String> = top_logprobs_entry
.token_ids .token_ids
@@ -928,7 +927,10 @@ pub fn convert_proto_to_openai_logprobs(
}); });
} }
} }
} top_logprobs
} else {
Vec::new()
};
content_items.push(ChatLogProbsContent { content_items.push(ChatLogProbsContent {
token: token_text, token: token_text,
+3 -2
View File
@@ -217,8 +217,8 @@ pub(super) fn prepare_mcp_payload_for_streaming(
} }
// Build function tools for all discovered MCP tools // Build function tools for all discovered MCP tools
let mut tools_json = Vec::new();
let tools = active_mcp.list_tools(); let tools = active_mcp.list_tools();
let mut tools_json = Vec::with_capacity(tools.len());
for t in tools { for t in tools {
let parameters = Value::Object((*t.input_schema).clone()); let parameters = Value::Object((*t.input_schema).clone());
let tool = serde_json::json!({ let tool = serde_json::json!({
@@ -252,7 +252,8 @@ pub(super) fn build_resume_payload(
.ok_or_else(|| "payload not an object".to_string())?; .ok_or_else(|| "payload not an object".to_string())?;
// Build input array: start with original user input // Build input array: start with original user input
let mut input_array = Vec::new(); // Pre-allocate: 1 for user message + conversation history
let mut input_array = Vec::with_capacity(1 + conversation_history.len());
// Add original user message // Add original user message
// For structured input, serialize the original input items // For structured input, serialize the original input items