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 InklingParser, JsonParser, KimiK2Parser, KimiK3Parser, LlamaParser, MinimaxM2Parser,
14 MinimaxM3Parser, MistralParser, PassthroughParser, PythonicParser, QwenParser,
15 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("cohere", || Box::new(CohereParser::new()));
422
423 Self::register_default_mappings(®istry);
425
426 Self { registry }
427 }
428
429 fn register_default_mappings(registry: &ParserRegistry) {
430 registry.map_model("gpt-4*", "json");
432 registry.map_model("gpt-3.5*", "json");
433 registry.map_model("gpt-4o*", "json");
434
435 registry.map_model("claude-*", "json");
437
438 registry.map_model("mistral-*", "mistral");
440 registry.map_model("mixtral-*", "mistral");
441
442 registry.map_model("Qwen/Qwen3.5*", "qwen_xml");
445 registry.map_model("Qwen3.5*", "qwen_xml");
446 registry.map_model("qwen3.5*", "qwen_xml");
447 registry.map_model("qwen/qwen3.5*", "qwen_xml");
448 registry.map_model("Qwen/Qwen3.6*", "qwen_xml");
449 registry.map_model("Qwen3.6*", "qwen_xml");
450 registry.map_model("qwen3.6*", "qwen_xml");
451 registry.map_model("qwen/qwen3.6*", "qwen_xml");
452 registry.map_model("Qwen/Qwen3.8*", "qwen_xml");
453 registry.map_model("Qwen3.8*", "qwen_xml");
454 registry.map_model("qwen3.8*", "qwen_xml");
455 registry.map_model("qwen/qwen3.8*", "qwen_xml");
456 registry.map_model("Qwen/Qwen3-Coder*", "qwen_xml");
457 registry.map_model("Qwen3-Coder*", "qwen_xml");
458 registry.map_model("qwen3-coder*", "qwen_xml");
459 registry.map_model("qwen/qwen3-coder*", "qwen_xml");
460 registry.map_model("qwen*", "qwen");
462 registry.map_model("Qwen*", "qwen");
463
464 registry.map_model("nemotron-3*", "qwen_xml");
471
472 registry.map_model("llama-4*", "pythonic");
474 registry.map_model("meta-llama-4*", "pythonic");
475 registry.map_model("llama-3.2*", "llama");
476 registry.map_model("meta-llama-3.2*", "llama");
477 registry.map_model("llama-*", "json");
478 registry.map_model("meta-llama-*", "json");
479
480 registry.map_model("deepseek-v3*", "deepseek");
482 registry.map_model("deepseek-ai/DeepSeek-V3*", "deepseek");
483 registry.map_model("deepseek-v3.1*", "deepseek31");
484 registry.map_model("deepseek-ai/DeepSeek-V3.1*", "deepseek31");
485 registry.map_model("deepseek-v3.2-exp*", "deepseek31");
487 registry.map_model("deepseek-ai/DeepSeek-V3.2-Exp*", "deepseek31");
488 registry.map_model("deepseek-v3.2*", "deepseek32");
490 registry.map_model("deepseek-ai/DeepSeek-V3.2*", "deepseek32");
491 registry.map_model("deepseek-v4*", "deepseek_v4");
493 registry.map_model("deepseek-ai/DeepSeek-V4*", "deepseek_v4");
494 registry.map_model("deepseek-v4.1*", "deepseek_v41");
496 registry.map_model("deepseek_v41*", "deepseek_v41");
497 registry.map_model("deepseek-v41*", "deepseek_v41");
498 registry.map_model("deepseek-ai/DeepSeek-V4.1*", "deepseek_v41");
499 registry.map_model("deepseek-ai/DeepSeek-V41*", "deepseek_v41");
502 registry.map_model("deepseek-*", "pythonic");
503
504 registry.map_model("glm-4.5*", "glm45_moe");
506 registry.map_model("glm-4.6*", "glm45_moe");
507 registry.map_model("glm-4.7*", "glm47_moe");
508 registry.map_model("glm-5*", "glm47_moe");
509 registry.map_model("glm-*", "json");
510
511 registry.map_model("sarashina*", "sarashina");
513 registry.map_model("Sarashina*", "sarashina");
514
515 registry.map_model("step3*", "step3");
517 registry.map_model("Step-3*", "step3");
518
519 registry.map_model("kimi-k2*", "kimik2");
521 registry.map_model("Kimi-K2*", "kimik2");
522 registry.map_model("moonshot*/Kimi-K2*", "kimik2");
523 registry.map_model("kimi-k3*", "kimi_k3");
524 registry.map_model("Kimi-K3*", "kimi_k3");
525 registry.map_model("moonshot*/Kimi-K3*", "kimi_k3");
526 registry.map_model("kimi_k3*", "kimi_k3");
530 registry.map_model("Kimi_K3*", "kimi_k3");
531 registry.map_model("moonshot*/Kimi_K3*", "kimi_k3");
532
533 registry.map_model("inkling*", "inkling");
535 registry.map_model("Inkling*", "inkling");
536
537 registry.map_model("minimax*", "minimax_m2");
540 registry.map_model("MiniMax*", "minimax_m2");
541 registry.map_model("minimax-m3*", "minimax_m3");
543 registry.map_model("MiniMax-M3*", "minimax_m3");
544 registry.map_model("MiniMaxAI/MiniMax-M3*", "minimax_m3");
545 registry.map_model("mm-m3*", "minimax_m3");
546
547 registry.map_model("command-r*", "cohere");
549 registry.map_model("command-r-plus*", "cohere");
550 registry.map_model("command-a*", "cohere");
551 registry.map_model("c4ai-command*", "cohere");
552 registry.map_model("cohere*", "cohere");
553 registry.map_model("CohereForAI*", "cohere");
554
555 registry.map_model("gemini-*", "json");
557 registry.map_model("palm-*", "json");
558 registry.map_model("gemma-*", "json");
559 }
560
561 #[expect(
565 clippy::expect_used,
566 reason = "passthrough parser is registered in new(); None indicates a bug in registration logic"
567 )]
568 pub fn get_pooled(&self, model_id: &str) -> PooledParser {
569 self.registry
570 .get_pooled_for_model(model_id)
571 .unwrap_or_else(|| {
572 self.registry
574 .get_pooled_parser("passthrough")
575 .expect("Passthrough parser should always be registered")
576 })
577 }
578
579 pub fn registry(&self) -> &ParserRegistry {
581 &self.registry
582 }
583
584 pub fn clear_pool(&self) {
586 self.registry.clear_pool();
587 }
588
589 pub fn get_parser(&self, model_id: &str) -> Option<Arc<dyn ToolParser>> {
592 let parser_type = self
593 .registry
594 .resolve_model_to_parser(model_id)
595 .unwrap_or_else(|| self.registry.default_parser.read().clone());
596 let entries = self.registry.entries.read();
597 entries.get(&parser_type).map(|entry| {
598 let boxed_parser = (entry.creator)();
599 Arc::from(boxed_parser)
600 })
601 }
602
603 pub fn has_parser(&self, name: &str) -> bool {
605 self.registry.entries.read().contains_key(name)
606 }
607
608 pub fn list_parsers(&self) -> Vec<String> {
610 self.registry.entries.read().keys().cloned().collect()
611 }
612}
613
614impl Default for ParserFactory {
615 fn default() -> Self {
616 Self::new()
617 }
618}
619
620fn wrap_in_reasoning_prefix(
624 mut tag: serde_json::Value,
625 prefix: serde_json::Value,
626) -> Result<serde_json::Value, String> {
627 let format = tag
628 .get_mut("format")
629 .map(serde_json::Value::take)
630 .ok_or_else(|| {
631 "structural tag has no `format` to wrap in a reasoning prefix".to_string()
632 })?;
633 Ok(json!({
634 "format": {
635 "type": "sequence",
636 "elements": [prefix, format],
637 }
638 }))
639}
640
641fn build_required_array_schema(tools: &[Tool]) -> Result<String, String> {
643 let mut any_of_schemas = Vec::with_capacity(tools.len());
644 for tool in tools {
645 let tool_schema = json!({
646 "properties": {
647 "name": {
648 "type": "string",
649 "enum": [tool.function.name]
650 },
651 "parameters": tool.function.parameters
652 },
653 "required": ["name", "parameters"]
654 });
655 any_of_schemas.push(tool_schema);
656 }
657
658 let mut all_defs: serde_json::Map<String, serde_json::Value> = serde_json::Map::new();
660 for tool in tools {
661 if let serde_json::Value::Object(params) = &tool.function.parameters {
662 if let Some(serde_json::Value::Object(defs)) = params.get("$defs") {
663 for (def_name, def_schema) in defs {
664 if let Some(existing) = all_defs.get(def_name) {
665 if existing != def_schema {
666 return Err(format!(
667 "Tool definition '{def_name}' has multiple conflicting schemas"
668 ));
669 }
670 } else {
671 all_defs.insert(def_name.clone(), def_schema.clone());
672 }
673 }
674 }
675 }
676 }
677
678 let mut array_schema = json!({
679 "type": "array",
680 "minItems": 1,
681 "items": {
682 "type": "object",
683 "anyOf": any_of_schemas
684 }
685 });
686
687 if !all_defs.is_empty() {
688 if let serde_json::Value::Object(ref mut obj) = array_schema {
689 obj.insert("$defs".to_string(), serde_json::Value::Object(all_defs));
690 }
691 }
692
693 serde_json::to_string(&array_schema)
694 .map_err(|e| format!("Failed to serialize tool schema: {e}"))
695}