langchain_rust/chain/constitutional/
mod.rs

1//! Constitutional chain: critiques and revises a chain's answer against a list of
2//! principles, ported from langchaingo / LangChain.
3
4use std::collections::HashMap;
5
6use serde::{Deserialize, Serialize};
7use serde_json::{json, Value};
8
9use crate::{
10    chain::{Chain, ChainError, LLMChain, DEFAULT_OUTPUT_KEY, DEFAULT_RESULT_KEY},
11    language_models::{llm::LLM, GenerateResult, TokenUsage},
12    output_parsers::strip_reasoning,
13    prompt::{PromptArgs, PromptTemplate, TemplateFormat},
14    prompt_args,
15};
16
17use super::LLMChainBuilder;
18
19mod principles;
20mod prompts;
21
22/// Output key with the chain's first answer, when
23/// [`with_return_intermediate_steps`](ConstitutionalChain::with_return_intermediate_steps) is on.
24pub const CONSTITUTIONAL_INITIAL_OUTPUT_KEY: &str = "initial_output";
25/// Output key with every [`CritiqueRevision`], when
26/// [`with_return_intermediate_steps`](ConstitutionalChain::with_return_intermediate_steps) is on.
27pub const CONSTITUTIONAL_STEPS_KEY: &str = "critiques_and_revisions";
28
29/// A rule the answer is checked against: what to look for in a critique, and how to
30/// revise the answer when the critique finds something.
31#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
32pub struct ConstitutionalPrinciple {
33    pub name: String,
34    pub critique_request: String,
35    pub revision_request: String,
36}
37
38impl ConstitutionalPrinciple {
39    pub fn new<C: Into<String>, R: Into<String>>(critique_request: C, revision_request: R) -> Self {
40        Self {
41            name: "Constitutional Principle".into(),
42            critique_request: critique_request.into(),
43            revision_request: revision_request.into(),
44        }
45    }
46
47    pub fn with_name<S: Into<String>>(mut self, name: S) -> Self {
48        self.name = name.into();
49        self
50    }
51
52    /// A built-in principle by name (`"harmful1"`, `"illegal"`, `"thoughtful"`,
53    /// `"uo-ethics-1"`, ...; see [`builtin_names`](Self::builtin_names)).
54    pub fn get(name: &str) -> Option<Self> {
55        principles::PRINCIPLES
56            .iter()
57            .find(|(n, _, _)| *n == name)
58            .map(|(name, critique, revision)| Self::new(*critique, *revision).with_name(*name))
59    }
60
61    /// Names of the built-in principles, the same set as langchaingo and LangChain.
62    pub fn builtin_names() -> impl Iterator<Item = &'static str> {
63        principles::PRINCIPLES.iter().map(|(name, _, _)| *name)
64    }
65}
66
67/// One principle's step: the critique and, if it asked for changes, the revision.
68#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
69pub struct CritiqueRevision {
70    pub principle: String,
71    pub critique: String,
72    /// `None` when the critique found nothing to change, or the reviser answered
73    /// "No revisions needed".
74    pub revision: Option<String>,
75}
76
77/// What [`ConstitutionalChain::critique_and_revise`] returns.
78#[derive(Clone, Debug, PartialEq)]
79pub struct ConstitutionalOutput {
80    /// The wrapped chain's answer, before any revision.
81    pub initial_output: String,
82    /// The answer after every principle.
83    pub output: String,
84    pub critiques_and_revisions: Vec<CritiqueRevision>,
85    /// Usage of all LLM calls.
86    pub tokens: Option<TokenUsage>,
87}
88
89/// Runs a chain, then for each principle asks the model to critique the answer and, when
90/// the critique finds a problem, to revise it. Each revision is what the next principle
91/// critiques.
92///
93/// ```rust,ignore
94/// let chain = ConstitutionalChain::new(
95///     qa_chain,
96///     llm,
97///     vec![ConstitutionalPrinciple::get("illegal").unwrap()],
98/// );
99/// let answer = chain.invoke(prompt_args! { "question" => "..." }).await?;
100/// ```
101pub struct ConstitutionalChain {
102    chain: LLMChain,
103    critique_chain: Box<dyn Chain>,
104    revision_chain: Box<dyn Chain>,
105    principles: Vec<ConstitutionalPrinciple>,
106    output_key: String,
107    return_intermediate_steps: bool,
108}
109
110const CRITIQUE_VARIABLES: [&str; 3] = ["input_prompt", "output_from_model", "critique_request"];
111const REVISION_VARIABLES: [&str; 5] = [
112    "input_prompt",
113    "output_from_model",
114    "critique_request",
115    "critique",
116    "revision_request",
117];
118
119impl ConstitutionalChain {
120    /// `chain` answers first; `llm` writes the critiques and revisions with the built-in
121    /// few-shot prompts.
122    pub fn new<L: Into<Box<dyn LLM>>>(
123        chain: LLMChain,
124        llm: L,
125        principles: Vec<ConstitutionalPrinciple>,
126    ) -> Self {
127        let llm = llm.into();
128        Self {
129            chain,
130            critique_chain: Box::new(llm_chain(
131                llm.clone_box(),
132                critique_template(),
133                &CRITIQUE_VARIABLES,
134            )),
135            revision_chain: Box::new(llm_chain(llm, revision_template(), &REVISION_VARIABLES)),
136            principles,
137            output_key: DEFAULT_OUTPUT_KEY.into(),
138            return_intermediate_steps: false,
139        }
140    }
141
142    /// Replaces the critique step. It gets `input_prompt`, `output_from_model` and
143    /// `critique_request`; an output containing "No critique needed" skips the revision.
144    pub fn with_critique_chain<C: Into<Box<dyn Chain>>>(mut self, chain: C) -> Self {
145        self.critique_chain = chain.into();
146        self
147    }
148
149    /// Replaces the revision step. It gets the critique step's inputs plus `critique` and
150    /// `revision_request`.
151    pub fn with_revision_chain<C: Into<Box<dyn Chain>>>(mut self, chain: C) -> Self {
152        self.revision_chain = chain.into();
153        self
154    }
155
156    /// Key of the final answer in [`execute`](Chain::execute)'s output (default `output`).
157    pub fn with_output_key<S: Into<String>>(mut self, key: S) -> Self {
158        self.output_key = key.into();
159        self
160    }
161
162    /// Also return the first answer and every [`CritiqueRevision`] from
163    /// [`execute`](Chain::execute), under [`CONSTITUTIONAL_INITIAL_OUTPUT_KEY`] and
164    /// [`CONSTITUTIONAL_STEPS_KEY`].
165    pub fn with_return_intermediate_steps(mut self, enabled: bool) -> Self {
166        self.return_intermediate_steps = enabled;
167        self
168    }
169
170    /// Runs the chain and every principle, returning all the steps.
171    pub async fn critique_and_revise(
172        &self,
173        input_variables: PromptArgs,
174    ) -> Result<ConstitutionalOutput, ChainError> {
175        let input_prompt = self.chain.prompt_text(input_variables.clone()).await?;
176        let first = self.chain.call(input_variables).await?;
177        let mut tokens = first.tokens;
178        let initial_output = strip_reasoning(&first.generation).trim().to_string();
179
180        let mut response = initial_output.clone();
181        let mut steps = Vec::with_capacity(self.principles.len());
182        for principle in &self.principles {
183            let critique = self
184                .critique_chain
185                .call(prompt_args! {
186                    "input_prompt" => input_prompt,
187                    "output_from_model" => response,
188                    "critique_request" => principle.critique_request,
189                })
190                .await?;
191            add_usage(&mut tokens, &critique.tokens);
192            let critique = parse_critique(&critique.generation);
193            if critique.to_lowercase().contains("no critique needed") {
194                steps.push(CritiqueRevision {
195                    principle: principle.name.clone(),
196                    critique,
197                    revision: None,
198                });
199                continue;
200            }
201
202            let revision = self
203                .revision_chain
204                .call(prompt_args! {
205                    "input_prompt" => input_prompt,
206                    "output_from_model" => response,
207                    "critique_request" => principle.critique_request,
208                    "critique" => critique,
209                    "revision_request" => principle.revision_request,
210                })
211                .await?;
212            add_usage(&mut tokens, &revision.tokens);
213            let revision = parse_revision(&revision.generation);
214            // The revision prompt tells the model to say so instead of changing nothing.
215            let revision = if revision.to_lowercase().starts_with("no revisions needed") {
216                None
217            } else {
218                response = revision.clone();
219                Some(revision)
220            };
221            steps.push(CritiqueRevision {
222                principle: principle.name.clone(),
223                critique,
224                revision,
225            });
226        }
227
228        Ok(ConstitutionalOutput {
229            initial_output,
230            output: response,
231            critiques_and_revisions: steps,
232            tokens,
233        })
234    }
235}
236
237#[async_trait::async_trait]
238impl Chain for ConstitutionalChain {
239    async fn call(&self, input_variables: PromptArgs) -> Result<GenerateResult, ChainError> {
240        let out = self.critique_and_revise(input_variables).await?;
241        Ok(GenerateResult {
242            generation: out.output,
243            tokens: out.tokens,
244            ..Default::default()
245        })
246    }
247
248    async fn execute(
249        &self,
250        input_variables: PromptArgs,
251    ) -> Result<HashMap<String, Value>, ChainError> {
252        let out = self.critique_and_revise(input_variables).await?;
253        let result = GenerateResult {
254            generation: out.output.clone(),
255            tokens: out.tokens,
256            ..Default::default()
257        };
258        let mut output = HashMap::new();
259        output.insert(self.output_key.clone(), json!(out.output));
260        output.insert(DEFAULT_RESULT_KEY.to_string(), json!(result));
261        if self.return_intermediate_steps {
262            output.insert(
263                CONSTITUTIONAL_INITIAL_OUTPUT_KEY.to_string(),
264                json!(out.initial_output),
265            );
266            output.insert(
267                CONSTITUTIONAL_STEPS_KEY.to_string(),
268                json!(out.critiques_and_revisions),
269            );
270        }
271        Ok(output)
272    }
273
274    fn get_input_keys(&self) -> Vec<String> {
275        self.chain.get_input_keys()
276    }
277
278    fn get_output_keys(&self) -> Vec<String> {
279        let mut keys = vec![self.output_key.clone(), DEFAULT_RESULT_KEY.to_string()];
280        if self.return_intermediate_steps {
281            keys.push(CONSTITUTIONAL_INITIAL_OUTPUT_KEY.to_string());
282            keys.push(CONSTITUTIONAL_STEPS_KEY.to_string());
283        }
284        keys
285    }
286}
287
288fn add_usage(total: &mut Option<TokenUsage>, usage: &Option<TokenUsage>) {
289    if let Some(usage) = usage {
290        total.get_or_insert_with(TokenUsage::default).add(usage);
291    }
292}
293
294/// The critique without reasoning blocks, an echoed `Critique:` label, or a
295/// `Revision request:` the model may have gone on to write.
296fn parse_critique(raw: &str) -> String {
297    let critique = strip_reasoning(raw);
298    // ASCII lowercasing keeps byte offsets valid for slicing `critique`.
299    let lower = critique.to_ascii_lowercase();
300    let end = lower.find("revision request:").unwrap_or(critique.len());
301    strip_label(&critique[..end], "critique:").to_string()
302}
303
304/// The revision without reasoning blocks or an echoed `Revision:` label.
305fn parse_revision(raw: &str) -> String {
306    strip_label(strip_reasoning(raw), "revision:").to_string()
307}
308
309/// `text` trimmed, without a leading `label` (any case), as models often repeat the
310/// prompt's last line.
311fn strip_label<'a>(text: &'a str, label: &str) -> &'a str {
312    let text = text.trim();
313    match text.get(..label.len()) {
314        Some(start) if start.eq_ignore_ascii_case(label) => text[label.len()..].trim(),
315        _ => text,
316    }
317}
318
319fn llm_chain(llm: Box<dyn LLM>, template: String, variables: &[&str]) -> LLMChain {
320    LLMChainBuilder::new()
321        .prompt(PromptTemplate::new(
322            template,
323            variables.iter().map(|v| v.to_string()).collect(),
324            TemplateFormat::Jinja2,
325        ))
326        .llm(llm)
327        .build()
328        .expect("prompt and llm are set")
329}
330
331const EXAMPLE_SEPARATOR: &str = "\n === \n";
332
333fn critique_example(e: &prompts::Example) -> String {
334    format!(
335        "Human: {}\n\nModel: {}\n\nCritique Request: {}\n\nCritique: {}",
336        e.input_prompt, e.output_from_model, e.critique_request, e.critique
337    )
338}
339
340pub(super) fn critique_template() -> String {
341    let mut parts = vec!["Below is a conversation between a human and an AI model. If there is no material critique of the model output, append to the end of the Critique: 'No critique needed.' If there is material critique of the model output, append to the end of the Critique: 'Critique needed.'".to_string()];
342    parts.extend(prompts::EXAMPLES.iter().map(critique_example));
343    parts.push(
344        "Human: {{input_prompt}}\n\nModel: {{output_from_model}}\n\nCritique Request: {{critique_request}}\n\nCritique:"
345            .to_string(),
346    );
347    parts.join(EXAMPLE_SEPARATOR)
348}
349
350pub(super) fn revision_template() -> String {
351    let mut parts = vec!["Below is a conversation between a human and an AI model.".to_string()];
352    parts.extend(prompts::EXAMPLES.iter().map(|e| {
353        format!(
354            "{}\n\nRevision Request: {}\n\nRevision: {}",
355            critique_example(e),
356            e.revision_request,
357            e.revision
358        )
359    }));
360    parts.push(
361        "Human: {{input_prompt}}\n\nModel: {{output_from_model}}\n\nCritique Request: {{critique_request}}\n\nCritique: {{critique}}\n\nIf the critique does not identify anything worth changing, ignore the Revision Request and do not make any revisions. Instead, return \"No revisions needed\".\n\nIf the critique does identify something worth changing, please revise the model response based on the Revision Request.\n\nRevision Request: {{revision_request}}\n\nRevision:"
362            .to_string(),
363    );
364    parts.join(EXAMPLE_SEPARATOR)
365}
366
367#[cfg(test)]
368mod tests;
langchain_rust/chain/constitutional/mod.rs

langchain_rust/chain/constitutional/
mod.rs
