langchain_rust/chain/router/
chain.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use futures::StreamExt;
5use serde_json::{json, Value};
6
7use crate::{
8    chain::{Chain, ChainError, ChainEventStream, DEFAULT_OUTPUT_KEY, DEFAULT_RESULT_KEY},
9    language_models::{llm::LLM, GenerateResult, LLMEvent, TokenUsage},
10    observability::traced,
11    prompt::PromptArgs,
12    schemas::{messages::value_to_text, Message},
13};
14
15/// Output key of [`RouterChain::execute`] holding the chosen destination's name
16/// (`DEFAULT` when the default chain ran).
17pub const ROUTER_DESTINATION_KEY: &str = "destination";
18/// What the router answers when no destination fits.
19pub const DEFAULT_DESTINATION: &str = "DEFAULT";
20const ROUTER_DEFAULT_INPUT_KEY: &str = "input";
21
22const ROUTER_PROMPT: &str = "Given a raw text input, select the destination best suited for it. \
23You will be given the names of the available destinations and a description of what each one \
24is good for.
25
26<< DESTINATIONS >>
27{destinations}
28
29Answer with a JSON object and nothing else:
30{\"destination\": \"<one of the destination names above, or DEFAULT if none of them fits>\"}";
31
32struct Destination {
33    name: String,
34    description: String,
35    chain: Box<dyn Chain>,
36}
37
38/// The router's decision.
39#[derive(Debug, Clone, PartialEq)]
40pub struct Route {
41    /// The chosen destination; `None` for the default chain.
42    pub destination: Option<String>,
43    /// Usage of the routing call.
44    pub tokens: Option<TokenUsage>,
45}
46
47/// Asks an LLM which of several named chains should handle the input, then runs that
48/// chain with the same input. If none fits (or the LLM names one that doesn't exist), the
49/// default chain runs; without a default that's an error.
50///
51/// For routing by embeddings instead of an LLM call, see
52/// [`semantic_router`](crate::semantic_router).
53///
54/// ```rust,ignore
55/// let router = RouterChainBuilder::new()
56///     .llm(OpenAI::default())
57///     .destination("physics", "Questions about physics", physics_chain)
58///     .destination("math", "Math problems and proofs", math_chain)
59///     .default_chain(general_chain)
60///     .build()?;
61/// let answer = router.invoke(prompt_args! { "input" => "What is a black hole?" }).await?;
62/// ```
63pub struct RouterChain {
64    llm: Box<dyn LLM>,
65    destinations: Vec<Destination>,
66    default_chain: Option<Box<dyn Chain>>,
67    input_key: String,
68    output_key: String,
69}
70
71pub struct RouterChainBuilder {
72    llm: Option<Box<dyn LLM>>,
73    destinations: Vec<Destination>,
74    default_chain: Option<Box<dyn Chain>>,
75    input_key: String,
76    output_key: String,
77}
78
79impl Default for RouterChainBuilder {
80    fn default() -> Self {
81        Self::new()
82    }
83}
84
85impl RouterChainBuilder {
86    pub fn new() -> Self {
87        Self {
88            llm: None,
89            destinations: Vec::new(),
90            default_chain: None,
91            input_key: ROUTER_DEFAULT_INPUT_KEY.to_string(),
92            output_key: DEFAULT_OUTPUT_KEY.to_string(),
93        }
94    }
95
96    /// The LLM that picks the destination.
97    pub fn llm<L: Into<Box<dyn LLM>>>(mut self, llm: L) -> Self {
98        self.llm = Some(llm.into());
99        self
100    }
101
102    /// Adds a destination. The LLM sees `name` and `description`, so make the description
103    /// say what inputs the chain is good for.
104    pub fn destination<C: Into<Box<dyn Chain>>>(
105        mut self,
106        name: &str,
107        description: &str,
108        chain: C,
109    ) -> Self {
110        self.destinations.push(Destination {
111            name: name.trim().to_string(),
112            description: description.trim().to_string(),
113            chain: chain.into(),
114        });
115        self
116    }
117
118    /// Runs when no destination fits.
119    pub fn default_chain<C: Into<Box<dyn Chain>>>(mut self, chain: C) -> Self {
120        self.default_chain = Some(chain.into());
121        self
122    }
123
124    /// Input key holding the text to route (default `input`). Every input is forwarded
125    /// to the destination as is.
126    pub fn input_key(mut self, key: &str) -> Self {
127        self.input_key = key.to_string();
128        self
129    }
130
131    /// Output key of the answer in [`execute`](Chain::execute) (default `output`).
132    pub fn output_key(mut self, key: &str) -> Self {
133        self.output_key = key.to_string();
134        self
135    }
136
137    pub fn build(self) -> Result<RouterChain, ChainError> {
138        let llm = self
139            .llm
140            .ok_or_else(|| ChainError::MissingObject("llm".into()))?;
141        if self.destinations.is_empty() {
142            return Err(ChainError::MissingObject("destination".into()));
143        }
144        for (i, dest) in self.destinations.iter().enumerate() {
145            if dest.name.is_empty() || dest.name.eq_ignore_ascii_case(DEFAULT_DESTINATION) {
146                return Err(ChainError::OtherError(format!(
147                    "invalid router destination name {:?}",
148                    dest.name
149                )));
150            }
151            if self.destinations[..i]
152                .iter()
153                .any(|d| d.name.eq_ignore_ascii_case(&dest.name))
154            {
155                return Err(ChainError::OtherError(format!(
156                    "duplicate router destination {:?}",
157                    dest.name
158                )));
159            }
160        }
161        Ok(RouterChain {
162            llm,
163            destinations: self.destinations,
164            default_chain: self.default_chain,
165            input_key: self.input_key,
166            output_key: self.output_key,
167        })
168    }
169}
170
171/// The destination name in the router's answer: a JSON object (possibly fenced or
172/// surrounded by text), or failing that the bare answer.
173pub(super) fn parse_destination(answer: &str) -> String {
174    if let (Some(start), Some(end)) = (answer.find('{'), answer.rfind('}')) {
175        if start < end {
176            if let Ok(Value::Object(obj)) = serde_json::from_str::<Value>(&answer[start..=end]) {
177                if let Some(Value::String(name)) = obj.get("destination") {
178                    return name.trim().to_string();
179                }
180            }
181        }
182    }
183    answer
184        .trim()
185        .trim_matches(|c: char| c == '"' || c == '\'' || c == '`' || c == '.')
186        .trim()
187        .to_string()
188}
189
190fn add_tokens(a: Option<TokenUsage>, b: Option<&TokenUsage>) -> Option<TokenUsage> {
191    match (a, b) {
192        (Some(a), Some(b)) => Some(a.sum(b)),
193        (a, b) => a.or_else(|| b.cloned()),
194    }
195}
196
197impl RouterChain {
198    pub(super) fn prompt(&self, input: &str) -> Vec<Message> {
199        let destinations = self
200            .destinations
201            .iter()
202            .map(|d| format!("{}: {}", d.name, d.description))
203            .collect::<Vec<_>>()
204            .join("\n");
205        vec![
206            Message::new_system_message(ROUTER_PROMPT.replace("{destinations}", &destinations)),
207            Message::new_human_message(input),
208        ]
209    }
210
211    /// Asks the LLM where `input_variables` should go, without running anything else.
212    pub async fn route(&self, input_variables: &PromptArgs) -> Result<Route, ChainError> {
213        let input = input_variables
214            .get(&self.input_key)
215            .map(value_to_text)
216            .ok_or_else(|| ChainError::MissingInputVariable(self.input_key.clone()))?;
217        let result = self.llm.generate(&self.prompt(&input)).await?;
218        let name = parse_destination(&result.generation);
219        let destination = self
220            .destinations
221            .iter()
222            .find(|d| d.name.eq_ignore_ascii_case(&name))
223            .map(|d| d.name.clone());
224        if destination.is_none() && !name.eq_ignore_ascii_case(DEFAULT_DESTINATION) {
225            tracing::warn!(answer = %result.generation, "router answered an unknown destination");
226            if self.default_chain.is_none() {
227                return Err(ChainError::OtherError(format!(
228                    "router picked unknown destination {name:?} and there is no default chain; \
229                     expected one of {:?}",
230                    self.destinations
231                        .iter()
232                        .map(|d| &d.name)
233                        .collect::<Vec<_>>()
234                )));
235            }
236        }
237        Ok(Route {
238            destination,
239            tokens: result.tokens,
240        })
241    }
242
243    fn chain_for(&self, route: &Route) -> Result<&dyn Chain, ChainError> {
244        match &route.destination {
245            Some(name) => Ok(self
246                .destinations
247                .iter()
248                .find(|d| &d.name == name)
249                .map(|d| d.chain.as_ref())
250                .expect("route() only returns known destinations")),
251            None => self.default_chain.as_deref().ok_or_else(|| {
252                ChainError::OtherError("no destination fits and there is no default chain".into())
253            }),
254        }
255    }
256
257    fn span(&self) -> tracing::Span {
258        tracing::info_span!(
259            "chain.call",
260            chain = "router",
261            destination = tracing::field::Empty,
262            error = tracing::field::Empty
263        )
264    }
265}
266
267fn destination_label(route: &Route) -> &str {
268    route.destination.as_deref().unwrap_or(DEFAULT_DESTINATION)
269}
270
271#[async_trait]
272impl Chain for RouterChain {
273    async fn call(&self, input_variables: PromptArgs) -> Result<GenerateResult, ChainError> {
274        let span = self.span();
275        traced(
276            span.clone(),
277            async {
278                let route = self.route(&input_variables).await?;
279                span.record("destination", destination_label(&route));
280                let mut result = self.chain_for(&route)?.call(input_variables).await?;
281                result.tokens = add_tokens(result.tokens, route.tokens.as_ref());
282                Ok(result)
283            },
284            |_, _| {},
285        )
286        .await
287    }
288
289    /// The destination's own outputs, plus the answer under the output key, the full
290    /// result (routing tokens included) and [`ROUTER_DESTINATION_KEY`].
291    async fn execute(
292        &self,
293        input_variables: PromptArgs,
294    ) -> Result<HashMap<String, Value>, ChainError> {
295        let span = self.span();
296        traced(
297            span.clone(),
298            async {
299                let route = self.route(&input_variables).await?;
300                span.record("destination", destination_label(&route));
301                let chain = self.chain_for(&route)?;
302                let mut output = chain.execute(input_variables).await?;
303                let mut result: GenerateResult = match output.get(DEFAULT_RESULT_KEY) {
304                    Some(result) => serde_json::from_value(result.clone())?,
305                    None => {
306                        let key = chain
307                            .get_output_keys()
308                            .first()
309                            .cloned()
310                            .unwrap_or_else(|| DEFAULT_OUTPUT_KEY.to_string());
311                        GenerateResult {
312                            generation: output.get(&key).map(value_to_text).unwrap_or_default(),
313                            ..Default::default()
314                        }
315                    }
316                };
317                result.tokens = add_tokens(result.tokens, route.tokens.as_ref());
318                output.insert(self.output_key.clone(), json!(result.generation));
319                output.insert(DEFAULT_RESULT_KEY.to_string(), json!(result));
320                output.insert(
321                    ROUTER_DESTINATION_KEY.to_string(),
322                    json!(destination_label(&route)),
323                );
324                Ok(output)
325            },
326            |_, _| {},
327        )
328        .await
329    }
330
331    /// Routes, then streams the destination's answer. The final `Done` includes the
332    /// routing tokens.
333    fn call_stream<'a>(&'a self, input_variables: PromptArgs) -> ChainEventStream<'a> {
334        Box::pin(
335            futures::stream::once(async move {
336                let routed = async {
337                    let route = self.route(&input_variables).await?;
338                    let chain = self.chain_for(&route)?;
339                    Ok::<_, ChainError>((route, chain))
340                };
341                match routed.await {
342                    Ok((route, chain)) => {
343                        let tokens = route.tokens;
344                        Box::pin(chain.call_stream(input_variables).map(move |event| {
345                            event.map(|event| match event {
346                                LLMEvent::Done(mut result) => {
347                                    result.tokens = add_tokens(result.tokens, tokens.as_ref());
348                                    LLMEvent::Done(result)
349                                }
350                                other => other,
351                            })
352                        })) as ChainEventStream<'a>
353                    }
354                    Err(e) => Box::pin(futures::stream::iter([Err(e)])) as ChainEventStream<'a>,
355                }
356            })
357            .flatten(),
358        )
359    }
360
361    fn get_input_keys(&self) -> Vec<String> {
362        vec![self.input_key.clone()]
363    }
364
365    fn get_output_keys(&self) -> Vec<String> {
366        vec![
367            self.output_key.clone(),
368            DEFAULT_RESULT_KEY.to_string(),
369            ROUTER_DESTINATION_KEY.to_string(),
370        ]
371    }
372}
langchain_rust/chain/router/chain.rs

langchain_rust/chain/router/
chain.rs
