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