lc_core/structured_output/parser.rs
1// src/core/structured_output/parser.rs
2//! Incremental JSON parser for handling partial/incomplete JSON from streaming LLM output.
3
4use serde_json::Value;
5
6/// Errors produced by `PartialJsonParser`.
7#[derive(Debug, Clone, thiserror::Error)]
8#[non_exhaustive]
9pub enum PartialJsonError {
10 /// The buffer does not yet contain parseable JSON.
11 #[error("Incomplete JSON: {0}")]
12 Incomplete(String),
13
14 /// The accumulated text is not valid JSON even after repair attempts.
15 #[error("Invalid JSON: {0}")]
16 Invalid(String),
17}
18
19/// Incremental JSON parser that can handle partial/incomplete JSON.
20///
21/// Builds up a string token by token and attempts to parse at each step,
22/// returning the best partial result possible. This is designed for streaming
23/// LLM output where JSON arrives in small chunks and may be incomplete until
24/// the stream finishes.
25///
26/// # Strategy
27///
28/// 1. Accumulate tokens into an internal buffer.
29/// 2. On each `push_and_parse`, attempt to parse the buffer as complete JSON.
30/// 3. If that fails, try to repair the partial JSON by closing unclosed
31/// brackets/braces and truncating incomplete string values.
32/// 4. If repair yields valid JSON, return it; otherwise return
33/// `PartialJsonError::Incomplete`.
34///
35/// # Example
36///
37/// ```ignore
38/// let mut parser = PartialJsonParser::new();
39/// // Simulating token-by-token LLM output
40/// let _ = parser.push_and_parse(r#"{"name":"#); // Incomplete
41/// let v = parser.push_and_parse(r#""Alice","age":30}"#); // Ok({"name":"Alice","age":30})
42/// ```
43pub struct PartialJsonParser {
44 buffer: String,
45 depth: usize,
46 in_string: bool,
47 escape_next: bool,
48}
49
50impl PartialJsonParser {
51 /// Create a new, empty parser.
52 pub fn new() -> Self {
53 Self {
54 buffer: String::new(),
55 depth: 0,
56 in_string: false,
57 escape_next: false,
58 }
59 }
60
61 /// Push a new token and attempt to parse the accumulated buffer.
62 ///
63 /// Returns `Ok(value)` if the buffer (after optional repair) yields valid
64 /// JSON, or `Err(PartialJsonError::Incomplete)` if it does not yet form
65 /// any parseable JSON.
66 pub fn push_and_parse(&mut self, token: &str) -> Result<Value, PartialJsonError> {
67 // Update parser state by scanning the new token
68 for ch in token.chars() {
69 if self.escape_next {
70 self.escape_next = false;
71 continue;
72 }
73 if ch == '\\' && self.in_string {
74 self.escape_next = true;
75 continue;
76 }
77 if ch == '"' {
78 self.in_string = !self.in_string;
79 continue;
80 }
81 if !self.in_string {
82 match ch {
83 '{' | '[' => self.depth += 1,
84 '}' | ']' if self.depth > 0 => self.depth -= 1,
85 _ => {}
86 }
87 }
88 }
89
90 // Ensure we only push at character boundaries (M37: UTF-8 boundary check)
91 if token.is_char_boundary(0) {
92 self.buffer.push_str(token);
93 } else {
94 // Find the first valid char boundary
95 let mut pos = 0;
96 while pos < token.len() && !token.is_char_boundary(pos) {
97 pos += 1;
98 }
99 self.buffer.push_str(&token[pos..]);
100 }
101
102 // H4: 模型常把 JSON 包在 ```json ... ``` 围栏里(或先输出"结果是:"等
103 // 前导文本)。解析时先剥掉围栏/前导文本,只看真正的 JSON 值;否则带
104 // 围栏的合法 JSON 会被判为无法解析,流式结构化输出整条路径不可用。
105 let json = Self::strip_markdown_fence(&self.buffer);
106
107 // Fast path: try full parse first
108 if let Ok(value) = serde_json::from_str::<Value>(json) {
109 return Ok(value);
110 }
111
112 // Only attempt repair if we have at least opened a structure
113 let trimmed = json.trim();
114 if self.depth > 0 || trimmed.starts_with('{') || trimmed.starts_with('[') {
115 let repaired = Self::repair_partial_json(json);
116 if let Ok(value) = serde_json::from_str::<Value>(&repaired) {
117 return Ok(value);
118 }
119 }
120
121 Err(PartialJsonError::Incomplete(format!(
122 "Buffer has {} chars, depth={}",
123 self.buffer.len(),
124 self.depth
125 )))
126 }
127
128 /// Get the final complete value.
129 ///
130 /// Call this when the stream has ended. It first tries to parse the
131 /// fence-stripped buffer, then falls back to the repaired version.
132 pub fn finalize(self) -> Result<Value, PartialJsonError> {
133 // H4: 同样先剥 ```json 围栏,否则带围栏的完整 JSON 会被误判为 Invalid。
134 let json = Self::strip_markdown_fence(&self.buffer);
135
136 // Try full parse
137 if let Ok(value) = serde_json::from_str::<Value>(json) {
138 return Ok(value);
139 }
140
141 // Try repaired
142 let repaired = Self::repair_partial_json(json);
143 serde_json::from_str::<Value>(&repaired).map_err(|e| {
144 PartialJsonError::Invalid(format!(
145 "Failed to parse final buffer ({} chars): {}. Buffer: {}",
146 self.buffer.len(),
147 e,
148 &self.buffer[..std::cmp::min(200, self.buffer.len())]
149 ))
150 })
151 }
152
153 /// Return a reference to the current buffer contents.
154 pub fn buffer(&self) -> &str {
155 &self.buffer
156 }
157
158 /// Whether the parser is currently inside a JSON string.
159 pub fn is_in_string(&self) -> bool {
160 self.in_string
161 }
162
163 /// Current nesting depth of brackets/braces.
164 pub fn depth(&self) -> usize {
165 self.depth
166 }
167
168 /// Strip a markdown code fence and any leading/trailing non-JSON text,
169 /// returning the slice that holds the top-level JSON value.
170 ///
171 /// Streaming-safe: only reads what has been accumulated so far, so a
172 /// partially-delivered object still yields its partial JSON (e.g. while the
173 /// model is still emitting the closing brace).
174 ///
175 /// # Rules
176 ///
177 /// - Leading text up to the first `{` or `[` is dropped (covers a ```json
178 /// fence line, "结果是:" prose, and whitespace).
179 /// - Trailing text after the top-level structure closes is dropped (covers
180 /// the closing ``` fence).
181 /// - Returns `""` when no `{`/`[` has been seen yet (e.g. the buffer is
182 /// still just "```json").
183 pub(crate) fn strip_markdown_fence(buffer: &str) -> &str {
184 let bytes = buffer.as_bytes();
185 // First byte that opens the top-level JSON value. `{`/`[` are ASCII, so
186 // this byte index is always a UTF-8 char boundary.
187 let start = match bytes.iter().position(|b| *b == b'{' || *b == b'[') {
188 Some(i) => i,
189 None => return "",
190 };
191
192 // Walk from `start` tracking string/escape state; the top-level value
193 // ends where depth returns to 0. Everything after it (the closing ```
194 // fence) is dropped. Multi-byte UTF-8 never matches the structural
195 // ASCII bytes below, so byte-wise scanning is safe.
196 let mut depth: i64 = 0;
197 let mut in_string = false;
198 let mut escape_next = false;
199 let mut end = bytes.len();
200 let mut idx = start;
201 while idx < bytes.len() {
202 let b = bytes[idx];
203 if escape_next {
204 escape_next = false;
205 } else if b == b'\\' && in_string {
206 escape_next = true;
207 } else if b == b'"' {
208 in_string = !in_string;
209 } else if !in_string {
210 match b {
211 b'{' | b'[' => depth += 1,
212 b'}' | b']' => {
213 depth -= 1;
214 if depth == 0 {
215 end = idx + 1;
216 break;
217 }
218 }
219 _ => {}
220 }
221 }
222 idx += 1;
223 }
224 &buffer[start..end]
225 }
226
227 /// Repair a partial JSON string by closing unclosed structures and
228 /// truncating incomplete values.
229 ///
230 /// 0.21.0 BUG-1: closing order is a LIFO stack (last opened, first closed),
231 /// not two independent counters. The old counter-based repair closed all
232 /// braces first, then all brackets — independent of actual nesting — so a
233 /// stream truncated at the end of an array produced illegal JSON
234 /// (`{"a": [1, 2` → `{"a": [1, 2}]` instead of `{"a": [1, 2]}`), failing
235 /// the whole streaming structured-output path in that scenario.
236 pub(crate) fn repair_partial_json(text: &str) -> String {
237 let mut repaired = text.trim().to_string();
238
239 // Scan the text tracking string state to correctly identify structure
240 // characters outside of strings (C20 + C21).
241 let mut in_string = false;
242 let mut escape_next = false;
243 // Expected closers for currently-open structures, in open order.
244 // LIFO order at the end reproduces the true nesting.
245 let mut expected_closers: Vec<u8> = Vec::new();
246 let mut unescaped_quote_count = 0usize;
247
248 for ch in repaired.chars() {
249 if escape_next {
250 escape_next = false;
251 continue;
252 }
253 if ch == '\\' && in_string {
254 escape_next = true;
255 continue;
256 }
257 if ch == '"' {
258 unescaped_quote_count += 1;
259 in_string = !in_string;
260 continue;
261 }
262 if !in_string {
263 match ch {
264 '{' => expected_closers.push(b'}'),
265 '[' => expected_closers.push(b']'),
266 '}' | ']' => {
267 // Ignore stray closers (more closes than opens);
268 // pop is a no-op on an empty stack.
269 expected_closers.pop();
270 }
271 _ => {}
272 }
273 }
274 }
275
276 // If we are in the middle of a string value, close it.
277 // Heuristic: odd number of unescaped quotes means an unclosed string.
278 if unescaped_quote_count % 2 != 0 {
279 repaired.push('"');
280 }
281
282 // Close unclosed structures in reverse open order (LIFO) — before
283 // removing trailing commas, so that commas before the newly-added
284 // closers get removed.
285 for closer in expected_closers.iter().rev() {
286 repaired.push(*closer as char);
287 }
288
289 // Remove trailing commas before closing brackets/braces
290 // (must come after closing braces/brackets so we can detect them)
291 repaired = Self::remove_trailing_commas(&repaired);
292
293 repaired
294 }
295
296 /// Remove trailing commas before closing braces/brackets (invalid in strict JSON).
297 ///
298 /// 0.20.0 K1: the scan must track string state — a comma inside a string
299 /// literal (e.g. `{"a": "text, }"}`) is content, not a trailing comma.
300 /// Mirrors the in-string/escape state machine used by
301 /// [`Self::repair_partial_json`]; without it the old version corrupted
302 /// string values whose text ended with a comma followed by `}`/`]`.
303 pub(crate) fn remove_trailing_commas(s: &str) -> String {
304 let mut result = String::with_capacity(s.len());
305 let chars: Vec<char> = s.chars().collect();
306 let mut i = 0;
307 let mut in_string = false;
308 let mut escape_next = false;
309 while i < chars.len() {
310 let ch = chars[i];
311 if in_string {
312 result.push(ch);
313 if escape_next {
314 escape_next = false;
315 } else if ch == '\\' {
316 escape_next = true;
317 } else if ch == '"' {
318 in_string = false;
319 }
320 i += 1;
321 continue;
322 }
323 if ch == '"' {
324 in_string = true;
325 result.push(ch);
326 i += 1;
327 continue;
328 }
329 if ch == ',' && i + 1 < chars.len() {
330 let next_non_ws = chars[i + 1..].iter().find(|c| !c.is_whitespace());
331 if next_non_ws == Some(&'}') || next_non_ws == Some(&']') {
332 // Skip the trailing comma
333 i += 1;
334 continue;
335 }
336 }
337 result.push(ch);
338 i += 1;
339 }
340 result
341 }
342}
343
344impl Default for PartialJsonParser {
345 fn default() -> Self {
346 Self::new()
347 }
348}