1use std::collections::HashMap;
9use std::sync::Arc;
10
11use lc_chains::base::BaseChain;
12
13pub trait SkillRouter: Send + Sync {
18 fn chain_for(&self, skill_id: &str) -> Option<Arc<dyn BaseChain>>;
20}
21
22#[derive(Default)]
27pub struct SkillMapRouter {
28 skills: HashMap<String, Arc<dyn BaseChain>>,
29}
30
31impl SkillMapRouter {
32 pub fn new() -> Self {
34 Self::default()
35 }
36
37 pub fn with_skill(mut self, skill_id: impl Into<String>, chain: Arc<dyn BaseChain>) -> Self {
39 self.skills.insert(skill_id.into(), chain);
40 self
41 }
42}
43
44impl SkillRouter for SkillMapRouter {
45 fn chain_for(&self, skill_id: &str) -> Option<Arc<dyn BaseChain>> {
46 self.skills.get(skill_id).cloned()
47 }
48}
49
50#[cfg(test)]
51mod tests {
52 use super::*;
53 use lc_chains::base::{ChainError, ChainResult};
54 use serde_json::Value;
55
56 struct NamedChain(String);
58
59 #[async_trait::async_trait]
60 impl BaseChain for NamedChain {
61 fn input_keys(&self) -> Vec<&str> {
62 vec!["input"]
63 }
64
65 fn output_keys(&self) -> Vec<&str> {
66 vec!["output"]
67 }
68
69 async fn invoke(&self, _inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
70 let mut out = HashMap::new();
71 out.insert("output".to_string(), Value::String(self.0.clone()));
72 Ok(out)
73 }
74
75 fn name(&self) -> &str {
76 &self.0
77 }
78 }
79
80 fn arc_named(name: &str) -> Arc<dyn BaseChain> {
81 Arc::new(NamedChain(name.to_string()))
82 }
83
84 #[test]
85 fn empty_router_falls_through() {
86 let router = SkillMapRouter::new();
87 assert!(router.chain_for("anything").is_none());
88 }
89
90 #[test]
91 fn routes_by_skill_id() {
92 let router = SkillMapRouter::new()
93 .with_skill("research", arc_named("research-chain"))
94 .with_skill("summarize", arc_named("summary-chain"));
95
96 assert!(router.chain_for("research").is_some());
97 assert!(router.chain_for("summarize").is_some());
98 assert!(router.chain_for("nope").is_none());
100 let a = router.chain_for("research").unwrap();
102 let b = router.chain_for("summarize").unwrap();
103 assert_ne!(a.name(), b.name());
104 }
105}