Skip to main content

langchain_rust/output_parsers/
structured_parser.rs

1use serde_json::{Map, Value};
2
3use super::{
4    typed::{extract_json, preview, string_output_parser},
5    OutputParserError, TypedOutputParser,
6};
7
8/// One field the model has to return, for [`StructuredOutputParser`].
9#[derive(Debug, Clone, PartialEq)]
10pub struct ResponseSchema {
11    pub name: String,
12    pub description: String,
13}
14
15impl ResponseSchema {
16    pub fn new<N: Into<String>, D: Into<String>>(name: N, description: D) -> Self {
17        Self {
18            name: name.into(),
19            description: description.into(),
20        }
21    }
22}
23
24/// Asks the model for a JSON object with the given fields and checks they're all there.
25///
26/// ```rust,ignore
27/// let parser = StructuredOutputParser::new([
28///     ResponseSchema::new("answer", "answer to the user's question"),
29///     ResponseSchema::new("source", "website used to answer"),
30/// ]);
31/// let fields = parser.parse_output(&generation)?;
32/// ```
33///
34/// The JSON may come in a fenced block or bare. Field values are kept as JSON (strings,
35/// numbers, lists...); extra fields are kept too. As a chain output parser it yields the
36/// object as JSON.
37#[derive(Debug, Clone)]
38pub struct StructuredOutputParser {
39    schemas: Vec<ResponseSchema>,
40}
41
42impl StructuredOutputParser {
43    pub fn new<I: IntoIterator<Item = ResponseSchema>>(schemas: I) -> Self {
44        Self {
45            schemas: schemas.into_iter().collect(),
46        }
47    }
48}
49
50impl TypedOutputParser for StructuredOutputParser {
51    type Output = Map<String, Value>;
52
53    fn format_instructions(&self) -> String {
54        let fields: String = self
55            .schemas
56            .iter()
57            .map(|s| format!("\t\"{}\": string // {}\n", s.name, s.description))
58            .collect();
59        format!(
60            "The output should be a markdown code snippet formatted in the following schema, \
61             including the leading and trailing \"```json\" and \"```\":\n\n```json\n{{\n{fields}}}\n```"
62        )
63    }
64
65    fn parse_output(&self, text: &str) -> Result<Self::Output, OutputParserError> {
66        let Value::Object(object) = extract_json(text)? else {
67            return Err(OutputParserError::ParsingError(format!(
68                "expected a JSON object: {}",
69                preview(text)
70            )));
71        };
72        let missing: Vec<_> = self
73            .schemas
74            .iter()
75            .filter(|s| !object.contains_key(&s.name))
76            .map(|s| s.name.as_str())
77            .collect();
78        if !missing.is_empty() {
79            return Err(OutputParserError::ParsingError(format!(
80                "output is missing the fields {missing:?}: {}",
81                preview(text)
82            )));
83        }
84        Ok(object)
85    }
86}
87
88string_output_parser!(StructuredOutputParser, |object| Value::Object(object)
89    .to_string());