router-grpc: Add tools processing and other paramters for apply_chat_template (#10877)

This commit is contained in:
Chang Su
2025-09-24 15:23:06 -07:00
committed by GitHub
parent d7b20dd65d
commit 4a87ba217f
6 changed files with 320 additions and 126 deletions

View File

@@ -1,6 +1,7 @@
use sglang_router_rs::protocols::spec;
use sglang_router_rs::tokenizer::chat_template::{
detect_chat_template_content_format, ChatTemplateContentFormat, ChatTemplateProcessor,
detect_chat_template_content_format, ChatTemplateContentFormat, ChatTemplateParams,
ChatTemplateProcessor,
};
#[test]
@@ -169,11 +170,7 @@ assistant:
{%- endif %}
"#;
let processor = ChatTemplateProcessor::new(
template.to_string(),
Some("<s>".to_string()),
Some("</s>".to_string()),
);
let processor = ChatTemplateProcessor::new(template.to_string());
let messages = vec![
spec::ChatMessage::System {
@@ -194,8 +191,12 @@ assistant:
.map(|msg| serde_json::to_value(msg).unwrap())
.collect();
let params = ChatTemplateParams {
add_generation_prompt: true,
..Default::default()
};
let result = processor
.apply_chat_template(&message_values, true)
.apply_chat_template(&message_values, params)
.unwrap();
assert!(result.contains("system: You are helpful"));
assert!(result.contains("user: Hello"));
@@ -204,19 +205,15 @@ assistant:
#[test]
fn test_chat_template_with_tokens_unit_test() {
// Template that uses special tokens
// Template that uses template kwargs for tokens (more realistic)
let template = r#"
{{ bos_token }}
{%- if start_token -%}{{ start_token }}{%- endif -%}
{%- for message in messages -%}
{{ message.role }}: {{ message.content }}{{ eos_token }}
{{ message.role }}: {{ message.content }}{%- if end_token -%}{{ end_token }}{%- endif -%}
{% endfor -%}
"#;
let processor = ChatTemplateProcessor::new(
template.to_string(),
Some("<s>".to_string()),
Some("</s>".to_string()),
);
let processor = ChatTemplateProcessor::new(template.to_string());
let messages = [spec::ChatMessage::User {
role: "user".to_string(),
@@ -230,8 +227,24 @@ fn test_chat_template_with_tokens_unit_test() {
.map(|msg| serde_json::to_value(msg).unwrap())
.collect();
// Use template_kwargs to pass tokens
let mut template_kwargs = std::collections::HashMap::new();
template_kwargs.insert(
"start_token".to_string(),
serde_json::Value::String("<s>".to_string()),
);
template_kwargs.insert(
"end_token".to_string(),
serde_json::Value::String("</s>".to_string()),
);
let params = ChatTemplateParams {
template_kwargs: Some(&template_kwargs),
..Default::default()
};
let result = processor
.apply_chat_template(&message_values, false)
.apply_chat_template(&message_values, params)
.unwrap();
assert!(result.contains("<s>"));
assert!(result.contains("</s>"));