ferrin_core/middleware/builtin/
extract_json.rs

1//! JSON extraction from fenced text.
2
3use std::collections::HashMap;
4use std::fmt;
5use std::sync::Arc;
6
7use ferrin_spec::BoxFuture;
8use ferrin_spec::CallOptions;
9use ferrin_spec::Content;
10use ferrin_spec::PartId;
11use ferrin_spec::StreamPart;
12use ferrin_spec::StreamResult;
13use ferrin_spec::error::ProviderError;
14use ferrin_spec::language_model::GenerateResult;
15use futures_util::StreamExt;
16use futures_util::stream;
17
18use crate::middleware::GenerateNext;
19use crate::middleware::LanguageModelMiddleware;
20use crate::middleware::MiddlewareContext;
21use crate::middleware::StreamNext;
22
23/// Rewrites the complete text of a text part.
24pub type JsonTransformFn = Arc<dyn Fn(&str) -> String + Send + Sync>;
25
26/// Middleware created by [`extract_json`].
27#[derive(Clone)]
28pub struct ExtractJson {
29    transform: Option<JsonTransformFn>,
30}
31
32impl fmt::Debug for ExtractJson {
33    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
34        f.debug_struct("ExtractJson")
35            .field("custom_transform", &self.transform.is_some())
36            .finish()
37    }
38}
39
40/// Strips a markdown code fence (```` ```json ```` .. ```` ``` ````) around
41/// text parts and trims the result.
42///
43/// Streams are rewritten incrementally: the opening fence is removed once a
44/// full line is available, the last twelve characters are held back until
45/// `text-end` so the closing fence can be removed, and `text-start` is
46/// delayed until the first delta is forwarded. A custom transform buffers
47/// the whole part and applies at `text-end`.
48#[must_use]
49pub fn extract_json() -> ExtractJson {
50    ExtractJson { transform: None }
51}
52
53/// Characters held back while streaming so a trailing fence can be removed.
54const SUFFIX_BUFFER_CHARS: usize = 12;
55
56impl ExtractJson {
57    /// Replaces the default fence stripping with `transform`.
58    #[must_use]
59    pub fn transform(mut self, transform: impl Fn(&str) -> String + Send + Sync + 'static) -> Self {
60        self.transform = Some(Arc::new(transform));
61        self
62    }
63
64    fn run_transform(&self, text: &str) -> String {
65        match &self.transform {
66            Some(transform) => transform(text),
67            None => strip_json_fences(text),
68        }
69    }
70
71    /// Applies the extraction to a complete result.
72    #[must_use]
73    pub fn apply(&self, mut result: GenerateResult) -> GenerateResult {
74        for part in &mut result.content {
75            if let Content::Text { text, .. } = part {
76                *text = self.run_transform(text);
77            }
78        }
79        result
80    }
81
82    /// Applies the extraction to a stream result.
83    #[must_use]
84    pub fn apply_stream(&self, result: StreamResult) -> StreamResult {
85        let StreamResult {
86            stream,
87            request,
88            response,
89        } = result;
90        let mut state = StreamState {
91            transform: self.transform.clone(),
92            blocks: HashMap::new(),
93        };
94        let stream = stream
95            .map(move |part| stream::iter(state.process(part)))
96            .flatten();
97        StreamResult {
98            stream: Box::pin(stream),
99            request,
100            response,
101        }
102    }
103}
104
105impl LanguageModelMiddleware for ExtractJson {
106    fn wrap_generate<'a>(
107        &'a self,
108        options: CallOptions,
109        next: GenerateNext<'a>,
110        _ctx: MiddlewareContext<'a>,
111    ) -> BoxFuture<'a, Result<GenerateResult, ProviderError>> {
112        Box::pin(async move { Ok(self.apply(next(options).await?)) })
113    }
114
115    fn wrap_stream<'a>(
116        &'a self,
117        options: CallOptions,
118        next: StreamNext<'a>,
119        _ctx: MiddlewareContext<'a>,
120    ) -> BoxFuture<'a, Result<StreamResult, ProviderError>> {
121        Box::pin(async move { Ok(self.apply_stream(next(options).await?)) })
122    }
123}
124
125/// Default rewrite: strip the opening and closing fences, then trim.
126#[must_use]
127pub fn strip_json_fences(text: &str) -> String {
128    strip_fence_suffix(strip_fence_prefix(text))
129        .trim_matches(is_json_fence_whitespace)
130        .to_owned()
131}
132
133/// Removes ```` ``` ```` or ```` ```json ```` and the following whitespace.
134fn strip_fence_prefix(text: &str) -> &str {
135    match text.strip_prefix("```") {
136        Some(rest) => rest
137            .strip_prefix("json")
138            .unwrap_or(rest)
139            .trim_start_matches(is_json_fence_whitespace),
140        None => text,
141    }
142}
143
144/// Removes a trailing ```` ``` ```` (with trailing whitespace and one
145/// preceding newline).
146fn strip_fence_suffix(text: &str) -> &str {
147    let trimmed = text.trim_end_matches(is_json_fence_whitespace);
148    match trimmed.strip_suffix("```") {
149        Some(rest) => rest.strip_suffix('\n').unwrap_or(rest),
150        None => text,
151    }
152}
153
154fn strip_markdown_code_fence_suffix(text: &str) -> String {
155    strip_fence_suffix(text)
156        .trim_end_matches(is_json_fence_whitespace)
157        .to_owned()
158}
159
160fn is_json_fence_whitespace(ch: char) -> bool {
161    matches!(ch, '\u{0009}'..='\u{000d}' | ' ' | '\u{00a0}' | '\u{1680}'
162        | '\u{2000}'..='\u{200a}' | '\u{2028}' | '\u{2029}' | '\u{202f}'
163        | '\u{205f}' | '\u{3000}' | '\u{feff}')
164}
165
166/// Length of a complete opening fence line (```` ```json ```` plus
167/// whitespace ending in a newline), when `text` starts with one.
168fn fence_prefix_len(text: &str) -> Option<usize> {
169    let rest = text.strip_prefix("```")?;
170    let rest = rest.strip_prefix("json").unwrap_or(rest);
171    let consumed = text.len() - rest.len();
172    let whitespace_len = rest
173        .char_indices()
174        .find(|(_, c)| !is_json_fence_whitespace(*c))
175        .map_or(rest.len(), |(index, _)| index);
176    let whitespace = &rest[..whitespace_len];
177    let last_newline = whitespace.rfind('\n')?;
178    Some(consumed + last_newline + 1)
179}
180
181#[derive(Debug, Clone, Copy, PartialEq, Eq)]
182enum Phase {
183    Prefix,
184    Streaming,
185    Buffering,
186}
187
188struct Block {
189    start: StreamPart,
190    phase: Phase,
191    buffer: String,
192    prefix_stripped: bool,
193}
194
195struct StreamState {
196    transform: Option<JsonTransformFn>,
197    blocks: HashMap<PartId, Block>,
198}
199
200impl StreamState {
201    fn process(&mut self, part: StreamPart) -> Vec<StreamPart> {
202        let mut out = Vec::new();
203        match part {
204            StreamPart::TextStart { ref id, .. } => {
205                let phase = if self.transform.is_some() {
206                    Phase::Buffering
207                } else {
208                    Phase::Prefix
209                };
210                self.blocks.insert(
211                    id.clone(),
212                    Block {
213                        start: part,
214                        phase,
215                        buffer: String::new(),
216                        prefix_stripped: false,
217                    },
218                );
219            }
220            StreamPart::TextDelta {
221                id,
222                delta,
223                provider_metadata,
224            } => {
225                let Some(block) = self.blocks.get_mut(&id) else {
226                    out.push(StreamPart::TextDelta {
227                        id,
228                        delta,
229                        provider_metadata,
230                    });
231                    return out;
232                };
233                block.buffer.push_str(&delta);
234                if block.phase == Phase::Buffering {
235                    return out;
236                }
237                if block.phase == Phase::Prefix {
238                    if !block.buffer.is_empty() && !block.buffer.starts_with('`') {
239                        block.phase = Phase::Streaming;
240                        out.push(block.start.clone());
241                    } else if block.buffer.starts_with("```") {
242                        // Strip the fence only once the line is complete.
243                        if block.buffer.contains('\n') {
244                            if let Some(len) = fence_prefix_len(&block.buffer) {
245                                block.buffer = block.buffer[len..].to_owned();
246                                block.prefix_stripped = true;
247                            }
248                            block.phase = Phase::Streaming;
249                            out.push(block.start.clone());
250                        }
251                    } else if block.buffer.chars().count() >= 3 {
252                        block.phase = Phase::Streaming;
253                        out.push(block.start.clone());
254                    }
255                }
256                if block.phase == Phase::Streaming {
257                    let count = block.buffer.chars().count();
258                    if count > SUFFIX_BUFFER_CHARS {
259                        let split = block
260                            .buffer
261                            .char_indices()
262                            .nth(count - SUFFIX_BUFFER_CHARS)
263                            .map_or(block.buffer.len(), |(index, _)| index);
264                        let to_stream = block.buffer[..split].to_owned();
265                        block.buffer = block.buffer[split..].to_owned();
266                        out.push(StreamPart::TextDelta {
267                            id,
268                            delta: to_stream,
269                            provider_metadata: None,
270                        });
271                    }
272                }
273            }
274            StreamPart::TextEnd { ref id, .. } => {
275                if let Some(block) = self.blocks.remove(id) {
276                    if matches!(block.phase, Phase::Prefix | Phase::Buffering) {
277                        out.push(block.start);
278                    }
279                    let remaining = match block.phase {
280                        Phase::Buffering => match &self.transform {
281                            Some(transform) => transform(&block.buffer),
282                            None => strip_json_fences(&block.buffer),
283                        },
284                        // Nothing streamed yet: the full transform is safe.
285                        Phase::Prefix => strip_json_fences(&block.buffer),
286                        // Earlier text already streamed: only strip the
287                        // suffix so leading whitespace at the boundary stays.
288                        Phase::Streaming => strip_markdown_code_fence_suffix(&block.buffer),
289                    };
290                    if !remaining.is_empty() {
291                        out.push(StreamPart::TextDelta {
292                            id: id.clone(),
293                            delta: remaining,
294                            provider_metadata: None,
295                        });
296                    }
297                }
298                out.push(part);
299            }
300            other => out.push(other),
301        }
302        out
303    }
304}
ferrin_core/middleware/builtin/extract_json.rs

ferrin_core/middleware/builtin/
extract_json.rs
