langchain_rust/chain/llm_math/
mod.rs

1//! LLM math chain: the model translates a word problem into an expression, and the
2//! [calculator](crate::tools::evaluate_expression) computes it.
3
4use std::{collections::HashMap, sync::LazyLock};
5
6use async_trait::async_trait;
7use regex::Regex;
8use serde_json::{json, Value};
9
10use crate::{
11    chain::{Chain, ChainError, LLMChainBuilder, DEFAULT_RESULT_KEY},
12    language_models::{llm::LLM, GenerateResult, TokenUsage},
13    output_parsers::strip_reasoning,
14    prompt::{PromptArgs, PromptTemplate, TemplateFormat},
15    tools::{evaluate_expression, format_number},
16};
17
18/// Default input key of [`LLMMathChain`].
19pub const LLM_MATH_DEFAULT_INPUT_KEY: &str = "question";
20/// Default output key of [`LLMMathChain`], as in langchaingo.
21pub const LLM_MATH_DEFAULT_OUTPUT_KEY: &str = "answer";
22/// [`execute`](Chain::execute) output key with the evaluated expression, when the model
23/// wrote one.
24pub const LLM_MATH_EXPRESSION_KEY: &str = "expression";
25
26const PROMPT: &str = r#"Translate a math problem into a single-line arithmetic expression and put it in a ```text block. A calculator evaluates it: it supports + - * / % // ^ (power), parentheses, the constants pi and e, and functions such as sqrt, cbrt, abs, exp, ln, log, log10, sin, cos, tan, floor, ceil, round, min, max, pow and factorial. Don't compute the result yourself.
27
28Question: ${Question with math problem.}
29```text
30${single line expression that solves the problem}
31```
32
33Question: What is 37593 * 67?
34```text
3537593 * 67
36```
37
38Question: What is the fifth root of 37593?
39```text
4037593 ^ (1 / 5)
41```
42
43Question: {{question}}
44"#;
45
46/// A fenced block (any language tag); the closing fence may be missing.
47static FENCED: LazyLock<Regex> = LazyLock::new(|| {
48    Regex::new(r"(?s)```[A-Za-z0-9_-]*[ \t]*\r?\n?(.*?)(?:```|\z)").expect("valid regex")
49});
50
51/// What [`LLMMathChain::solve`] returns.
52#[derive(Clone, Debug, PartialEq)]
53pub struct MathAnswer {
54    /// The formatted result (`"2518731"`, `"0.333333333333"`), or the model's own answer
55    /// when it replied with `Answer:` instead of an expression.
56    pub answer: String,
57    /// The expression the model wrote, if any.
58    pub expression: Option<String>,
59    pub tokens: Option<TokenUsage>,
60}
61
62/// Answers math questions by asking the model for an expression and evaluating it with the
63/// built-in [calculator](crate::tools::evaluate_expression), so the arithmetic is exact
64/// instead of guessed. Nothing the model writes is executed as code.
65///
66/// The model's expression goes in a fenced block (any language tag, so ```` ```python ````
67/// works too). An `Answer: ...` reply is returned as is.
68///
69/// ```rust,ignore
70/// let chain = LLMMathChain::new(OpenAI::default());
71/// let answer = chain
72///     .invoke(prompt_args! { "question" => "What is 43 * 10000 / 7324.3?" })
73///     .await?; // "58.7086820584"
74/// ```
75pub struct LLMMathChain {
76    llm_chain: Box<dyn Chain>,
77    input_key: String,
78    output_key: String,
79}
80
81impl LLMMathChain {
82    /// Uses the built-in prompt, which takes `question`.
83    pub fn new<L: Into<Box<dyn LLM>>>(llm: L) -> Self {
84        let chain = LLMChainBuilder::new()
85            .prompt(PromptTemplate::new(
86                PROMPT.to_string(),
87                vec![LLM_MATH_DEFAULT_INPUT_KEY.to_string()],
88                TemplateFormat::Jinja2,
89            ))
90            .llm(llm)
91            .build()
92            .expect("prompt and llm are set");
93        Self::from_chain(chain)
94    }
95
96    /// Uses your own chain to write the expression. Its output must hold the expression in
97    /// a fenced block, or an `Answer:` line.
98    pub fn from_chain<C: Into<Box<dyn Chain>>>(chain: C) -> Self {
99        Self {
100            llm_chain: chain.into(),
101            input_key: LLM_MATH_DEFAULT_INPUT_KEY.into(),
102            output_key: LLM_MATH_DEFAULT_OUTPUT_KEY.into(),
103        }
104    }
105
106    /// Input key with the question (default `question`). The prompt still gets it as
107    /// `question`.
108    pub fn with_input_key<S: Into<String>>(mut self, key: S) -> Self {
109        self.input_key = key.into();
110        self
111    }
112
113    /// Key of the answer in [`execute`](Chain::execute)'s output (default `answer`).
114    pub fn with_output_key<S: Into<String>>(mut self, key: S) -> Self {
115        self.output_key = key.into();
116        self
117    }
118
119    /// Asks the model for an expression and evaluates it.
120    pub async fn solve(&self, input_variables: PromptArgs) -> Result<MathAnswer, ChainError> {
121        let mut inputs = input_variables;
122        if self.input_key != LLM_MATH_DEFAULT_INPUT_KEY {
123            let question = inputs
124                .remove(&self.input_key)
125                .ok_or_else(|| ChainError::MissingInputVariable(self.input_key.clone()))?;
126            inputs.insert(LLM_MATH_DEFAULT_INPUT_KEY.to_string(), question);
127        }
128        let result = self.llm_chain.call(inputs).await?;
129        let (answer, expression) = parse_math_output(&result.generation)?;
130        Ok(MathAnswer {
131            answer,
132            expression,
133            tokens: result.tokens,
134        })
135    }
136}
137
138/// `(answer, expression)` from the model's output.
139pub(crate) fn parse_math_output(output: &str) -> Result<(String, Option<String>), ChainError> {
140    let output = strip_reasoning(output).trim();
141    if let Some(caps) = FENCED.captures(output) {
142        let expression = caps[1].trim();
143        if !expression.is_empty() {
144            let value = evaluate_expression(expression).map_err(|e| {
145                ChainError::OtherError(format!("llm math: evaluating `{expression}`: {e}"))
146            })?;
147            return Ok((format_number(value), Some(expression.to_string())));
148        }
149    }
150    if let Some(start) = output.rfind("Answer:") {
151        let answer = output[start + "Answer:".len()..].trim();
152        if !answer.is_empty() {
153            return Ok((answer.to_string(), None));
154        }
155    }
156    Err(ChainError::OtherError(format!(
157        "llm math: expected an expression in a ``` block or an `Answer:` line, got: {output:?}"
158    )))
159}
160
161#[async_trait]
162impl Chain for LLMMathChain {
163    async fn call(&self, input_variables: PromptArgs) -> Result<GenerateResult, ChainError> {
164        let answer = self.solve(input_variables).await?;
165        Ok(GenerateResult {
166            generation: answer.answer,
167            tokens: answer.tokens,
168            ..Default::default()
169        })
170    }
171
172    async fn execute(
173        &self,
174        input_variables: PromptArgs,
175    ) -> Result<HashMap<String, Value>, ChainError> {
176        let answer = self.solve(input_variables).await?;
177        let result = GenerateResult {
178            generation: answer.answer.clone(),
179            tokens: answer.tokens,
180            ..Default::default()
181        };
182        let mut output = HashMap::new();
183        output.insert(self.output_key.clone(), json!(answer.answer));
184        output.insert(DEFAULT_RESULT_KEY.to_string(), json!(result));
185        if let Some(expression) = answer.expression {
186            output.insert(LLM_MATH_EXPRESSION_KEY.to_string(), json!(expression));
187        }
188        Ok(output)
189    }
190
191    fn get_input_keys(&self) -> Vec<String> {
192        vec![self.input_key.clone()]
193    }
194
195    fn get_output_keys(&self) -> Vec<String> {
196        vec![
197            self.output_key.clone(),
198            DEFAULT_RESULT_KEY.to_string(),
199            LLM_MATH_EXPRESSION_KEY.to_string(),
200        ]
201    }
202}
203
204#[cfg(test)]
205mod tests;
langchain_rust/chain/llm_math/mod.rs

langchain_rust/chain/llm_math/
mod.rs
