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, ChainStream};
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    /// Stream execution for RouterChain.
259    ///
260    /// After routing (keyword matching), delegates to the selected chain's
261    /// `stream()` method.
262    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
263        self.validate_inputs(&inputs)?;
264
265        let input = inputs
266            .get(&self.input_key)
267            .and_then(|v| v.as_str())
268            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
269
270        let route_result = self.select_route(input)?;
271
272        let chain = match route_result {
273            Some(dest) => dest.chain(),
274            None => self
275                .default_chain
276                .as_ref()
277                .ok_or_else(|| ChainError::ExecutionError(
278                    "No matching route destination and no default Chain configured".to_string(),
279                ))?,
280        };
281
282        chain.stream(inputs).await
283    }
284
285    fn name(&self) -> &str {
286        &self.name
287    }
288}
289
290/// LLM Router Chain
291///
292/// Uses an LLM to intelligently determine the routing destination.
293pub struct LLMRouterChain<M: BaseChatModel> {
294    /// LLM used for routing decisions.
295    llm: M,
296
297    /// Route destinations.
298    destinations: Vec<RouteDestination>,
299
300    /// Default Chain.
301    default_chain: Option<Arc<dyn BaseChain>>,
302
303    /// Input key name.
304    input_key: String,
305
306    /// Chain name.
307    name: String,
308
309    /// Whether to print verbose information.
310    verbose: bool,
311}
312
313impl<M: BaseChatModel> LLMRouterChain<M> {
314    pub fn new(llm: M) -> Self {
315        Self {
316            llm,
317            destinations: Vec::new(),
318            default_chain: None,
319            input_key: "input".to_string(),
320            name: "llm_router_chain".to_string(),
321            verbose: false,
322        }
323    }
324
325    pub fn add_route(
326        mut self,
327        name: impl Into<String>,
328        description: impl Into<String>,
329        chain: Arc<dyn BaseChain>,
330    ) -> Self {
331        self.destinations
332            .push(RouteDestination::new(name, description, chain));
333        self
334    }
335
336    pub fn add_route_with_keywords(
337        mut self,
338        name: impl Into<String>,
339        description: impl Into<String>,
340        chain: Arc<dyn BaseChain>,
341        keywords: Vec<&str>,
342    ) -> Self {
343        self.destinations
344            .push(RouteDestination::new(name, description, chain).with_keywords(keywords));
345        self
346    }
347
348    pub fn with_default(mut self, chain: Arc<dyn BaseChain>) -> Self {
349        self.default_chain = Some(chain);
350        self
351    }
352
353    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
354        self.input_key = key.into();
355        self
356    }
357
358    pub fn with_name(mut self, name: impl Into<String>) -> Self {
359        self.name = name.into();
360        self
361    }
362
363    pub fn with_verbose(mut self, verbose: bool) -> Self {
364        self.verbose = verbose;
365        self
366    }
367
368    pub fn destinations(&self) -> &[RouteDestination] {
369        &self.destinations
370    }
371
372    pub fn default_chain(&self) -> Option<&Arc<dyn BaseChain>> {
373        self.default_chain.as_ref()
374    }
375
376    /// Build the LLM routing prompt.
377    fn build_router_prompt(&self, input: &str) -> String {
378        let mut prompt =
379            String::from("Based on the user input, select the most appropriate handler.\n\n");
380        prompt.push_str("Available handlers:\n");
381
382        for (i, dest) in self.destinations.iter().enumerate() {
383            prompt.push_str(&format!(
384                "{}. {}: {}\n",
385                i + 1,
386                dest.name(),
387                dest.description()
388            ));
389        }
390
391        prompt.push_str("\nUser input: ");
392        prompt.push_str(input);
393        prompt
394            .push_str("\n\nReturn only the name of the most appropriate handler (no explanation).");
395
396        prompt
397    }
398
399    /// Use LLM to determine the route.
400    async fn route_with_llm(&self, input: &str) -> Result<String, ChainError> {
401        let prompt = self.build_router_prompt(input);
402
403        let messages = vec![Message::human(&prompt)];
404
405        let result = self
406            .llm
407            .invoke(messages, None)
408            .await
409            .map_err(|e| ChainError::ExecutionError(format!("LLM call failed: {}", e)))?;
410
411        Ok(result.content.trim().to_string())
412    }
413
414    /// Find a route destination by name.
415    fn find_destination(&self, name: &str) -> Option<&RouteDestination> {
416        let name_lower = name.to_lowercase();
417        // 1. Exact case-insensitive match
418        if let Some(dest) = self
419            .destinations
420            .iter()
421            .find(|d| d.name().eq_ignore_ascii_case(name))
422        {
423            return Some(dest);
424        }
425        // 2. The LLM result starts or ends with the destination name
426        self.destinations.iter().find(|d| {
427            let d_lower = d.name().to_lowercase();
428            name_lower.starts_with(&d_lower)
429                || name_lower.ends_with(&d_lower)
430                || name_lower
431                    .split_whitespace()
432                    .any(|word| word.eq_ignore_ascii_case(&d_lower))
433        })
434    }
435
436    /// LLM routing takes priority over keyword matching.
437    async fn select_route(&self, input: &str) -> Result<&RouteDestination, ChainError> {
438        if self.destinations.is_empty() {
439            return Err(ChainError::ExecutionError(
440                "No route destinations configured".to_string(),
441            ));
442        }
443
444        if self.destinations.len() == 1 {
445            return Ok(&self.destinations[0]);
446        }
447
448        // Try LLM routing first (primary strategy)
449        let llm_result = self.route_with_llm(input).await;
450        match llm_result {
451            Ok(route_name) => {
452                if let Some(dest) = self.find_destination(&route_name) {
453                    return Ok(dest);
454                }
455            }
456            Err(_) => {
457                // LLM call failed; fall through to keyword matching
458            }
459        }
460
461        // Fallback: keyword matching (longest match first)
462        let mut best_match: Option<(&RouteDestination, usize)> = None;
463        for dest in &self.destinations {
464            for keyword in dest.keywords() {
465                if input.contains(keyword) {
466                    let len = keyword.len();
467                    if best_match.is_none() || len > best_match.unwrap().1 {
468                        best_match = Some((dest, len));
469                    }
470                }
471            }
472        }
473        if let Some((dest, _)) = best_match {
474            return Ok(dest);
475        }
476
477        Err(ChainError::ExecutionError(
478            "No matching route destination found (LLM and keyword matching both failed)"
479                .to_string(),
480        ))
481    }
482}
483
484#[async_trait]
485impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for LLMRouterChain<M>
486where
487    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
488{
489    fn input_keys(&self) -> Vec<&str> {
490        vec![&self.input_key]
491    }
492
493    fn output_keys(&self) -> Vec<&str> {
494        let mut seen = std::collections::HashSet::new();
495        let mut result: Vec<&str> = Vec::new();
496
497        for dest in &self.destinations {
498            for key in dest.chain().output_keys() {
499                if seen.insert(key.to_string()) {
500                    result.push(key);
501                }
502            }
503        }
504        if let Some(default) = &self.default_chain {
505            for key in default.output_keys() {
506                if seen.insert(key.to_string()) {
507                    result.push(key);
508                }
509            }
510        }
511
512        if result.is_empty() {
513            vec!["output"]
514        } else {
515            result
516        }
517    }
518
519    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
520        self.validate_inputs(&inputs)?;
521
522        let input = inputs
523            .get(&self.input_key)
524            .and_then(|v| v.as_str())
525            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
526
527        if self.verbose {
528            println!("\n=== LLMRouterChain execution ===");
529            println!("Input: {}", input);
530            println!("Route destination count: {}", self.destinations.len());
531        }
532
533        let route_result = self.select_route(input).await;
534
535        let chain = match route_result {
536            Ok(dest) => {
537                if self.verbose {
538                    println!("Routed to: {} ({})", dest.name(), dest.description());
539                }
540                dest.chain()
541            }
542            Err(e) => {
543                if let Some(default) = &self.default_chain {
544                    if self.verbose {
545                        println!("Routing failed: {}, using default Chain", e);
546                    }
547                    default
548                } else {
549                    return Err(e);
550                }
551            }
552        };
553
554        let result = chain.invoke(inputs).await?;
555
556        if self.verbose {
557            println!("=== LLMRouterChain complete ===\n");
558        }
559
560        Ok(result)
561    }
562
563    /// Stream execution for LLMRouterChain.
564    ///
565    /// The routing LLM call must complete first (to determine the destination),
566    /// then delegates to the selected chain's `stream()` method.
567    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
568        self.validate_inputs(&inputs)?;
569
570        let input = inputs
571            .get(&self.input_key)
572            .and_then(|v| v.as_str())
573            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
574
575        let route_result = self.select_route(input).await;
576
577        let chain = match route_result {
578            Ok(dest) => dest.chain(),
579            Err(e) => {
580                if let Some(default) = &self.default_chain {
581                    default
582                } else {
583                    return Err(e);
584                }
585            }
586        };
587
588        chain.stream(inputs).await
589    }
590
591    fn name(&self) -> &str {
592        &self.name
593    }
594}
lc_chains/router_chain.rs

lc_chains/
router_chain.rs
