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

lc_core/structured_output/
extract.rs
