[router] add reasoning and tool parser argument in router (#11290)

This commit is contained in:
Simo Lin
2025-10-07 09:08:32 -04:00
committed by GitHub
parent 1519a89cfd
commit 79d3495177
13 changed files with 188 additions and 20 deletions

View File

@@ -53,6 +53,8 @@ pub struct GrpcPDRouter {
dp_aware: bool,
api_key: Option<String>,
retry_config: RetryConfig,
configured_reasoning_parser: Option<String>,
configured_tool_parser: Option<String>,
}
impl GrpcPDRouter {
@@ -88,6 +90,8 @@ impl GrpcPDRouter {
dp_aware: ctx.router_config.dp_aware,
api_key: ctx.router_config.api_key.clone(),
retry_config: ctx.router_config.effective_retry_config(),
configured_reasoning_parser: ctx.configured_reasoning_parser.clone(),
configured_tool_parser: ctx.configured_tool_parser.clone(),
})
}
@@ -1179,9 +1183,13 @@ impl GrpcPDRouter {
created: u64,
) -> (String, Option<ChatCompletionStreamResponse>, bool) {
// Get or create parser for this index
reasoning_parsers
.entry(index)
.or_insert_with(|| self.reasoning_parser_factory.get_pooled(model));
reasoning_parsers.entry(index).or_insert_with(|| {
utils::get_reasoning_parser(
&self.reasoning_parser_factory,
self.configured_reasoning_parser.as_ref(),
model,
)
});
if let Some(pooled_parser) = reasoning_parsers.get(&index) {
let (parse_result, in_reasoning) = {
@@ -1248,9 +1256,13 @@ impl GrpcPDRouter {
let mut chunks = Vec::new();
// Get or create parser for this index
tool_parsers
.entry(index)
.or_insert_with(|| self.tool_parser_factory.get_pooled(model));
tool_parsers.entry(index).or_insert_with(|| {
utils::get_tool_parser(
&self.tool_parser_factory,
self.configured_tool_parser.as_ref(),
model,
)
});
if let Some(pooled_parser) = tool_parsers.get(&index) {
let mut parser = pooled_parser.lock().await;
@@ -1737,9 +1749,11 @@ impl GrpcPDRouter {
// Check if reasoning parsing is enabled and separate_reasoning is requested
if original_request.separate_reasoning {
let pooled_parser = self
.reasoning_parser_factory
.get_pooled(&original_request.model);
let pooled_parser = utils::get_reasoning_parser(
&self.reasoning_parser_factory,
self.configured_reasoning_parser.as_ref(),
&original_request.model,
);
let mut parser = pooled_parser
.lock()
@@ -1860,7 +1874,11 @@ impl GrpcPDRouter {
history_tool_calls_count: usize,
) -> (Option<Vec<ToolCall>>, String) {
// Get pooled parser for this model
let pooled_parser = self.tool_parser_factory.get_pooled(model);
let pooled_parser = utils::get_tool_parser(
&self.tool_parser_factory,
self.configured_tool_parser.as_ref(),
model,
);
// Check format detection first
let can_parse = {

View File

@@ -53,6 +53,8 @@ pub struct GrpcRouter {
dp_aware: bool,
api_key: Option<String>,
retry_config: RetryConfig,
configured_reasoning_parser: Option<String>,
configured_tool_parser: Option<String>,
}
impl GrpcRouter {
@@ -87,6 +89,8 @@ impl GrpcRouter {
dp_aware: ctx.router_config.dp_aware,
api_key: ctx.router_config.api_key.clone(),
retry_config: ctx.router_config.effective_retry_config(),
configured_reasoning_parser: ctx.configured_reasoning_parser.clone(),
configured_tool_parser: ctx.configured_tool_parser.clone(),
})
}
@@ -301,7 +305,11 @@ impl GrpcRouter {
history_tool_calls_count: usize,
) -> (Option<Vec<ToolCall>>, String) {
// Get pooled parser for this model
let pooled_parser = self.tool_parser_factory.get_pooled(model);
let pooled_parser = utils::get_tool_parser(
&self.tool_parser_factory,
self.configured_tool_parser.as_ref(),
model,
);
// Check format detection first
let can_parse = {
@@ -496,9 +504,13 @@ impl GrpcRouter {
created: u64,
) -> (String, Option<ChatCompletionStreamResponse>, bool) {
// Get or create parser for this index
reasoning_parsers
.entry(index)
.or_insert_with(|| self.reasoning_parser_factory.get_pooled(model));
reasoning_parsers.entry(index).or_insert_with(|| {
utils::get_reasoning_parser(
&self.reasoning_parser_factory,
self.configured_reasoning_parser.as_ref(),
model,
)
});
if let Some(pooled_parser) = reasoning_parsers.get(&index) {
let (parse_result, in_reasoning) = {
@@ -569,9 +581,13 @@ impl GrpcRouter {
let mut chunks = Vec::new();
// Get or create parser for this index
tool_parsers
.entry(index)
.or_insert_with(|| self.tool_parser_factory.get_pooled(model));
tool_parsers.entry(index).or_insert_with(|| {
utils::get_tool_parser(
&self.tool_parser_factory,
self.configured_tool_parser.as_ref(),
model,
)
});
if let Some(pooled_parser) = tool_parsers.get(&index) {
let mut parser = pooled_parser.lock().await;
@@ -1615,9 +1631,11 @@ impl GrpcRouter {
// Check if reasoning parsing is enabled and separate_reasoning is requested
if original_request.separate_reasoning {
let pooled_parser = self
.reasoning_parser_factory
.get_pooled(&original_request.model);
let pooled_parser = utils::get_reasoning_parser(
&self.reasoning_parser_factory,
self.configured_reasoning_parser.as_ref(),
&original_request.model,
);
let mut parser = pooled_parser
.lock()

View File

@@ -641,6 +641,64 @@ pub fn generate_tool_call_id(
}
}
/// Get the appropriate reasoning parser for a model
///
/// If a parser name is explicitly configured, use that parser.
/// Otherwise, auto-detect based on the model name.
pub fn get_reasoning_parser(
reasoning_parser_factory: &crate::reasoning_parser::ReasoningParserFactory,
configured_parser: Option<&String>,
model: &str,
) -> crate::reasoning_parser::PooledParser {
use tracing::warn;
if let Some(parser_name) = configured_parser {
// Use configured parser if specified
reasoning_parser_factory
.registry()
.get_pooled_parser(parser_name)
.unwrap_or_else(|| {
warn!(
"Configured reasoning parser '{}' not found, falling back to model-based selection",
parser_name
);
reasoning_parser_factory.get_pooled(model)
})
} else {
// Auto-detect based on model
reasoning_parser_factory.get_pooled(model)
}
}
/// Get the appropriate tool parser for a model
///
/// If a parser name is explicitly configured, use that parser.
/// Otherwise, auto-detect based on the model name.
pub fn get_tool_parser(
tool_parser_factory: &crate::tool_parser::ToolParserFactory,
configured_parser: Option<&String>,
model: &str,
) -> crate::tool_parser::PooledToolParser {
use tracing::warn;
if let Some(parser_name) = configured_parser {
// Use configured parser if specified
tool_parser_factory
.registry()
.get_pooled_parser(parser_name)
.unwrap_or_else(|| {
warn!(
"Configured tool parser '{}' not found, falling back to model-based selection",
parser_name
);
tool_parser_factory.get_pooled(model)
})
} else {
// Auto-detect based on model
tool_parser_factory.get_pooled(model)
}
}
#[cfg(test)]
mod tests {
use super::*;