lc_chains/
router_chain.rs

1// lc-chains/src/router_chain.rs
2//! Router Chain
3//!
4//! Automatically routes to different Chains based on input content.
5
6use async_trait::async_trait;
7use lc_core::language_models::LLMResult;
8use lc_core::{BaseChatModel, Runnable};
9use lc_schema::Message;
10use serde_json::Value;
11use std::collections::HashMap;
12use std::sync::Arc;
13
14use crate::base::{BaseChain, ChainError, ChainResult};
15
16/// Route destination.
17pub struct RouteDestination {
18    /// Destination name.
19    name: String,
20    /// Destination description (used for routing decisions).
21    description: String,
22    /// Destination Chain.
23    chain: Arc<dyn BaseChain>,
24    /// Keyword list (used for keyword-based routing).
25    keywords: Vec<String>,
26}
27
28impl RouteDestination {
29    pub fn new(
30        name: impl Into<String>,
31        description: impl Into<String>,
32        chain: Arc<dyn BaseChain>,
33    ) -> Self {
34        Self {
35            name: name.into(),
36            description: description.into(),
37            chain,
38            keywords: Vec::new(),
39        }
40    }
41
42    pub fn with_keywords(mut self, keywords: Vec<&str>) -> Self {
43        self.keywords = keywords.into_iter().map(String::from).collect();
44        self
45    }
46
47    pub fn name(&self) -> &str {
48        &self.name
49    }
50
51    pub fn description(&self) -> &str {
52        &self.description
53    }
54
55    pub fn chain(&self) -> &Arc<dyn BaseChain> {
56        &self.chain
57    }
58
59    pub fn keywords(&self) -> &[String] {
60        &self.keywords
61    }
62}
63
64/// Router Chain
65///
66/// Automatically routes to different Chains based on input content.
67pub struct RouterChain {
68    /// Route destination list.
69    destinations: Vec<RouteDestination>,
70
71    /// Default Chain (used when no match is found).
72    default_chain: Option<Arc<dyn BaseChain>>,
73
74    /// Input key name.
75    input_key: String,
76
77    /// Chain name.
78    name: String,
79
80    /// Whether to print verbose information.
81    verbose: bool,
82}
83
84impl RouterChain {
85    pub fn new() -> Self {
86        Self {
87            destinations: Vec::new(),
88            default_chain: None,
89            input_key: "input".to_string(),
90            name: "router_chain".to_string(),
91            verbose: false,
92        }
93    }
94
95    pub fn add_route(
96        mut self,
97        name: impl Into<String>,
98        description: impl Into<String>,
99        chain: Arc<dyn BaseChain>,
100    ) -> Self {
101        self.destinations
102            .push(RouteDestination::new(name, description, chain));
103        self
104    }
105
106    pub fn add_route_with_keywords(
107        mut self,
108        name: impl Into<String>,
109        description: impl Into<String>,
110        chain: Arc<dyn BaseChain>,
111        keywords: Vec<&str>,
112    ) -> Self {
113        self.destinations
114            .push(RouteDestination::new(name, description, chain).with_keywords(keywords));
115        self
116    }
117
118    pub fn with_default(mut self, chain: Arc<dyn BaseChain>) -> Self {
119        self.default_chain = Some(chain);
120        self
121    }
122
123    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
124        self.input_key = key.into();
125        self
126    }
127
128    pub fn with_name(mut self, name: impl Into<String>) -> Self {
129        self.name = name.into();
130        self
131    }
132
133    pub fn with_verbose(mut self, verbose: bool) -> Self {
134        self.verbose = verbose;
135        self
136    }
137
138    pub fn destinations(&self) -> &[RouteDestination] {
139        &self.destinations
140    }
141
142    pub fn default_chain(&self) -> Option<&Arc<dyn BaseChain>> {
143        self.default_chain.as_ref()
144    }
145
146    /// Keyword-based routing.
147    ///
148    /// Longest-match-first instead of first-match-wins.
149    fn route_by_keywords(&self, input: &str) -> Option<&RouteDestination> {
150        let mut best_match: Option<(&RouteDestination, usize)> = None;
151        for dest in &self.destinations {
152            for keyword in &dest.keywords {
153                if input.contains(keyword) {
154                    let len = keyword.len();
155                    if best_match.is_none() || len > best_match.unwrap().1 {
156                        best_match = Some((dest, len));
157                    }
158                }
159            }
160        }
161        best_match.map(|(dest, _)| dest)
162    }
163
164    /// Select a route destination.
165    fn select_route(&self, input: &str) -> Result<Option<&RouteDestination>, ChainError> {
166        if let Some(dest) = self.route_by_keywords(input) {
167            return Ok(Some(dest));
168        }
169
170        Ok(None)
171    }
172}
173
174impl Default for RouterChain {
175    fn default() -> Self {
176        Self::new()
177    }
178}
179
180#[async_trait]
181impl BaseChain for RouterChain {
182    fn input_keys(&self) -> Vec<&str> {
183        vec![&self.input_key]
184    }
185
186    fn output_keys(&self) -> Vec<&str> {
187        let mut seen = std::collections::HashSet::new();
188        let mut result: Vec<&str> = Vec::new();
189
190        for dest in &self.destinations {
191            for key in dest.chain().output_keys() {
192                if seen.insert(key.to_string()) {
193                    result.push(key);
194                }
195            }
196        }
197        if let Some(default) = &self.default_chain {
198            for key in default.output_keys() {
199                if seen.insert(key.to_string()) {
200                    result.push(key);
201                }
202            }
203        }
204
205        if result.is_empty() {
206            vec!["output"]
207        } else {
208            result
209        }
210    }
211
212    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
213        self.validate_inputs(&inputs)?;
214
215        let input = inputs
216            .get(&self.input_key)
217            .and_then(|v| v.as_str())
218            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
219
220        if self.verbose {
221            println!("\n=== RouterChain execution ===");
222            println!("Input: {}", input);
223            println!("Route destination count: {}", self.destinations.len());
224        }
225
226        let route_result = self.select_route(input)?;
227
228        let chain = match route_result {
229            Some(dest) => {
230                if self.verbose {
231                    println!("Routed to: {} ({})", dest.name(), dest.description());
232                }
233                dest.chain()
234            }
235            None => {
236                if let Some(default) = &self.default_chain {
237                    if self.verbose {
238                        println!("No keyword match, using default Chain");
239                    }
240                    default
241                } else {
242                    return Err(ChainError::ExecutionError(
243                        "No matching route destination and no default Chain configured".to_string(),
244                    ));
245                }
246            }
247        };
248
249        let result = chain.invoke(inputs).await?;
250
251        if self.verbose {
252            println!("=== RouterChain complete ===\n");
253        }
254
255        Ok(result)
256    }
257
258    fn name(&self) -> &str {
259        &self.name
260    }
261}
262
263/// LLM Router Chain
264///
265/// Uses an LLM to intelligently determine the routing destination.
266pub struct LLMRouterChain<M: BaseChatModel> {
267    /// LLM used for routing decisions.
268    llm: M,
269
270    /// Route destinations.
271    destinations: Vec<RouteDestination>,
272
273    /// Default Chain.
274    default_chain: Option<Arc<dyn BaseChain>>,
275
276    /// Input key name.
277    input_key: String,
278
279    /// Chain name.
280    name: String,
281
282    /// Whether to print verbose information.
283    verbose: bool,
284}
285
286impl<M: BaseChatModel> LLMRouterChain<M> {
287    pub fn new(llm: M) -> Self {
288        Self {
289            llm,
290            destinations: Vec::new(),
291            default_chain: None,
292            input_key: "input".to_string(),
293            name: "llm_router_chain".to_string(),
294            verbose: false,
295        }
296    }
297
298    pub fn add_route(
299        mut self,
300        name: impl Into<String>,
301        description: impl Into<String>,
302        chain: Arc<dyn BaseChain>,
303    ) -> Self {
304        self.destinations
305            .push(RouteDestination::new(name, description, chain));
306        self
307    }
308
309    pub fn add_route_with_keywords(
310        mut self,
311        name: impl Into<String>,
312        description: impl Into<String>,
313        chain: Arc<dyn BaseChain>,
314        keywords: Vec<&str>,
315    ) -> Self {
316        self.destinations
317            .push(RouteDestination::new(name, description, chain).with_keywords(keywords));
318        self
319    }
320
321    pub fn with_default(mut self, chain: Arc<dyn BaseChain>) -> Self {
322        self.default_chain = Some(chain);
323        self
324    }
325
326    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
327        self.input_key = key.into();
328        self
329    }
330
331    pub fn with_name(mut self, name: impl Into<String>) -> Self {
332        self.name = name.into();
333        self
334    }
335
336    pub fn with_verbose(mut self, verbose: bool) -> Self {
337        self.verbose = verbose;
338        self
339    }
340
341    pub fn destinations(&self) -> &[RouteDestination] {
342        &self.destinations
343    }
344
345    pub fn default_chain(&self) -> Option<&Arc<dyn BaseChain>> {
346        self.default_chain.as_ref()
347    }
348
349    /// Build the LLM routing prompt.
350    fn build_router_prompt(&self, input: &str) -> String {
351        let mut prompt =
352            String::from("Based on the user input, select the most appropriate handler.\n\n");
353        prompt.push_str("Available handlers:\n");
354
355        for (i, dest) in self.destinations.iter().enumerate() {
356            prompt.push_str(&format!(
357                "{}. {}: {}\n",
358                i + 1,
359                dest.name(),
360                dest.description()
361            ));
362        }
363
364        prompt.push_str("\nUser input: ");
365        prompt.push_str(input);
366        prompt
367            .push_str("\n\nReturn only the name of the most appropriate handler (no explanation).");
368
369        prompt
370    }
371
372    /// Use LLM to determine the route.
373    async fn route_with_llm(&self, input: &str) -> Result<String, ChainError> {
374        let prompt = self.build_router_prompt(input);
375
376        let messages = vec![Message::human(&prompt)];
377
378        let result = self
379            .llm
380            .invoke(messages, None)
381            .await
382            .map_err(|e| ChainError::ExecutionError(format!("LLM call failed: {}", e)))?;
383
384        Ok(result.content.trim().to_string())
385    }
386
387    /// Find a route destination by name.
388    fn find_destination(&self, name: &str) -> Option<&RouteDestination> {
389        let name_lower = name.to_lowercase();
390        // 1. Exact case-insensitive match
391        if let Some(dest) = self
392            .destinations
393            .iter()
394            .find(|d| d.name().eq_ignore_ascii_case(name))
395        {
396            return Some(dest);
397        }
398        // 2. The LLM result starts or ends with the destination name
399        self.destinations.iter().find(|d| {
400            let d_lower = d.name().to_lowercase();
401            name_lower.starts_with(&d_lower)
402                || name_lower.ends_with(&d_lower)
403                || name_lower
404                    .split_whitespace()
405                    .any(|word| word.eq_ignore_ascii_case(&d_lower))
406        })
407    }
408
409    /// LLM routing takes priority over keyword matching.
410    async fn select_route(&self, input: &str) -> Result<&RouteDestination, ChainError> {
411        if self.destinations.is_empty() {
412            return Err(ChainError::ExecutionError(
413                "No route destinations configured".to_string(),
414            ));
415        }
416
417        if self.destinations.len() == 1 {
418            return Ok(&self.destinations[0]);
419        }
420
421        // Try LLM routing first (primary strategy)
422        let llm_result = self.route_with_llm(input).await;
423        match llm_result {
424            Ok(route_name) => {
425                if let Some(dest) = self.find_destination(&route_name) {
426                    return Ok(dest);
427                }
428            }
429            Err(_) => {
430                // LLM call failed; fall through to keyword matching
431            }
432        }
433
434        // Fallback: keyword matching (longest match first)
435        let mut best_match: Option<(&RouteDestination, usize)> = None;
436        for dest in &self.destinations {
437            for keyword in dest.keywords() {
438                if input.contains(keyword) {
439                    let len = keyword.len();
440                    if best_match.is_none() || len > best_match.unwrap().1 {
441                        best_match = Some((dest, len));
442                    }
443                }
444            }
445        }
446        if let Some((dest, _)) = best_match {
447            return Ok(dest);
448        }
449
450        Err(ChainError::ExecutionError(
451            "No matching route destination found (LLM and keyword matching both failed)"
452                .to_string(),
453        ))
454    }
455}
456
457#[async_trait]
458impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for LLMRouterChain<M>
459where
460    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
461{
462    fn input_keys(&self) -> Vec<&str> {
463        vec![&self.input_key]
464    }
465
466    fn output_keys(&self) -> Vec<&str> {
467        let mut seen = std::collections::HashSet::new();
468        let mut result: Vec<&str> = Vec::new();
469
470        for dest in &self.destinations {
471            for key in dest.chain().output_keys() {
472                if seen.insert(key.to_string()) {
473                    result.push(key);
474                }
475            }
476        }
477        if let Some(default) = &self.default_chain {
478            for key in default.output_keys() {
479                if seen.insert(key.to_string()) {
480                    result.push(key);
481                }
482            }
483        }
484
485        if result.is_empty() {
486            vec!["output"]
487        } else {
488            result
489        }
490    }
491
492    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
493        self.validate_inputs(&inputs)?;
494
495        let input = inputs
496            .get(&self.input_key)
497            .and_then(|v| v.as_str())
498            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
499
500        if self.verbose {
501            println!("\n=== LLMRouterChain execution ===");
502            println!("Input: {}", input);
503            println!("Route destination count: {}", self.destinations.len());
504        }
505
506        let route_result = self.select_route(input).await;
507
508        let chain = match route_result {
509            Ok(dest) => {
510                if self.verbose {
511                    println!("Routed to: {} ({})", dest.name(), dest.description());
512                }
513                dest.chain()
514            }
515            Err(e) => {
516                if let Some(default) = &self.default_chain {
517                    if self.verbose {
518                        println!("Routing failed: {}, using default Chain", e);
519                    }
520                    default
521                } else {
522                    return Err(e);
523                }
524            }
525        };
526
527        let result = chain.invoke(inputs).await?;
528
529        if self.verbose {
530            println!("=== LLMRouterChain complete ===\n");
531        }
532
533        Ok(result)
534    }
535
536    fn name(&self) -> &str {
537        &self.name
538    }
539}
lc_chains/router_chain.rs

lc_chains/
router_chain.rs
