1use std::{collections::HashMap, sync::Arc};
4
5use openai_protocol::common::{Tool, ToolChoice, ToolChoiceValue};
6use parking_lot::RwLock;
7use serde_json::json;
8use tokio::sync::Mutex;
9
10use crate::{
11 parsers::{
12 CohereParser, DeepSeek31Parser, DeepSeekDsmlParser, DeepSeekParser, Glm4MoeParser,
13 HyV4Parser, InklingParser, JsonParser, KimiK2Parser, KimiK3Parser, LlamaParser,
14 MinimaxM2Parser, MinimaxM3Parser, MistralParser, PassthroughParser, PythonicParser,
15 QwenParser, QwenXmlParser, SarashinaParser, Step3Parser,
16 },
17 traits::ToolParser,
18};
19
20pub type PooledParser = Arc<Mutex<Box<dyn ToolParser>>>;
22
23type ParserCreator = Arc<dyn Fn() -> Box<dyn ToolParser> + Send + Sync>;
25
26type BuildStructuralTagFn = Arc<dyn Fn(&[Tool], bool) -> serde_json::Value + Send + Sync>;
29
30type ReasoningPrefixFn = fn() -> serde_json::Value;
35
36#[derive(Debug, Clone)]
38pub enum ToolConstraint {
39 JsonSchema(String),
42 StructuralTag(String),
45}
46
47impl ToolConstraint {
48 pub fn to_tuple(&self) -> (String, String) {
50 match self {
51 ToolConstraint::JsonSchema(s) => ("json_schema".to_string(), s.clone()),
52 ToolConstraint::StructuralTag(s) => ("structural_tag".to_string(), s.clone()),
53 }
54 }
55
56 pub fn is_json_schema(&self) -> bool {
58 matches!(self, ToolConstraint::JsonSchema(_))
59 }
60}
61
62struct ParserEntry {
67 creator: ParserCreator,
68 build_structural_tag: Option<BuildStructuralTagFn>,
69 reasoning_prefix: Option<ReasoningPrefixFn>,
70}
71
72#[derive(Clone)]
74pub struct ParserRegistry {
75 entries: Arc<RwLock<HashMap<String, Arc<ParserEntry>>>>,
77 pool: Arc<RwLock<HashMap<String, PooledParser>>>,
79 model_mapping: Arc<RwLock<HashMap<String, String>>>,
81 default_parser: Arc<RwLock<String>>,
83}
84
85impl ParserRegistry {
86 pub fn new() -> Self {
88 Self {
89 entries: Arc::new(RwLock::new(HashMap::new())),
90 pool: Arc::new(RwLock::new(HashMap::new())),
91 model_mapping: Arc::new(RwLock::new(HashMap::new())),
92 default_parser: Arc::new(RwLock::new("passthrough".to_string())),
93 }
94 }
95
96 pub fn register_parser<F>(&self, name: &str, creator: F)
98 where
99 F: Fn() -> Box<dyn ToolParser> + Send + Sync + 'static,
100 {
101 let mut entries = self.entries.write();
102 entries.insert(
103 name.to_string(),
104 Arc::new(ParserEntry {
105 creator: Arc::new(creator),
106 build_structural_tag: None,
107 reasoning_prefix: None,
108 }),
109 );
110 }
111
112 pub fn register_parser_with_structural_tag<F>(
117 &self,
118 name: &str,
119 creator: F,
120 build_structural_tag: fn(&[Tool], bool) -> serde_json::Value,
121 ) where
122 F: Fn() -> Box<dyn ToolParser> + Send + Sync + 'static,
123 {
124 let mut entries = self.entries.write();
125 entries.insert(
126 name.to_string(),
127 Arc::new(ParserEntry {
128 creator: Arc::new(creator),
129 build_structural_tag: Some(Arc::new(build_structural_tag)),
130 reasoning_prefix: None,
131 }),
132 );
133 }
134
135 pub fn register_reasoning_prefix(&self, name: &str, reasoning_prefix: ReasoningPrefixFn) {
145 let mut entries = self.entries.write();
146 let Some(existing) = entries.get(name).map(Arc::clone) else {
147 debug_assert!(false, "reasoning prefix for unregistered parser {name}");
148 return;
149 };
150 debug_assert!(
151 existing.build_structural_tag.is_some(),
152 "reasoning prefix for parser {name} without a structural tag"
153 );
154 entries.insert(
155 name.to_string(),
156 Arc::new(ParserEntry {
157 creator: Arc::clone(&existing.creator),
158 build_structural_tag: existing.build_structural_tag.clone(),
159 reasoning_prefix: Some(reasoning_prefix),
160 }),
161 );
162 }
163
164 pub fn has_reasoning_prefix(&self, configured: Option<&str>) -> bool {
168 configured.is_some_and(|name| {
169 self.entries
170 .read()
171 .get(name)
172 .is_some_and(|entry| entry.reasoning_prefix.is_some())
173 })
174 }
175
176 pub fn map_model(&self, model: impl Into<String>, parser: impl Into<String>) {
178 let mut mapping = self.model_mapping.write();
179 mapping.insert(model.into(), parser.into());
180 }
181
182 pub fn get_pooled_parser(&self, name: &str) -> Option<PooledParser> {
185 {
187 let pool = self.pool.read();
188 if let Some(parser) = pool.get(name) {
189 return Some(Arc::clone(parser));
190 }
191 }
192
193 let entries = self.entries.read();
195 if let Some(entry) = entries.get(name) {
196 let parser = Arc::new(Mutex::new((entry.creator)()));
197
198 let mut pool = self.pool.write();
200 pool.insert(name.to_string(), Arc::clone(&parser));
201
202 Some(parser)
203 } else {
204 None
205 }
206 }
207
208 pub fn has_parser(&self, name: &str) -> bool {
210 let entries = self.entries.read();
211 entries.contains_key(name)
212 }
213
214 pub fn create_parser(&self, name: &str) -> Option<Box<dyn ToolParser>> {
217 let entries = self.entries.read();
218 entries.get(name).map(|entry| (entry.creator)())
219 }
220
221 pub fn has_structural_tag(&self, name: &str) -> bool {
223 let entries = self.entries.read();
224 entries
225 .get(name)
226 .is_some_and(|e| e.build_structural_tag.is_some())
227 }
228
229 pub fn has_structural_tag_for_parser(&self, configured: Option<&str>) -> bool {
233 configured.is_some_and(|p| self.has_structural_tag(p))
234 }
235
236 pub fn generate_tool_constraint(
248 &self,
249 configured_parser: Option<&str>,
250 tools: &[Tool],
251 tool_choice: &ToolChoice,
252 reasoning: bool,
253 ) -> Result<Option<ToolConstraint>, String> {
254 if tools.is_empty() {
255 return Ok(None);
256 }
257
258 let at_least_one = match tool_choice {
259 ToolChoice::Value(ToolChoiceValue::Required) => true,
260 ToolChoice::Function { .. } => true,
261 ToolChoice::AllowedTools { mode, .. } if mode == "required" => true,
262 _ => return Ok(None),
263 };
264
265 if let Some(name) = configured_parser {
267 let entries = self.entries.read();
268 if let Some(entry) = entries.get(name) {
269 if let Some(build_fn) = entry.build_structural_tag.as_ref() {
270 let mut tag = build_fn(tools, at_least_one);
271 if let (true, Some(prefix)) = (reasoning, entry.reasoning_prefix) {
272 tag = wrap_in_reasoning_prefix(tag, prefix())?;
273 }
274 let json_str = serde_json::to_string(&tag)
275 .map_err(|e| format!("Failed to serialize structural tag: {e}"))?;
276 return Ok(Some(ToolConstraint::StructuralTag(json_str)));
277 }
278 }
279 }
280
281 match tool_choice {
283 ToolChoice::Function { .. } => {
284 let params_schema = serde_json::to_string(&tools[0].function.parameters)
285 .map_err(|e| format!("Failed to serialize tool parameters: {e}"))?;
286 Ok(Some(ToolConstraint::JsonSchema(params_schema)))
287 }
288 _ => {
289 let schema = build_required_array_schema(tools)?;
290 Ok(Some(ToolConstraint::JsonSchema(schema)))
291 }
292 }
293 }
294
295 pub fn resolve_model_to_parser(&self, model: &str) -> Option<String> {
298 let mapping = self.model_mapping.read();
299 if let Some(parser_name) = mapping.get(model) {
301 return Some(parser_name.clone());
302 }
303 let model_lower = model.to_lowercase();
306 mapping
307 .iter()
308 .filter_map(|(pattern, parser_name)| {
309 let stem = pattern.strip_suffix('*')?;
310 model_lower
311 .contains(&stem.to_lowercase())
312 .then_some((stem, parser_name))
313 })
314 .max_by_key(|(stem, _)| stem.len())
315 .map(|(_, parser_name)| parser_name.clone())
316 }
317
318 pub fn has_parser_for_model(&self, model: &str) -> bool {
320 self.resolve_model_to_parser(model)
321 .is_some_and(|name| self.has_parser(&name))
322 }
323
324 pub fn create_for_model(&self, model: &str) -> Option<Box<dyn ToolParser>> {
326 let parser_name = self
327 .resolve_model_to_parser(model)
328 .unwrap_or_else(|| self.default_parser.read().clone());
329 self.create_parser(&parser_name)
330 }
331
332 pub fn get_pooled_for_model(&self, model: &str) -> Option<PooledParser> {
334 let parser_name = self
335 .resolve_model_to_parser(model)
336 .unwrap_or_else(|| self.default_parser.read().clone());
337 self.get_pooled_parser(&parser_name)
338 }
339
340 pub fn clear_pool(&self) {
342 let mut pool = self.pool.write();
343 pool.clear();
344 }
345
346 pub fn set_default_parser(&self, name: impl Into<String>) {
348 let mut default = self.default_parser.write();
349 *default = name.into();
350 }
351}
352
353impl Default for ParserRegistry {
354 fn default() -> Self {
355 Self::new()
356 }
357}
358
359#[derive(Clone)]
361pub struct ParserFactory {
362 registry: ParserRegistry,
363}
364
365impl ParserFactory {
366 pub fn new() -> Self {
368 let registry = ParserRegistry::new();
369
370 registry.register_parser("passthrough", || Box::new(PassthroughParser::new()));
372 registry.register_parser("json", || Box::new(JsonParser::new()));
373 registry.register_parser_with_structural_tag(
374 "mistral",
375 || Box::new(MistralParser::new()),
376 MistralParser::build_structural_tag,
377 );
378 registry.register_parser("qwen", || Box::new(QwenParser::new()));
379 registry.register_parser("qwen_xml", || Box::new(QwenXmlParser::new()));
380 registry.register_parser("qwen_coder", || Box::new(QwenXmlParser::new()));
381 registry.register_parser("nemotron", || Box::new(QwenXmlParser::new()));
384 registry.register_parser("pythonic", || Box::new(PythonicParser::new()));
385 registry.register_parser("llama", || Box::new(LlamaParser::new()));
386 registry.register_parser("deepseek", || Box::new(DeepSeekParser::new()));
387 registry.register_parser("deepseek31", || Box::new(DeepSeek31Parser::new()));
388 registry.register_parser("deepseek32", || Box::new(DeepSeekDsmlParser::v32()));
389 registry.register_parser("deepseek_v4", || Box::new(DeepSeekDsmlParser::v4()));
390 registry.register_parser_with_structural_tag(
391 "deepseek_v41",
392 || Box::new(DeepSeekDsmlParser::v41()),
393 DeepSeekDsmlParser::build_v41_structural_tag,
394 );
395 registry.register_parser("glm45_moe", || Box::new(Glm4MoeParser::glm45()));
396 registry.register_parser_with_structural_tag(
397 "glm47_moe",
398 || Box::new(Glm4MoeParser::glm47()),
399 Glm4MoeParser::build_structural_tag,
400 );
401 registry.register_reasoning_prefix("glm47_moe", Glm4MoeParser::reasoning_prefix);
402 registry.register_parser("step3", || Box::new(Step3Parser::new()));
403 registry.register_parser("sarashina", || Box::new(SarashinaParser::new()));
404 registry.register_parser_with_structural_tag(
405 "kimik2",
406 || Box::new(KimiK2Parser::new()),
407 KimiK2Parser::build_structural_tag,
408 );
409 registry.register_parser_with_structural_tag(
410 "kimi_k3",
411 || Box::new(KimiK3Parser::new()),
412 KimiK3Parser::build_structural_tag,
413 );
414 registry.register_parser_with_structural_tag(
415 "inkling",
416 || Box::new(InklingParser::new()),
417 InklingParser::build_structural_tag,
418 );
419 registry.register_parser("minimax_m2", || Box::new(MinimaxM2Parser::new()));
420 registry.register_parser("minimax_m3", || Box::new(MinimaxM3Parser::new()));
421 registry.register_parser_with_structural_tag(
422 "hy_v4",
423 || Box::new(HyV4Parser::new()),
424 HyV4Parser::build_structural_tag,
425 );
426 registry.register_reasoning_prefix("hy_v4", HyV4Parser::reasoning_prefix);
427 registry.register_parser("cohere", || Box::new(CohereParser::new()));
428
429 Self::register_default_mappings(®istry);
431
432 Self { registry }
433 }
434
435 fn register_default_mappings(registry: &ParserRegistry) {
436 registry.map_model("gpt-4*", "json");
438 registry.map_model("gpt-3.5*", "json");
439 registry.map_model("gpt-4o*", "json");
440
441 registry.map_model("claude-*", "json");
443
444 registry.map_model("mistral-*", "mistral");
446 registry.map_model("mixtral-*", "mistral");
447
448 registry.map_model("Qwen/Qwen3.5*", "qwen_xml");
451 registry.map_model("Qwen3.5*", "qwen_xml");
452 registry.map_model("qwen3.5*", "qwen_xml");
453 registry.map_model("qwen/qwen3.5*", "qwen_xml");
454 registry.map_model("Qwen/Qwen3.6*", "qwen_xml");
455 registry.map_model("Qwen3.6*", "qwen_xml");
456 registry.map_model("qwen3.6*", "qwen_xml");
457 registry.map_model("qwen/qwen3.6*", "qwen_xml");
458 registry.map_model("Qwen/Qwen3.8*", "qwen_xml");
459 registry.map_model("Qwen3.8*", "qwen_xml");
460 registry.map_model("qwen3.8*", "qwen_xml");
461 registry.map_model("qwen/qwen3.8*", "qwen_xml");
462 registry.map_model("Qwen/Qwen3-Coder*", "qwen_xml");
463 registry.map_model("Qwen3-Coder*", "qwen_xml");
464 registry.map_model("qwen3-coder*", "qwen_xml");
465 registry.map_model("qwen/qwen3-coder*", "qwen_xml");
466 registry.map_model("qwen*", "qwen");
468 registry.map_model("Qwen*", "qwen");
469
470 registry.map_model("nemotron-3*", "qwen_xml");
477
478 registry.map_model("llama-4*", "pythonic");
480 registry.map_model("meta-llama-4*", "pythonic");
481 registry.map_model("llama-3.2*", "llama");
482 registry.map_model("meta-llama-3.2*", "llama");
483 registry.map_model("llama-*", "json");
484 registry.map_model("meta-llama-*", "json");
485
486 registry.map_model("deepseek-v3*", "deepseek");
488 registry.map_model("deepseek-ai/DeepSeek-V3*", "deepseek");
489 registry.map_model("deepseek-v3.1*", "deepseek31");
490 registry.map_model("deepseek-ai/DeepSeek-V3.1*", "deepseek31");
491 registry.map_model("deepseek-v3.2-exp*", "deepseek31");
493 registry.map_model("deepseek-ai/DeepSeek-V3.2-Exp*", "deepseek31");
494 registry.map_model("deepseek-v3.2*", "deepseek32");
496 registry.map_model("deepseek-ai/DeepSeek-V3.2*", "deepseek32");
497 registry.map_model("deepseek-v4*", "deepseek_v4");
499 registry.map_model("deepseek-ai/DeepSeek-V4*", "deepseek_v4");
500 registry.map_model("deepseek-v4.1*", "deepseek_v41");
502 registry.map_model("deepseek_v41*", "deepseek_v41");
503 registry.map_model("deepseek-v41*", "deepseek_v41");
504 registry.map_model("deepseek-ai/DeepSeek-V4.1*", "deepseek_v41");
505 registry.map_model("deepseek-ai/DeepSeek-V41*", "deepseek_v41");
508 registry.map_model("deepseek-*", "pythonic");
509
510 registry.map_model("glm-4.5*", "glm45_moe");
512 registry.map_model("glm-4.6*", "glm45_moe");
513 registry.map_model("glm-4.7*", "glm47_moe");
514 registry.map_model("glm-5*", "glm47_moe");
515 registry.map_model("glm-*", "json");
516
517 registry.map_model("sarashina*", "sarashina");
519 registry.map_model("Sarashina*", "sarashina");
520
521 registry.map_model("step3*", "step3");
523 registry.map_model("Step-3*", "step3");
524
525 registry.map_model("kimi-k2*", "kimik2");
527 registry.map_model("Kimi-K2*", "kimik2");
528 registry.map_model("moonshot*/Kimi-K2*", "kimik2");
529 registry.map_model("hy4*", "hy_v4");
530 registry.map_model("hy_v4", "hy_v4");
531 registry.map_model("kimi-k3*", "kimi_k3");
532 registry.map_model("Kimi-K3*", "kimi_k3");
533 registry.map_model("moonshot*/Kimi-K3*", "kimi_k3");
534 registry.map_model("kimi_k3*", "kimi_k3");
538 registry.map_model("Kimi_K3*", "kimi_k3");
539 registry.map_model("moonshot*/Kimi_K3*", "kimi_k3");
540
541 registry.map_model("inkling*", "inkling");
543 registry.map_model("Inkling*", "inkling");
544
545 registry.map_model("minimax*", "minimax_m2");
548 registry.map_model("MiniMax*", "minimax_m2");
549 registry.map_model("minimax-m3*", "minimax_m3");
551 registry.map_model("MiniMax-M3*", "minimax_m3");
552 registry.map_model("MiniMaxAI/MiniMax-M3*", "minimax_m3");
553 registry.map_model("mm-m3*", "minimax_m3");
554
555 registry.map_model("command-r*", "cohere");
557 registry.map_model("command-r-plus*", "cohere");
558 registry.map_model("command-a*", "cohere");
559 registry.map_model("c4ai-command*", "cohere");
560 registry.map_model("cohere*", "cohere");
561 registry.map_model("CohereForAI*", "cohere");
562
563 registry.map_model("gemini-*", "json");
565 registry.map_model("palm-*", "json");
566 registry.map_model("gemma-*", "json");
567 }
568
569 #[expect(
573 clippy::expect_used,
574 reason = "passthrough parser is registered in new(); None indicates a bug in registration logic"
575 )]
576 pub fn get_pooled(&self, model_id: &str) -> PooledParser {
577 self.registry
578 .get_pooled_for_model(model_id)
579 .unwrap_or_else(|| {
580 self.registry
582 .get_pooled_parser("passthrough")
583 .expect("Passthrough parser should always be registered")
584 })
585 }
586
587 pub fn registry(&self) -> &ParserRegistry {
589 &self.registry
590 }
591
592 pub fn clear_pool(&self) {
594 self.registry.clear_pool();
595 }
596
597 pub fn get_parser(&self, model_id: &str) -> Option<Arc<dyn ToolParser>> {
600 let parser_type = self
601 .registry
602 .resolve_model_to_parser(model_id)
603 .unwrap_or_else(|| self.registry.default_parser.read().clone());
604 let entries = self.registry.entries.read();
605 entries.get(&parser_type).map(|entry| {
606 let boxed_parser = (entry.creator)();
607 Arc::from(boxed_parser)
608 })
609 }
610
611 pub fn has_parser(&self, name: &str) -> bool {
613 self.registry.entries.read().contains_key(name)
614 }
615
616 pub fn list_parsers(&self) -> Vec<String> {
618 self.registry.entries.read().keys().cloned().collect()
619 }
620}
621
622impl Default for ParserFactory {
623 fn default() -> Self {
624 Self::new()
625 }
626}
627
628fn wrap_in_reasoning_prefix(
632 mut tag: serde_json::Value,
633 prefix: serde_json::Value,
634) -> Result<serde_json::Value, String> {
635 let format = tag
636 .get_mut("format")
637 .map(serde_json::Value::take)
638 .ok_or_else(|| {
639 "structural tag has no `format` to wrap in a reasoning prefix".to_string()
640 })?;
641 Ok(json!({
642 "format": {
643 "type": "sequence",
644 "elements": [prefix, format],
645 }
646 }))
647}
648
649fn build_required_array_schema(tools: &[Tool]) -> Result<String, String> {
651 let mut any_of_schemas = Vec::with_capacity(tools.len());
652 for tool in tools {
653 let tool_schema = json!({
654 "properties": {
655 "name": {
656 "type": "string",
657 "enum": [tool.function.name]
658 },
659 "parameters": tool.function.parameters
660 },
661 "required": ["name", "parameters"]
662 });
663 any_of_schemas.push(tool_schema);
664 }
665
666 let mut all_defs: serde_json::Map<String, serde_json::Value> = serde_json::Map::new();
668 for tool in tools {
669 if let serde_json::Value::Object(params) = &tool.function.parameters {
670 if let Some(serde_json::Value::Object(defs)) = params.get("$defs") {
671 for (def_name, def_schema) in defs {
672 if let Some(existing) = all_defs.get(def_name) {
673 if existing != def_schema {
674 return Err(format!(
675 "Tool definition '{def_name}' has multiple conflicting schemas"
676 ));
677 }
678 } else {
679 all_defs.insert(def_name.clone(), def_schema.clone());
680 }
681 }
682 }
683 }
684 }
685
686 let mut array_schema = json!({
687 "type": "array",
688 "minItems": 1,
689 "items": {
690 "type": "object",
691 "anyOf": any_of_schemas
692 }
693 });
694
695 if !all_defs.is_empty() {
696 if let serde_json::Value::Object(ref mut obj) = array_schema {
697 obj.insert("$defs".to_string(), serde_json::Value::Object(all_defs));
698 }
699 }
700
701 serde_json::to_string(&array_schema)
702 .map_err(|e| format!("Failed to serialize tool schema: {e}"))
703}