minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[router][grpc] Replace Vec<(String, String, String)> with ExtractedToolCall (#16598)

Browse Source
This commit is contained in:
Chang Su
2026-01-06 16:32:59 -08:00
committed by GitHub
parent 4f443f445a
commit 05b54b6d7b
3 changed files with 73 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/routers/grpc/regular/responses/common.rs
+17 -11
View File
@@ -114,10 +114,18 @@ pub(super) fn prepare_chat_tools_and_choice(
}; };
} }
/// Tool call extracted from a ChatCompletionResponse
#[derive(Debug, Clone)]
pub(super) struct ExtractedToolCall {
pub call_id: String,
pub name: String,
pub arguments: String,
}
/// Extract all tool calls from chat response (for parallel tool call support) /// Extract all tool calls from chat response (for parallel tool call support)
pub(super) fn extract_all_tool_calls_from_chat( pub(super) fn extract_all_tool_calls_from_chat(
response: &crate::protocols::chat::ChatCompletionResponse, response: &crate::protocols::chat::ChatCompletionResponse,
) -> Vec<(String, String, String)> { ) -> Vec<ExtractedToolCall> {
// Check if response has choices with tool calls // Check if response has choices with tool calls
let Some(choice) = response.choices.first() else { let Some(choice) = response.choices.first() else {
return Vec::new(); return Vec::new();
@@ -128,16 +136,14 @@ pub(super) fn extract_all_tool_calls_from_chat(
if let Some(tool_calls) = &message.tool_calls { if let Some(tool_calls) = &message.tool_calls {
tool_calls tool_calls
.iter() .iter()
.map(|tool_call| { .map(|tool_call| ExtractedToolCall {
( call_id: tool_call.id.clone(),
tool_call.id.clone(), name: tool_call.function.name.clone(),
tool_call.function.name.clone(), arguments: tool_call
tool_call .function
.function .arguments
.arguments .clone()
.clone() .unwrap_or_else(|| "{}".to_string()),
.unwrap_or_else(|| "{}".to_string()),
)
}) })
.collect() .collect()
} else { } else {
sgl-model-gateway/src/routers/grpc/regular/responses/non_streaming.rs
+15 -14
View File
@@ -15,7 +15,7 @@ use super::{
common::{ common::{
build_mcp_list_tools_item, build_next_request, convert_mcp_tools_to_chat_tools, build_mcp_list_tools_item, build_next_request, convert_mcp_tools_to_chat_tools,
extract_all_tool_calls_from_chat, load_conversation_history, prepare_chat_tools_and_choice, extract_all_tool_calls_from_chat, load_conversation_history, prepare_chat_tools_and_choice,
ToolLoopState, ExtractedToolCall, ToolLoopState,
}, },
conversions, conversions,
}; };
@@ -232,9 +232,10 @@ pub(super) async fn execute_tool_loop(
// Separate MCP and function tool calls // Separate MCP and function tool calls
let mcp_tool_names: std::collections::HashSet<&str> = let mcp_tool_names: std::collections::HashSet<&str> =
mcp_tools.iter().map(|t| t.name.as_ref()).collect(); mcp_tools.iter().map(|t| t.name.as_ref()).collect();
let (mcp_tool_calls, function_tool_calls): (Vec<_>, Vec<_>) = tool_calls let (mcp_tool_calls, function_tool_calls): (Vec<ExtractedToolCall>, Vec<_>) =
.into_iter() tool_calls
.partition(|(_, tool_name, _)| mcp_tool_names.contains(tool_name.as_str())); .into_iter()
.partition(|tc| mcp_tool_names.contains(tc.name.as_str()));
trace!( trace!(
"Separated tool calls: {} MCP, {} function", "Separated tool calls: {} MCP, {} function",
@@ -312,18 +313,18 @@ pub(super) async fn execute_tool_loop(
} }
// Execute all MCP tools // Execute all MCP tools
for (call_id, tool_name, args_json_str) in mcp_tool_calls { for tool_call in mcp_tool_calls {
trace!( trace!(
"Calling MCP tool '{}' (call_id: {}) with args: {}", "Calling MCP tool '{}' (call_id: {}) with args: {}",
tool_name, tool_call.name,
call_id, tool_call.call_id,
args_json_str tool_call.arguments
); );
let tool_start = Instant::now(); let tool_start = Instant::now();
let (output_str, success, error) = match ctx let (output_str, success, error) = match ctx
.mcp_manager .mcp_manager
.call_tool(tool_name.as_str(), args_json_str.as_str()) .call_tool(tool_call.name.as_str(), tool_call.arguments.as_str())
.await .await
{ {
Ok(result) => match serde_json::to_string(&result) { Ok(result) => match serde_json::to_string(&result) {
@@ -348,12 +349,12 @@ pub(super) async fn execute_tool_loop(
// Record MCP tool metrics // Record MCP tool metrics
Metrics::record_mcp_tool_duration( Metrics::record_mcp_tool_duration(
&current_request.model, &current_request.model,
&tool_name, &tool_call.name,
tool_duration, tool_duration,
); );
Metrics::record_mcp_tool_call( Metrics::record_mcp_tool_call(
&current_request.model, &current_request.model,
&tool_name, &tool_call.name,
if success { if success {
metrics_labels::RESULT_SUCCESS metrics_labels::RESULT_SUCCESS
} else { } else {
@@ -363,9 +364,9 @@ pub(super) async fn execute_tool_loop(
// Record the call in state // Record the call in state
state.record_call( state.record_call(
call_id, tool_call.call_id,
tool_name, tool_call.name,
args_json_str, tool_call.arguments,
output_str, output_str,
success, success,
error, error,
sgl-model-gateway/src/routers/grpc/regular/responses/streaming.rs
+41 -34
View File
@@ -27,7 +27,7 @@ use uuid::Uuid;
use super::{ use super::{
common::{ common::{
build_next_request, convert_mcp_tools_to_chat_tools, extract_all_tool_calls_from_chat, build_next_request, convert_mcp_tools_to_chat_tools, extract_all_tool_calls_from_chat,
prepare_chat_tools_and_choice, ToolLoopState, prepare_chat_tools_and_choice, ExtractedToolCall, ToolLoopState,
}, },
conversions, conversions,
}; };
@@ -630,9 +630,10 @@ async fn execute_tool_loop_streaming_internal(
// Separate MCP and function tool calls // Separate MCP and function tool calls
let mcp_tool_names: std::collections::HashSet<&str> = let mcp_tool_names: std::collections::HashSet<&str> =
mcp_tools.iter().map(|t| t.name.as_ref()).collect(); mcp_tools.iter().map(|t| t.name.as_ref()).collect();
let (mcp_tool_calls, function_tool_calls): (Vec<_>, Vec<_>) = tool_calls let (mcp_tool_calls, function_tool_calls): (Vec<ExtractedToolCall>, Vec<_>) =
.into_iter() tool_calls
.partition(|(_, tool_name, _)| mcp_tool_names.contains(tool_name.as_str())); .into_iter()
.partition(|tc| mcp_tool_names.contains(tc.name.as_str()));
trace!( trace!(
"Separated tool calls: {} MCP, {} function", "Separated tool calls: {} MCP, {} function",
@@ -659,15 +660,15 @@ async fn execute_tool_loop_streaming_internal(
} }
// Process each MCP tool call // Process each MCP tool call
for (call_id, tool_name, args_json_str) in mcp_tool_calls { for tool_call in mcp_tool_calls {
state.total_calls += 1; state.total_calls += 1;
trace!( trace!(
"Executing tool call {}/{}: {} (call_id: {})", "Executing tool call {}/{}: {} (call_id: {})",
state.total_calls, state.total_calls,
state.total_calls, state.total_calls,
tool_name, tool_call.name,
call_id tool_call.call_id
); );
// Allocate output_index for this mcp_call item // Allocate output_index for this mcp_call item
@@ -678,7 +679,7 @@ async fn execute_tool_loop_streaming_internal(
let item = json!({ let item = json!({
"id": item_id, "id": item_id,
"type": "mcp_call", "type": "mcp_call",
"name": tool_name, "name": tool_call.name,
"server_label": state.server_label, "server_label": state.server_label,
"status": "in_progress", "status": "in_progress",
"arguments": "" "arguments": ""
@@ -693,25 +694,31 @@ async fn execute_tool_loop_streaming_internal(
emitter.send_event(&event, &tx)?; emitter.send_event(&event, &tx)?;
// Emit mcp_call_arguments.delta (simulate streaming by sending full arguments) // Emit mcp_call_arguments.delta (simulate streaming by sending full arguments)
let event = let event = emitter.emit_mcp_call_arguments_delta(
emitter.emit_mcp_call_arguments_delta(output_index, &item_id, &args_json_str); output_index,
&item_id,
&tool_call.arguments,
);
emitter.send_event(&event, &tx)?; emitter.send_event(&event, &tx)?;
// Emit mcp_call_arguments.done // Emit mcp_call_arguments.done
let event = let event = emitter.emit_mcp_call_arguments_done(
emitter.emit_mcp_call_arguments_done(output_index, &item_id, &args_json_str); output_index,
&item_id,
&tool_call.arguments,
);
emitter.send_event(&event, &tx)?; emitter.send_event(&event, &tx)?;
// Execute the MCP tool - manager handles parsing and type coercion // Execute the MCP tool - manager handles parsing and type coercion
trace!( trace!(
"Calling MCP tool '{}' with args: {}", "Calling MCP tool '{}' with args: {}",
tool_name, tool_call.name,
args_json_str tool_call.arguments
); );
let tool_start = Instant::now(); let tool_start = Instant::now();
let (output_str, success, error) = match ctx let (output_str, success, error) = match ctx
.mcp_manager .mcp_manager
.call_tool(tool_name.as_str(), args_json_str.as_str()) .call_tool(tool_call.name.as_str(), tool_call.arguments.as_str())
.await .await
{ {
Ok(result) => match serde_json::to_string(&result) { Ok(result) => match serde_json::to_string(&result) {
@@ -724,10 +731,10 @@ async fn execute_tool_loop_streaming_internal(
let item_done = json!({ let item_done = json!({
"id": item_id, "id": item_id,
"type": "mcp_call", "type": "mcp_call",
"name": tool_name, "name": tool_call.name,
"server_label": state.server_label, "server_label": state.server_label,
"status": "completed", "status": "completed",
"arguments": args_json_str, "arguments": tool_call.arguments,
"output": output "output": output
}); });
@@ -749,10 +756,10 @@ async fn execute_tool_loop_streaming_internal(
let item_done = json!({ let item_done = json!({
"id": item_id, "id": item_id,
"type": "mcp_call", "type": "mcp_call",
"name": tool_name, "name": tool_call.name,
"server_label": state.server_label, "server_label": state.server_label,
"status": "failed", "status": "failed",
"arguments": args_json_str, "arguments": tool_call.arguments,
"error": &err "error": &err
}); });
@@ -776,10 +783,10 @@ async fn execute_tool_loop_streaming_internal(
let item_done = json!({ let item_done = json!({
"id": item_id, "id": item_id,
"type": "mcp_call", "type": "mcp_call",
"name": tool_name, "name": tool_call.name,
"server_label": state.server_label, "server_label": state.server_label,
"status": "failed", "status": "failed",
"arguments": args_json_str, "arguments": tool_call.arguments,
"error": &err_str "error": &err_str
}); });
@@ -795,10 +802,10 @@ async fn execute_tool_loop_streaming_internal(
let tool_duration = tool_start.elapsed(); let tool_duration = tool_start.elapsed();
// Record MCP tool metrics // Record MCP tool metrics
Metrics::record_mcp_tool_duration(&model, &tool_name, tool_duration); Metrics::record_mcp_tool_duration(&model, &tool_call.name, tool_duration);
Metrics::record_mcp_tool_call( Metrics::record_mcp_tool_call(
&model, &model,
&tool_name, &tool_call.name,
if success { if success {
metrics_labels::RESULT_SUCCESS metrics_labels::RESULT_SUCCESS
} else { } else {
@@ -808,9 +815,9 @@ async fn execute_tool_loop_streaming_internal(
// Record the call in state // Record the call in state
state.record_call( state.record_call(
call_id, tool_call.call_id,
tool_name, tool_call.name,
args_json_str, tool_call.arguments,
output_str, output_str,
success, success,
error, error,
@@ -825,7 +832,7 @@ async fn execute_tool_loop_streaming_internal(
); );
// Emit function_tool_call events for each function tool // Emit function_tool_call events for each function tool
for (call_id, tool_name, args_json_str) in function_tool_calls { for tool_call in function_tool_calls {
// Allocate output_index for this function_tool_call item // Allocate output_index for this function_tool_call item
let (output_index, item_id) = let (output_index, item_id) =
emitter.allocate_output_index(OutputItemType::FunctionCall); emitter.allocate_output_index(OutputItemType::FunctionCall);
@@ -834,8 +841,8 @@ async fn execute_tool_loop_streaming_internal(
let item = json!({ let item = json!({
"id": item_id, "id": item_id,
"type": "function_tool_call", "type": "function_tool_call",
"call_id": call_id, "call_id": tool_call.call_id,
"name": tool_name, "name": tool_call.name,
"status": "in_progress", "status": "in_progress",
"arguments": "" "arguments": ""
}); });
@@ -848,7 +855,7 @@ async fn execute_tool_loop_streaming_internal(
let event = emitter.emit_function_call_arguments_delta( let event = emitter.emit_function_call_arguments_delta(
output_index, output_index,
&item_id, &item_id,
&args_json_str, &tool_call.arguments,
); );
emitter.send_event(&event, &tx)?; emitter.send_event(&event, &tx)?;
@@ -856,7 +863,7 @@ async fn execute_tool_loop_streaming_internal(
let event = emitter.emit_function_call_arguments_done( let event = emitter.emit_function_call_arguments_done(
output_index, output_index,
&item_id, &item_id,
&args_json_str, &tool_call.arguments,
); );
emitter.send_event(&event, &tx)?; emitter.send_event(&event, &tx)?;
@@ -864,10 +871,10 @@ async fn execute_tool_loop_streaming_internal(
let item_complete = json!({ let item_complete = json!({
"id": item_id, "id": item_id,
"type": "function_tool_call", "type": "function_tool_call",
"call_id": call_id, "call_id": tool_call.call_id,
"name": tool_name, "name": tool_call.name,
"status": "completed", "status": "completed",
"arguments": args_json_str "arguments": tool_call.arguments
}); });
// Emit output_item.done // Emit output_item.done