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::runnables::RunnableConfig;
9use lc_core::tools::ToolDefinition;
10use lc_core::{BaseChatModel, Runnable};
11use lc_schema::Message;
12use serde::Deserialize;
13use serde_json::json;
14use serde_json::Value;
15use std::collections::HashMap;
16use std::sync::Arc;
17
18use crate::base::{
19    run_chain_with_callbacks, stream_chain_with_callbacks, BaseChain, ChainError, ChainResult,
20    ChainStream,
21};
22
23/// Route destination.
24pub struct RouteDestination {
25    /// Destination name.
26    name: String,
27    /// Destination description (used for routing decisions).
28    description: String,
29    /// Destination Chain.
30    chain: Arc<dyn BaseChain>,
31    /// Keyword list (used for keyword-based routing).
32    keywords: Vec<String>,
33}
34
35impl RouteDestination {
36    pub fn new(
37        name: impl Into<String>,
38        description: impl Into<String>,
39        chain: Arc<dyn BaseChain>,
40    ) -> Self {
41        Self {
42            name: name.into(),
43            description: description.into(),
44            chain,
45            keywords: Vec::new(),
46        }
47    }
48
49    pub fn with_keywords(mut self, keywords: Vec<&str>) -> Self {
50        self.keywords = keywords.into_iter().map(String::from).collect();
51        self
52    }
53
54    pub fn name(&self) -> &str {
55        &self.name
56    }
57
58    pub fn description(&self) -> &str {
59        &self.description
60    }
61
62    pub fn chain(&self) -> &Arc<dyn BaseChain> {
63        &self.chain
64    }
65
66    pub fn keywords(&self) -> &[String] {
67        &self.keywords
68    }
69}
70
71/// Router Chain
72///
73/// Automatically routes to different Chains based on input content.
74pub struct RouterChain {
75    /// Route destination list.
76    destinations: Vec<RouteDestination>,
77
78    /// Default Chain (used when no match is found).
79    default_chain: Option<Arc<dyn BaseChain>>,
80
81    /// Input key name.
82    input_key: String,
83
84    /// Chain name.
85    name: String,
86
87    /// Whether to print verbose information.
88    verbose: bool,
89}
90
91impl RouterChain {
92    pub fn new() -> Self {
93        Self {
94            destinations: Vec::new(),
95            default_chain: None,
96            input_key: "input".to_string(),
97            name: "router_chain".to_string(),
98            verbose: false,
99        }
100    }
101
102    pub fn add_route(
103        mut self,
104        name: impl Into<String>,
105        description: impl Into<String>,
106        chain: Arc<dyn BaseChain>,
107    ) -> Self {
108        self.destinations
109            .push(RouteDestination::new(name, description, chain));
110        self
111    }
112
113    pub fn add_route_with_keywords(
114        mut self,
115        name: impl Into<String>,
116        description: impl Into<String>,
117        chain: Arc<dyn BaseChain>,
118        keywords: Vec<&str>,
119    ) -> Self {
120        self.destinations
121            .push(RouteDestination::new(name, description, chain).with_keywords(keywords));
122        self
123    }
124
125    pub fn with_default(mut self, chain: Arc<dyn BaseChain>) -> Self {
126        self.default_chain = Some(chain);
127        self
128    }
129
130    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
131        self.input_key = key.into();
132        self
133    }
134
135    pub fn with_name(mut self, name: impl Into<String>) -> Self {
136        self.name = name.into();
137        self
138    }
139
140    pub fn with_verbose(mut self, verbose: bool) -> Self {
141        self.verbose = verbose;
142        self
143    }
144
145    pub fn destinations(&self) -> &[RouteDestination] {
146        &self.destinations
147    }
148
149    pub fn default_chain(&self) -> Option<&Arc<dyn BaseChain>> {
150        self.default_chain.as_ref()
151    }
152
153    /// Keyword-based routing.
154    ///
155    /// Longest-match-first instead of first-match-wins.
156    fn route_by_keywords(&self, input: &str) -> Option<&RouteDestination> {
157        let mut best_match: Option<(&RouteDestination, usize)> = None;
158        for dest in &self.destinations {
159            for keyword in &dest.keywords {
160                if input.contains(keyword) {
161                    let len = keyword.len();
162                    if best_match.is_none() || len > best_match.unwrap().1 {
163                        best_match = Some((dest, len));
164                    }
165                }
166            }
167        }
168        best_match.map(|(dest, _)| dest)
169    }
170
171    /// Select a route destination.
172    fn select_route(&self, input: &str) -> Result<Option<&RouteDestination>, ChainError> {
173        if let Some(dest) = self.route_by_keywords(input) {
174            return Ok(Some(dest));
175        }
176
177        Ok(None)
178    }
179
180    /// Route to a destination/default chain and invoke it, threading `config`
181    /// through `invoke_with_config` (never silently dropping it).
182    async fn route_and_invoke(
183        &self,
184        inputs: HashMap<String, Value>,
185        config: Option<RunnableConfig>,
186    ) -> Result<ChainResult, ChainError> {
187        self.validate_inputs(&inputs)?;
188
189        let input = inputs
190            .get(&self.input_key)
191            .and_then(|v| v.as_str())
192            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
193
194        if self.verbose {
195            println!("\n=== RouterChain execution ===");
196            println!("Input: {}", input);
197            println!("Route destination count: {}", self.destinations.len());
198        }
199
200        let route_result = self.select_route(input)?;
201
202        let chain = match route_result {
203            Some(dest) => {
204                if self.verbose {
205                    println!("Routed to: {} ({})", dest.name(), dest.description());
206                }
207                dest.chain()
208            }
209            None => {
210                if let Some(default) = &self.default_chain {
211                    if self.verbose {
212                        println!("No keyword match, using default Chain");
213                    }
214                    default
215                } else {
216                    return Err(ChainError::ExecutionError(
217                        "No matching route destination and no default Chain configured".to_string(),
218                    ));
219                }
220            }
221        };
222
223        let result = chain.invoke_with_config(inputs, config).await?;
224
225        if self.verbose {
226            println!("=== RouterChain complete ===\n");
227        }
228
229        Ok(result)
230    }
231
232    /// Route to a destination/default chain and stream it, threading `config`
233    /// through `stream_with_config`.
234    async fn route_and_stream(
235        &self,
236        inputs: HashMap<String, Value>,
237        config: Option<RunnableConfig>,
238    ) -> Result<ChainStream, ChainError> {
239        self.validate_inputs(&inputs)?;
240
241        let input = inputs
242            .get(&self.input_key)
243            .and_then(|v| v.as_str())
244            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
245
246        let route_result = self.select_route(input)?;
247
248        let chain = match route_result {
249            Some(dest) => dest.chain(),
250            None => self.default_chain.as_ref().ok_or_else(|| {
251                ChainError::ExecutionError(
252                    "No matching route destination and no default Chain configured".to_string(),
253                )
254            })?,
255        };
256
257        chain.stream_with_config(inputs, config).await
258    }
259}
260
261impl Default for RouterChain {
262    fn default() -> Self {
263        Self::new()
264    }
265}
266
267#[async_trait]
268impl BaseChain for RouterChain {
269    fn input_keys(&self) -> Vec<&str> {
270        vec![&self.input_key]
271    }
272
273    fn output_keys(&self) -> Vec<&str> {
274        let mut seen = std::collections::HashSet::new();
275        let mut result: Vec<&str> = Vec::new();
276
277        for dest in &self.destinations {
278            for key in dest.chain().output_keys() {
279                if seen.insert(key.to_string()) {
280                    result.push(key);
281                }
282            }
283        }
284        if let Some(default) = &self.default_chain {
285            for key in default.output_keys() {
286                if seen.insert(key.to_string()) {
287                    result.push(key);
288                }
289            }
290        }
291
292        if result.is_empty() {
293            vec!["output"]
294        } else {
295            result
296        }
297    }
298
299    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
300        self.route_and_invoke(inputs, None).await
301    }
302
303    /// Execute the Chain with config propagation.
304    ///
305    /// Dispatches this chain's `on_chain_start`/`on_chain_end` and threads
306    /// `config` into the routed destination chain via `invoke_with_config`.
307    async fn invoke_with_config(
308        &self,
309        inputs: HashMap<String, Value>,
310        config: Option<RunnableConfig>,
311    ) -> Result<ChainResult, ChainError> {
312        run_chain_with_callbacks(self.name(), inputs, config.clone(), |inputs| async move {
313            self.route_and_invoke(inputs, config).await
314        })
315        .await
316    }
317
318    /// Stream execution for RouterChain.
319    ///
320    /// After routing (keyword matching), delegates to the selected chain's
321    /// `stream()` method.
322    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
323        self.route_and_stream(inputs, None).await
324    }
325
326    /// Stream execute the Chain with config propagation.
327    async fn stream_with_config(
328        &self,
329        inputs: HashMap<String, Value>,
330        config: Option<RunnableConfig>,
331    ) -> Result<ChainStream, ChainError> {
332        stream_chain_with_callbacks(self.name(), inputs, config.clone(), |inputs| async move {
333            self.route_and_stream(inputs, config).await
334        })
335        .await
336    }
337
338    fn name(&self) -> &str {
339        &self.name
340    }
341}
342
343/// Structured routing decision returned by the LLM (P2-5).
344///
345/// Preferred source is the `route_to_destination` tool call's JSON arguments
346/// (`{"destination": ..., "reason": ...}`). `from_text` is a lenient fallback
347/// for providers without tool binding: it tries the same JSON object shape,
348/// then a bare destination name.
349#[derive(Debug, Clone, Deserialize)]
350pub(crate) struct RouteDecision {
351    /// Handler name — must match a configured destination.
352    pub destination: String,
353    /// Optional explanation for the choice (used in verbose diagnostics).
354    pub reason: Option<String>,
355}
356
357impl RouteDecision {
358    /// Parse a lenient text reply: JSON object first, bare name fallback.
359    fn from_text(text: &str) -> Self {
360        let trimmed = text.trim();
361        if let Ok(decision) = serde_json::from_str::<RouteDecision>(trimmed) {
362            return decision;
363        }
364        Self {
365            destination: trimmed.to_string(),
366            reason: None,
367        }
368    }
369}
370
371/// Tool definition that forces the routing LLM to emit a structured
372/// `{destination, reason}` object instead of free text.
373fn route_tool() -> ToolDefinition {
374    ToolDefinition::new(
375        "route_to_destination",
376        "根据用户输入选择最合适的处理 handler,返回目标名称与理由",
377    )
378    .with_parameters(json!({
379        "type": "object",
380        "properties": {
381            "destination": { "type": "string", "description": "目标 handler 名称" },
382            "reason": { "type": "string", "description": "选择该 handler 的理由" }
383        },
384        "required": ["destination"]
385    }))
386}
387
388/// LLM Router Chain
389///
390/// Uses an LLM to intelligently determine the routing destination.
391pub struct LLMRouterChain<M: BaseChatModel> {
392    /// LLM used for routing decisions.
393    llm: M,
394
395    /// Route destinations.
396    destinations: Vec<RouteDestination>,
397
398    /// Default Chain.
399    default_chain: Option<Arc<dyn BaseChain>>,
400
401    /// Input key name.
402    input_key: String,
403
404    /// Chain name.
405    name: String,
406
407    /// Whether to print verbose information.
408    verbose: bool,
409}
410
411impl<M: BaseChatModel + Send + Sync + 'static> LLMRouterChain<M> {
412    pub fn new(llm: M) -> Self {
413        Self {
414            llm,
415            destinations: Vec::new(),
416            default_chain: None,
417            input_key: "input".to_string(),
418            name: "llm_router_chain".to_string(),
419            verbose: false,
420        }
421    }
422
423    pub fn add_route(
424        mut self,
425        name: impl Into<String>,
426        description: impl Into<String>,
427        chain: Arc<dyn BaseChain>,
428    ) -> Self {
429        self.destinations
430            .push(RouteDestination::new(name, description, chain));
431        self
432    }
433
434    pub fn add_route_with_keywords(
435        mut self,
436        name: impl Into<String>,
437        description: impl Into<String>,
438        chain: Arc<dyn BaseChain>,
439        keywords: Vec<&str>,
440    ) -> Self {
441        self.destinations
442            .push(RouteDestination::new(name, description, chain).with_keywords(keywords));
443        self
444    }
445
446    pub fn with_default(mut self, chain: Arc<dyn BaseChain>) -> Self {
447        self.default_chain = Some(chain);
448        self
449    }
450
451    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
452        self.input_key = key.into();
453        self
454    }
455
456    pub fn with_name(mut self, name: impl Into<String>) -> Self {
457        self.name = name.into();
458        self
459    }
460
461    pub fn with_verbose(mut self, verbose: bool) -> Self {
462        self.verbose = verbose;
463        self
464    }
465
466    pub fn destinations(&self) -> &[RouteDestination] {
467        &self.destinations
468    }
469
470    pub fn default_chain(&self) -> Option<&Arc<dyn BaseChain>> {
471        self.default_chain.as_ref()
472    }
473
474    /// Build the LLM routing prompt.
475    fn build_router_prompt(&self, input: &str) -> String {
476        let mut prompt =
477            String::from("Based on the user input, select the most appropriate handler.\n\n");
478        prompt.push_str("Available handlers:\n");
479
480        for (i, dest) in self.destinations.iter().enumerate() {
481            prompt.push_str(&format!(
482                "{}. {}: {}\n",
483                i + 1,
484                dest.name(),
485                dest.description()
486            ));
487        }
488
489        prompt.push_str("\nUser input: ");
490        prompt.push_str(input);
491        // P2-5: ask for a structured object rather than a bare name so the
492        // decision (destination + reason) survives parsing reliably.
493        prompt.push_str(
494            "\n\nReturn only the chosen handler as JSON: {\"destination\": \"<handler name>\", \"reason\": \"<why this handler>\"}",
495        );
496
497        prompt
498    }
499
500    /// Use LLM to determine the route (P2-5: structured output).
501    ///
502    /// Binds `route_to_destination` so the provider emits a structured
503    /// `{destination, reason}` tool-call argument instead of free text; the
504    /// first tool call's parsed arguments win. Providers without tool binding
505    /// (or a response without `tool_calls`) fall back to the same call's text,
506    /// parsed leniently by [`RouteDecision::from_text`]. One LLM call, no
507    /// retry.
508    async fn route_with_llm(
509        &self,
510        input: &str,
511        config: Option<RunnableConfig>,
512    ) -> Result<RouteDecision, ChainError> {
513        let prompt = self.build_router_prompt(input);
514
515        let messages = vec![Message::human(&prompt)];
516
517        let map_err = |e| ChainError::Nested {
518            context: "LLM routing call failed".to_string(),
519            source: Box::new(e),
520        };
521
522        let result = match self.llm.bind_tools(vec![route_tool()]) {
523            Some(bound) => bound.chat(messages, config).await.map_err(map_err)?,
524            None => self.llm.chat(messages, config).await.map_err(map_err)?,
525        };
526
527        if let Some(decision) = result
528            .tool_calls
529            .as_ref()
530            .and_then(|calls| calls.first())
531            .and_then(|call| call.parse_arguments::<RouteDecision>().ok())
532        {
533            return Ok(decision);
534        }
535
536        Ok(RouteDecision::from_text(&result.content))
537    }
538
539    /// Find a route destination by name.
540    fn find_destination(&self, name: &str) -> Option<&RouteDestination> {
541        let name_lower = name.to_lowercase();
542        // 1. Exact case-insensitive match
543        if let Some(dest) = self
544            .destinations
545            .iter()
546            .find(|d| d.name().eq_ignore_ascii_case(name))
547        {
548            return Some(dest);
549        }
550        // 2. The LLM result starts or ends with the destination name
551        self.destinations.iter().find(|d| {
552            let d_lower = d.name().to_lowercase();
553            name_lower.starts_with(&d_lower)
554                || name_lower.ends_with(&d_lower)
555                || name_lower
556                    .split_whitespace()
557                    .any(|word| word.eq_ignore_ascii_case(&d_lower))
558        })
559    }
560
561    /// LLM routing takes priority over keyword matching.
562    async fn select_route(
563        &self,
564        input: &str,
565        config: Option<RunnableConfig>,
566    ) -> Result<&RouteDestination, ChainError> {
567        if self.destinations.is_empty() {
568            return Err(ChainError::ExecutionError(
569                "No route destinations configured".to_string(),
570            ));
571        }
572
573        if self.destinations.len() == 1 {
574            return Ok(&self.destinations[0]);
575        }
576
577        // Try LLM routing first (primary strategy).
578        // P1-6: the LLM error/unknown-name is retained rather than swallowed, so
579        // keyword fallback stays as a legitimate safety net but the final error
580        // carries the real routing diagnostics.
581        // P2-5: the LLM now returns a structured decision {destination, reason};
582        // the reason is surfaced in verbose mode only.
583        let llm_note: Option<String> = {
584            let llm_result = self.route_with_llm(input, config).await;
585            match llm_result {
586                Ok(decision) => {
587                    if let Some(reason) = &decision.reason {
588                        if self.verbose {
589                            println!("LLM route reason: {}", reason);
590                        }
591                    }
592                    if let Some(dest) = self.find_destination(&decision.destination) {
593                        return Ok(dest);
594                    }
595                    Some(format!(
596                        "LLM returned an unknown route destination {:?}",
597                        decision.destination
598                    ))
599                }
600                Err(e) => Some(format!("LLM routing call failed: {}", e)),
601            }
602        };
603
604        // Fallback: keyword matching (longest match first)
605        let mut best_match: Option<(&RouteDestination, usize)> = None;
606        for dest in &self.destinations {
607            for keyword in dest.keywords() {
608                if input.contains(keyword) {
609                    let len = keyword.len();
610                    if best_match.is_none() || len > best_match.unwrap().1 {
611                        best_match = Some((dest, len));
612                    }
613                }
614            }
615        }
616        if let Some((dest, _)) = best_match {
617            return Ok(dest);
618        }
619
620        Err(ChainError::ExecutionError(format!(
621            "No matching route destination found ({})",
622            llm_note.unwrap_or_else(|| "LLM and keyword matching both failed".to_string())
623        )))
624    }
625
626    /// Route to a destination/default chain and invoke it, threading `config`
627    /// through `invoke_with_config` (never silently dropping it).
628    async fn route_and_invoke(
629        &self,
630        inputs: HashMap<String, Value>,
631        config: Option<RunnableConfig>,
632    ) -> Result<ChainResult, ChainError> {
633        self.validate_inputs(&inputs)?;
634
635        let input = inputs
636            .get(&self.input_key)
637            .and_then(|v| v.as_str())
638            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
639
640        if self.verbose {
641            println!("\n=== LLMRouterChain execution ===");
642            println!("Input: {}", input);
643            println!("Route destination count: {}", self.destinations.len());
644        }
645
646        let route_result = self.select_route(input, config.clone()).await;
647
648        let chain = match route_result {
649            Ok(dest) => {
650                if self.verbose {
651                    println!("Routed to: {} ({})", dest.name(), dest.description());
652                }
653                dest.chain()
654            }
655            Err(e) => {
656                if let Some(default) = &self.default_chain {
657                    // 路由失败走默认链:不静默,记 error 日志说明原因,
658                    // 避免调用方把 fallback 答案当成路由选择的正确结果
659                    log::error!("路由失败,改用默认链(调用方可能得到与输入不匹配的答案): {e}");
660                    default
661                } else {
662                    return Err(e);
663                }
664            }
665        };
666
667        let result = chain.invoke_with_config(inputs, config).await?;
668
669        if self.verbose {
670            println!("=== LLMRouterChain complete ===\n");
671        }
672
673        Ok(result)
674    }
675
676    /// Route to a destination/default chain and stream it, threading `config`
677    /// through `stream_with_config`.
678    async fn route_and_stream(
679        &self,
680        inputs: HashMap<String, Value>,
681        config: Option<RunnableConfig>,
682    ) -> Result<ChainStream, ChainError> {
683        self.validate_inputs(&inputs)?;
684
685        let input = inputs
686            .get(&self.input_key)
687            .and_then(|v| v.as_str())
688            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
689
690        let route_result = self.select_route(input, config.clone()).await;
691
692        let chain = match route_result {
693            Ok(dest) => dest.chain(),
694            Err(e) => {
695                if let Some(default) = &self.default_chain {
696                    default
697                } else {
698                    return Err(e);
699                }
700            }
701        };
702
703        chain.stream_with_config(inputs, config).await
704    }
705}
706
707#[async_trait]
708impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for LLMRouterChain<M>
709where
710    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
711{
712    fn input_keys(&self) -> Vec<&str> {
713        vec![&self.input_key]
714    }
715
716    fn output_keys(&self) -> Vec<&str> {
717        let mut seen = std::collections::HashSet::new();
718        let mut result: Vec<&str> = Vec::new();
719
720        for dest in &self.destinations {
721            for key in dest.chain().output_keys() {
722                if seen.insert(key.to_string()) {
723                    result.push(key);
724                }
725            }
726        }
727        if let Some(default) = &self.default_chain {
728            for key in default.output_keys() {
729                if seen.insert(key.to_string()) {
730                    result.push(key);
731                }
732            }
733        }
734
735        if result.is_empty() {
736            vec!["output"]
737        } else {
738            result
739        }
740    }
741
742    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
743        self.route_and_invoke(inputs, None).await
744    }
745
746    /// Execute the Chain with config propagation.
747    ///
748    /// Dispatches this chain's `on_chain_start`/`on_chain_end`, threads
749    /// `config` into the routing LLM call, and into the routed destination
750    /// chain via `invoke_with_config`.
751    async fn invoke_with_config(
752        &self,
753        inputs: HashMap<String, Value>,
754        config: Option<RunnableConfig>,
755    ) -> Result<ChainResult, ChainError> {
756        run_chain_with_callbacks(self.name(), inputs, config.clone(), |inputs| async move {
757            self.route_and_invoke(inputs, config).await
758        })
759        .await
760    }
761
762    /// Stream execution for LLMRouterChain.
763    ///
764    /// The routing LLM call must complete first (to determine the destination),
765    /// then delegates to the selected chain's `stream()` method.
766    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
767        self.route_and_stream(inputs, None).await
768    }
769
770    /// Stream execute the Chain with config propagation.
771    async fn stream_with_config(
772        &self,
773        inputs: HashMap<String, Value>,
774        config: Option<RunnableConfig>,
775    ) -> Result<ChainStream, ChainError> {
776        stream_chain_with_callbacks(self.name(), inputs, config.clone(), |inputs| async move {
777            self.route_and_stream(inputs, config).await
778        })
779        .await
780    }
781
782    fn name(&self) -> &str {
783        &self.name
784    }
785}
786
787#[cfg(test)]
788mod tests {
789    use super::*;
790    use async_trait::async_trait;
791    use futures_util::Stream;
792    use lc_core::runnables::RunnableConfig;
793    use lc_core::{BaseLanguageModel, Runnable};
794    use std::pin::Pin;
795    use std::sync::Arc;
796
797    #[derive(Debug)]
798    struct MockError(String);
799    impl std::fmt::Display for MockError {
800        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
801            write!(f, "{}", self.0)
802        }
803    }
804    impl std::error::Error for MockError {}
805
806    /// Mock chat model that returns a canned `LLMResult`. `supports_tools`
807    /// controls whether `bind_tools` yields a bound copy (tool-call path) or
808    /// `None` (plain-text fallback path) — so both P2-5 branches are testable.
809    #[derive(Clone)]
810    struct MockRouterLLM {
811        response: LLMResult,
812        supports_tools: bool,
813    }
814
815    impl MockRouterLLM {
816        fn tools(response: LLMResult) -> Self {
817            Self {
818                response,
819                supports_tools: true,
820            }
821        }
822        fn plain(response: LLMResult) -> Self {
823            Self {
824                response,
825                supports_tools: false,
826            }
827        }
828    }
829
830    #[async_trait]
831    impl Runnable<Vec<Message>, LLMResult> for MockRouterLLM {
832        type Error = MockError;
833        async fn invoke(
834            &self,
835            _input: Vec<Message>,
836            _config: Option<RunnableConfig>,
837        ) -> Result<LLMResult, Self::Error> {
838            Ok(self.response.clone())
839        }
840    }
841
842    #[async_trait]
843    impl BaseLanguageModel<Vec<Message>, LLMResult> for MockRouterLLM {
844        fn model_name(&self) -> &str {
845            "mock"
846        }
847        fn get_num_tokens(&self, t: &str) -> usize {
848            t.len()
849        }
850        fn with_temperature(self, _: f32) -> Self {
851            self
852        }
853        fn with_max_tokens(self, _: usize) -> Self {
854            self
855        }
856    }
857
858    #[async_trait]
859    impl BaseChatModel for MockRouterLLM {
860        async fn chat(
861            &self,
862            _messages: Vec<Message>,
863            _config: Option<RunnableConfig>,
864        ) -> Result<LLMResult, Self::Error> {
865            Ok(self.response.clone())
866        }
867        async fn stream_chat(
868            &self,
869            _messages: Vec<Message>,
870            _config: Option<RunnableConfig>,
871        ) -> Result<Pin<Box<dyn Stream<Item = Result<String, Self::Error>> + Send>>, Self::Error>
872        {
873            let tokens = [Ok(self.response.content.clone())];
874            Ok(Box::pin(futures_util::stream::iter(tokens)))
875        }
876        fn bind_tools(
877            &self,
878            _tools: Vec<ToolDefinition>,
879        ) -> Option<Box<dyn BaseChatModel<Error = Self::Error> + Send + Sync>> {
880            if self.supports_tools {
881                Some(Box::new(self.clone()))
882            } else {
883                None
884            }
885        }
886    }
887
888    /// Simple destination chain that echoes the input under `output`.
889    struct EchoChain;
890
891    #[async_trait]
892    impl BaseChain for EchoChain {
893        fn input_keys(&self) -> Vec<&str> {
894            vec!["input"]
895        }
896        fn output_keys(&self) -> Vec<&str> {
897            vec!["output"]
898        }
899        async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
900            let mut result = HashMap::new();
901            if let Some(v) = inputs.get("input") {
902                result.insert("output".to_string(), v.clone());
903            }
904            Ok(result)
905        }
906    }
907
908    fn router_with(llm: MockRouterLLM) -> LLMRouterChain<MockRouterLLM> {
909        LLMRouterChain::new(llm)
910            .add_route(
911                "math",
912                "handles mathematical questions",
913                Arc::new(EchoChain),
914            )
915            .add_route(
916                "science",
917                "handles scientific questions",
918                Arc::new(EchoChain),
919            )
920    }
921
922    fn decision_response(destination: &str, reason: &str) -> LLMResult {
923        LLMResult {
924            content: String::new(),
925            model: "mock".to_string(),
926            token_usage: None,
927            tool_calls: Some(vec![lc_core::tools::ToolCall::new(
928                "call_1",
929                "route_to_destination",
930                format!(
931                    r#"{{"destination": "{}", "reason": "{}"}}"#,
932                    destination, reason
933                ),
934            )]),
935            thinking_content: None,
936        }
937    }
938
939    fn text_response(content: &str) -> LLMResult {
940        LLMResult {
941            content: content.to_string(),
942            model: "mock".to_string(),
943            token_usage: None,
944            tool_calls: None,
945            thinking_content: None,
946        }
947    }
948
949    /// P2-5: the routing LLM's `route_to_destination` tool-call arguments (the
950    /// structured `{destination, reason}` object) drive the route.
951    #[tokio::test]
952    async fn test_llm_router_routes_via_tool_call() {
953        let chain = router_with(MockRouterLLM::tools(decision_response(
954            "math",
955            "user asked a calculation",
956        )));
957        let inputs = HashMap::from([("input".to_string(), json!("what is 2 + 2?"))]);
958
959        let result = chain.invoke(inputs).await.unwrap();
960        assert_eq!(
961            result.get("output").unwrap(),
962            &json!("what is 2 + 2?"),
963            "routed to the math destination"
964        );
965    }
966
967    /// P2-5: a provider without tool binding falls back to the same call's
968    /// text, parsed as a JSON `{destination, reason}` object.
969    #[tokio::test]
970    async fn test_llm_router_json_text_fallback() {
971        let chain = router_with(MockRouterLLM::plain(text_response(
972            r#"{"destination": "science", "reason": "explains a phenomenon"}"#,
973        )));
974        let inputs = HashMap::from([("input".to_string(), json!("why is the sky blue?"))]);
975
976        let result = chain.invoke(inputs).await.unwrap();
977        assert_eq!(
978            result.get("output").unwrap(),
979            &json!("why is the sky blue?"),
980            "routed to the science destination"
981        );
982    }
983
984    /// P2-5: a bare destination name still works (lenient text fallback).
985    #[tokio::test]
986    async fn test_llm_router_bare_name_fallback() {
987        let chain = router_with(MockRouterLLM::plain(text_response("math")));
988        let inputs = HashMap::from([("input".to_string(), json!("calculate something"))]);
989
990        let result = chain.invoke(inputs).await.unwrap();
991        assert_eq!(result.get("output").unwrap(), &json!("calculate something"));
992    }
993
994    /// P2-5: an unknown destination from the LLM does not abort routing — the
995    /// real routing diagnostics stay in the error (P1-6).
996    #[tokio::test]
997    async fn test_llm_router_unknown_destination_reports_diagnostics() {
998        let chain = router_with(MockRouterLLM::tools(decision_response(
999            "physics", "closest",
1000        )));
1001        let inputs = HashMap::from([("input".to_string(), json!("what is 2+2"))]);
1002
1003        let err = match chain.invoke(inputs).await {
1004            Ok(_) => panic!("expected a routing failure"),
1005            Err(e) => e,
1006        };
1007        let msg = format!("{err:?}");
1008        assert!(
1009            msg.contains("unknown route destination"),
1010            "expected the LLM diagnostics in the error, got: {msg}"
1011        );
1012    }
1013
1014    /// P2-5: when the LLM names a nonexistent destination but a keyword
1015    /// matches, the keyword fallback still lands a destination.
1016    #[tokio::test]
1017    async fn test_llm_router_unknown_destination_keyword_still_routes() {
1018        let chain = LLMRouterChain::new(MockRouterLLM::tools(decision_response(
1019            "physics", "nonsense",
1020        )))
1021        .add_route_with_keywords(
1022            "math",
1023            "handles mathematical questions",
1024            Arc::new(EchoChain),
1025            vec!["calculate", "2+2"],
1026        )
1027        .add_route(
1028            "science",
1029            "handles scientific questions",
1030            Arc::new(EchoChain),
1031        );
1032
1033        let inputs = HashMap::from([("input".to_string(), json!("please calculate 2+2"))]);
1034        let result = chain.invoke(inputs).await.unwrap();
1035        assert_eq!(
1036            result.get("output").unwrap(),
1037            &json!("please calculate 2+2")
1038        );
1039    }
1040
1041    #[test]
1042    fn test_route_decision_from_text_json() {
1043        let d = RouteDecision::from_text(r#"{"destination": "math", "reason": "why"}"#);
1044        assert_eq!(d.destination, "math");
1045        assert_eq!(d.reason.as_deref(), Some("why"));
1046    }
1047
1048    #[test]
1049    fn test_route_decision_from_text_bare_name() {
1050        let d = RouteDecision::from_text("  science  ");
1051        assert_eq!(d.destination, "science");
1052        assert!(d.reason.is_none());
1053    }
1054
1055    #[test]
1056    fn test_route_tool_schema() {
1057        let tool = route_tool();
1058        assert_eq!(tool.function.name, "route_to_destination");
1059        assert!(tool.function.parameters.is_some());
1060    }
1061}
lc_chains/router_chain.rs

lc_chains/
router_chain.rs
