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

lc_core/structured_output/
extract.rs
