lc_core/structured_output/
extract.rs

1// src/core/structured_output/extract.rs
2//! Structured output extraction from LLM responses.
3
4use async_trait::async_trait;
5use serde::de::DeserializeOwned;
6use serde::Serialize;
7use serde_json::Value;
8
9use crate::language_models::{BaseChatModel, LLMResult};
10use crate::output_parsers::BaseOutputParser;
11use crate::output_parsers::JsonOutputParser;
12use lc_schema::Message;
13
14/// Errors that can occur during structured output extraction.
15#[derive(Debug, Clone, thiserror::Error)]
16#[non_exhaustive]
17pub enum StructuredOutputError {
18    /// The provided JSON schema is invalid or malformed.
19    #[error("Schema error: {0}")]
20    SchemaError(String),
21
22    /// The LLM response could not be parsed as the target type.
23    #[error("Parse error: {0}")]
24    ParseError(String),
25
26    /// The provider does not support the requested structured output method.
27    #[error("Provider unsupported: {0}")]
28    ProviderUnsupported(String),
29
30    /// The LLM call itself failed.
31    #[error("LLM error: {0}")]
32    LLMError(String),
33
34    /// The stream ended before a complete JSON object could be parsed.
35    #[error("Stream incomplete: {0}")]
36    StreamIncomplete(String),
37}
38
39/// Trait that extends `BaseChatModel` with structured output capabilities.
40///
41/// Implementors can override the default prompt-injection strategy with
42/// provider-specific mechanisms (e.g., OpenAI function calling, Ollama JSON mode).
43#[async_trait]
44pub trait StructuredOutputExt: BaseChatModel {
45    /// Call the LLM with a JSON schema and prompt, returning a parsed result of type `T`.
46    ///
47    /// The default implementation uses prompt injection: it embeds the schema
48    /// into the system prompt and parses the JSON response with `JsonOutputParser`.
49    ///
50    /// # Arguments
51    ///
52    /// * `schema` - A JSON Schema (`serde_json::Value`) describing the expected output shape.
53    /// * `prompt` - The user prompt / question to send to the LLM.
54    ///
55    /// # Returns
56    ///
57    /// A `Result<T, StructuredOutputError>` where `T` is the deserialized output.
58    async fn with_structured_output<T: DeserializeOwned + Serialize + Send + Sync + 'static>(
59        &self,
60        schema: Value,
61        prompt: &str,
62    ) -> Result<T, StructuredOutputError> {
63        with_structured_output(self, schema, prompt).await
64    }
65}
66
67/// Blanket implementation: every `BaseChatModel` automatically gets `StructuredOutputExt`.
68impl<M: BaseChatModel> StructuredOutputExt for M {}
69
70/// Standalone function to extract structured output from any `BaseChatModel`.
71///
72/// This is the core implementation that works with any chat model by:
73/// 1. Building a system prompt that includes the JSON schema and format instructions
74/// 2. Calling `llm.chat()` with the combined messages
75/// 3. Parsing the LLM's JSON response into the target type `T`
76///
77/// # Arguments
78///
79/// * `llm` - Any type implementing `BaseChatModel`.
80/// * `schema` - A JSON Schema describing the expected output.
81/// * `prompt` - The user prompt to send to the LLM.
82///
83/// # Returns
84///
85/// A `Result<T, StructuredOutputError>` where `T` is the deserialized output.
86///
87/// # Errors
88///
89/// - `StructuredOutputError::SchemaError` if the schema is not a valid JSON object.
90/// - `StructuredOutputError::LLMError` if the underlying `chat()` call fails.
91/// - `StructuredOutputError::ParseError` if the response cannot be parsed as JSON
92///   or deserialized into type `T`.
93pub async fn with_structured_output<T, M>(
94    llm: &M,
95    schema: Value,
96    prompt: &str,
97) -> Result<T, StructuredOutputError>
98where
99    T: DeserializeOwned + Serialize + Send + Sync + 'static,
100    M: BaseChatModel + ?Sized,
101{
102    // Validate the schema is an object
103    if !schema.is_object() {
104        return Err(StructuredOutputError::SchemaError(format!(
105            "Schema must be a JSON object, got: {}",
106            schema
107        )));
108    }
109
110    // Build the system prompt with schema and format instructions
111    let system_prompt = build_structured_system_prompt(&schema);
112
113    let messages = vec![Message::system(system_prompt), Message::human(prompt)];
114
115    // Call the LLM
116    let result: LLMResult = llm
117        .chat(messages, None)
118        .await
119        .map_err(|e| StructuredOutputError::LLMError(e.to_string()))?;
120
121    // Parse the response
122    parse_structured_response::<T>(&result.content).await
123}
124
125/// Build a system prompt that instructs the LLM to output JSON conforming to the schema.
126pub(crate) fn build_structured_system_prompt(schema: &Value) -> String {
127    let schema_str = serde_json::to_string_pretty(schema).unwrap_or_else(|_| schema.to_string());
128
129    format!(
130        "You are a helpful assistant that responds exclusively in valid JSON format.\n\
131         \n\
132         You must respond with a JSON object that conforms to the following JSON Schema:\n\
133         ```json\n\
134         {schema_str}\n\
135         ```\n\
136         \n\
137         Important rules:\n\
138         1. Respond ONLY with valid JSON. Do not include any explanatory text before or after the JSON.\n\
139         2. The JSON must conform exactly to the schema above.\n\
140         3. All required fields must be present.\n\
141         4. Do not include fields that are not in the schema.\n\
142         5. If you cannot satisfy the schema, respond with the closest valid JSON you can produce."
143    )
144}
145
146/// Parse the LLM response content into the target type `T`.
147///
148/// Uses `JsonOutputParser` to handle markdown code blocks and other common
149/// LLM output formatting, then deserializes into `T`.
150pub(crate) async fn parse_structured_response<
151    T: DeserializeOwned + Serialize + Send + Sync + 'static,
152>(
153    content: &str,
154) -> Result<T, StructuredOutputError> {
155    let parser = JsonOutputParser::new();
156
157    let json_value: Value = parser.parse(content).await.map_err(|e| {
158        StructuredOutputError::ParseError(format!("Failed to parse LLM response as JSON: {}", e))
159    })?;
160
161    serde_json::from_value::<T>(json_value).map_err(|e| {
162        StructuredOutputError::ParseError(format!(
163            "Failed to deserialize JSON into target type: {}. Response was: {}",
164            e,
165            &content[..std::cmp::min(200, content.len())]
166        ))
167    })
168}
lc_core/structured_output/extract.rs

lc_core/structured_output/
extract.rs
