pagesift_api/jina/
respond.rs

1//! 按 Jina 的输出参数写响应正文：`x-respond-with` 的各种形状、图片/链接/媒体的保留方式、
2//! 链接和图片汇总、`x-max-tokens` 截断（报 `output_truncated`）、`x-token-budget` 超限拒绝。
3//! token 数和分块模块用同一个 `TokenIndex`。
4
5use regex::Regex;
6use std::sync::LazyLock;
7
8use super::headers::{JinaParams, RespondWith, Retain};
9use crate::error::{ApiError, ErrorCode};
10use crate::types::Document;
11
12static IMAGE: LazyLock<Regex> =
13    LazyLock::new(|| Regex::new(r#"!\[([^\]]*)\]\(([^)\s]+)(?:\s+"[^"]*")?\)"#).expect("regex"));
14static LINK: LazyLock<Regex> = LazyLock::new(|| {
15    Regex::new(r#"(^|[^!])\[([^\]]*)\]\(([^)\s]+)(?:\s+"[^"]*")?\)"#).expect("regex")
16});
17
18/// 视频、音频：文件扩展名，或常见视频站的观看/嵌入地址。
19fn is_media(url: &str) -> bool {
20    let path = url
21        .split(['?', '#'])
22        .next()
23        .unwrap_or(url)
24        .to_ascii_lowercase();
25    const EXT: &[&str] = &[
26        ".mp4", ".webm", ".mov", ".m4v", ".ogv", ".mp3", ".m4a", ".ogg", ".oga", ".wav", ".flac",
27        ".m3u8", ".mpd",
28    ];
29    const HOSTS: &[&str] = &[
30        "youtube.com/watch",
31        "youtube.com/embed/",
32        "youtu.be/",
33        "vimeo.com/",
34        "player.vimeo.com/",
35        "bilibili.com/video/",
36        "player.bilibili.com/",
37    ];
38    EXT.iter().any(|e| path.ends_with(e)) || HOSTS.iter().any(|h| path.contains(h))
39}
40
41/// 汇总里的一项：文字和地址。
42type Ref = (String, String);
43
44/// 应用保留方式，并收集汇总用的链接和图片（去掉之前收集）。
45pub fn reshape(md: &str, p: &JinaParams) -> (String, Vec<Ref>, Vec<Ref>) {
46    let images: Vec<Ref> = IMAGE
47        .captures_iter(md)
48        .filter(|c| !is_media(&c[2]))
49        .map(|c| (c[1].to_string(), c[2].to_string()))
50        .collect();
51    let links: Vec<Ref> = LINK
52        .captures_iter(md)
53        .filter(|c| !is_media(&c[3]))
54        .map(|c| (c[2].to_string(), c[3].to_string()))
55        .collect();
56    let out = IMAGE.replace_all(md, |c: &regex::Captures<'_>| {
57        let (alt, url) = (&c[1], &c[2]);
58        let retain = if is_media(url) {
59            p.retain_media
60        } else {
61            p.retain_images
62        };
63        match retain {
64            Retain::All | Retain::Image => c[0].to_string(),
65            Retain::None => String::new(),
66            Retain::Text => alt.to_string(),
67        }
68    });
69    let out = LINK.replace_all(&out, |c: &regex::Captures<'_>| {
70        let (lead, text, url) = (&c[1], &c[2], &c[3]);
71        if is_media(url) {
72            return match p.retain_media {
73                Retain::All => c[0].to_string(),
74                Retain::None => lead.to_string(),
75                Retain::Text => format!("{lead}{text}"),
76                Retain::Image => format!("{lead}![{text}]({url})"),
77            };
78        }
79        match p.retain_links {
80            Retain::All | Retain::Image => c[0].to_string(),
81            Retain::None => lead.to_string(),
82            Retain::Text => format!("{lead}{text}"),
83        }
84    });
85    (out.into_owned(), links, images)
86}
87
88fn summary(title: &str, refs: &[Ref], all: bool, image: bool) -> String {
89    let mut seen = std::collections::HashSet::new();
90    let mut s = format!("\n\n{title}:\n");
91    for (text, url) in refs {
92        if !all && !seen.insert(url.clone()) {
93            continue;
94        }
95        if image {
96            s.push_str(&format!("- ![{text}]({url})\n"));
97        } else {
98            s.push_str(&format!("- [{text}]({url})\n"));
99        }
100    }
101    s
102}
103
104/// 粗略的 Markdown → 纯文本：去掉标题井号、强调符号、链接和图片的语法（保留文字）。
105pub fn markdown_text(md: &str) -> String {
106    let no_img = IMAGE.replace_all(md, "$1");
107    let no_link = LINK.replace_all(&no_img, "$1$2");
108    no_link
109        .lines()
110        .map(|l| {
111            l.trim_start_matches('#')
112                .trim_start()
113                .replace("**", "")
114                .replace("__", "")
115        })
116        .collect::<Vec<_>>()
117        .join("\n")
118}
119
120pub fn token_count(s: &str) -> usize {
121    pagesift_extract::TokenIndex::new(s).len()
122}
123
124/// 截到 `max` 个 token（在 token 边界）；返回截断后的文本和是否截断了。
125pub fn truncate_tokens(s: &str, max: usize) -> (String, bool) {
126    let idx = pagesift_extract::TokenIndex::new(s);
127    match idx.starts.get(max) {
128        Some(&cut) => (s[..cut].trim_end().to_string(), true),
129        None => (s.to_string(), false),
130    }
131}
132
133fn yaml_str(v: &str) -> String {
134    serde_json::to_string(v).unwrap_or_else(|_| "\"\"".into())
135}
136
137/// 响应：内容类型、正文、警告。
138#[derive(Debug)]
139pub struct Rendered {
140    pub content_type: &'static str,
141    pub body: String,
142    /// 纯文本响应：不分块时同 `body`；分块时是 front matter、用分隔线连起来的各块、汇总。
143    pub text: String,
144    pub warnings: Vec<String>,
145    /// 分块时的各块（`Accept: application/json` 时返回数组）。
146    pub chunks: Option<Vec<String>>,
147}
148
149pub fn render(doc: &Document, url: &str, p: &JinaParams) -> Result<Rendered, ApiError> {
150    if let Some(want) = p.assert_status
151        && doc.metadata.status_code != want
152    {
153        return Err(ApiError::new(
154            ErrorCode::AssertionFailed,
155            format!(
156                "the page answered {}, not the asserted {want}",
157                doc.metadata.status_code
158            ),
159        ));
160    }
161    let title = doc.metadata.title.as_deref().unwrap_or_default();
162    let mut warnings = p.notes.clone();
163    // 分块时纯文本响应的开头和结尾（front matter、汇总）：块替换的只是正文
164    let (mut head, mut tail) = (String::new(), String::new());
165    let (content_type, body) = match p.respond_with {
166        RespondWith::Html => (
167            "text/html; charset=utf-8",
168            doc.raw_html.clone().unwrap_or_default(),
169        ),
170        RespondWith::Screenshot { .. } => {
171            let shot = doc
172                .actions
173                .as_ref()
174                .and_then(|a| a["screenshots"].get(0))
175                .and_then(|v| v.as_str())
176                .ok_or_else(|| {
177                    ApiError::unsupported(
178                        "screenshot needs browser rendering, which is not available here",
179                    )
180                })?;
181            ("text/plain; charset=utf-8", shot.to_string())
182        }
183        _ => {
184            let md = doc.markdown.clone().unwrap_or_default();
185            let (mut md, links, images) = reshape(&md, p);
186            if let Some(all) = p.images_summary {
187                tail.push_str(&summary("Images", &images, all, true));
188            }
189            if let Some(all) = p.links_summary {
190                tail.push_str(&summary("Links/Buttons", &links, all, false));
191            }
192            md.push_str(&tail);
193            match p.respond_with {
194                RespondWith::Text => ("text/plain; charset=utf-8", markdown_text(&md)),
195                RespondWith::Frontmatter | RespondWith::MarkdownFrontmatter => {
196                    let desc = doc.metadata.description.as_deref().unwrap_or_default();
197                    head = format!(
198                        "---\ntitle: {}\ndescription: {}\nurl: {}\n---\n\n",
199                        yaml_str(title),
200                        yaml_str(desc),
201                        yaml_str(url)
202                    );
203                    ("text/markdown; charset=utf-8", format!("{head}{md}"))
204                }
205                _ => ("text/markdown; charset=utf-8", md),
206            }
207        }
208    };
209    let mut body = body;
210    if let Some(budget) = p.token_budget {
211        let n = token_count(&body);
212        if n > budget {
213            return Err(ApiError::new(
214                ErrorCode::TooLarge,
215                format!("the content has {n} tokens, over x-token-budget {budget}"),
216            )
217            .with_details(serde_json::json!({ "tokens": n, "budget": budget })));
218        }
219    }
220    // 分块来自原 Markdown：每块同样应用保留方式；纯文本时块里也要是纯文本
221    let mut chunks = p.chunking.as_ref().and(doc.chunks.as_ref()).map(|c| {
222        c.iter()
223            .map(|ch| {
224                let md = reshape(&ch.text, p).0;
225                match p.respond_with {
226                    RespondWith::Text => markdown_text(&md),
227                    _ => md,
228                }
229            })
230            .collect::<Vec<_>>()
231    });
232    if let Some(max) = p.max_tokens {
233        let (cut, truncated) = truncate_tokens(&body, max);
234        if truncated {
235            warnings.push(format!(
236                "output_truncated: {} of {} tokens",
237                max,
238                token_count(&body)
239            ));
240            body = cut;
241        }
242        // 块和正文共用同一个上限：超出的块丢掉，跨线的那块截断
243        if let Some(c) = &mut chunks {
244            truncate_chunks(c, max);
245        }
246    }
247    let text = match &chunks {
248        Some(c) => {
249            let (joined, _) = truncate_tokens(
250                &format!("{head}{}{tail}", c.join(CHUNK_SEPARATOR)),
251                p.max_tokens.unwrap_or(usize::MAX),
252            );
253            joined
254        }
255        None => body.clone(),
256    };
257    Ok(Rendered {
258        content_type,
259        body,
260        text,
261        warnings,
262        chunks,
263    })
264}
265
266/// 纯文本分块响应里块与块之间的分隔：一个空行加 `---` 再一个空行（Markdown 的分隔线，不和正文混淆）。
267const CHUNK_SEPARATOR: &str = "\n\n---\n\n";
268
269/// 让所有块加起来不超过 `max` 个 token。
270fn truncate_chunks(chunks: &mut Vec<String>, max: usize) {
271    let mut left = max;
272    let mut keep = 0;
273    for c in chunks.iter_mut() {
274        if left == 0 {
275            break;
276        }
277        let n = token_count(c);
278        if n > left {
279            *c = truncate_tokens(c, left).0;
280            keep += 1;
281            break;
282        }
283        left -= n;
284        keep += 1;
285    }
286    chunks.truncate(keep);
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292    use crate::jina::headers::parse;
293    use axum::http::HeaderMap;
294
295    fn params(pairs: &[(&str, &str)]) -> JinaParams {
296        let mut h = HeaderMap::new();
297        for (k, v) in pairs {
298            h.insert(
299                axum::http::HeaderName::from_bytes(k.as_bytes()).unwrap(),
300                v.parse().unwrap(),
301            );
302        }
303        parse(&h).unwrap().1
304    }
305
306    const MD: &str = "# T\n\nSee [docs](https://a.test/docs) and [docs again](https://a.test/docs).\n\n![logo](https://a.test/l.png) [clip](https://youtu.be/x1)";
307
308    #[test]
309    fn retain_modes_and_summaries() {
310        let (out, ..) = reshape(
311            MD,
312            &params(&[
313                ("x-retain-links", "text"),
314                ("x-retain-images", "none"),
315                ("x-retain-media", "none"),
316            ]),
317        );
318        assert_eq!(out, "# T\n\nSee docs and docs again.\n\n ");
319        let p = params(&[("x-retain-images", "alt"), ("x-with-links-summary", "true")]);
320        let doc = Document {
321            markdown: Some(MD.into()),
322            ..Default::default()
323        };
324        let r = render(&doc, "https://a.test/", &p).unwrap();
325        assert!(
326            r.body.contains("\n\nlogo [clip](https://youtu.be/x1)"),
327            "{}",
328            r.body
329        );
330        assert_eq!(
331            r.body.matches("- [docs").count(),
332            1,
333            "deduplicated: {}",
334            r.body
335        );
336    }
337
338    fn chunk(text: &str) -> crate::types::DocChunk {
339        crate::types::DocChunk {
340            text: text.into(),
341            heading_path: Vec::new(),
342            start: 0,
343            end: 0,
344            tokens: 0,
345        }
346    }
347
348    fn chunked_doc() -> Document {
349        Document {
350            markdown: Some(
351                "## A\n\nalpha words\n\n## B\n\nbeta words [l](https://a.test/x)".into(),
352            ),
353            chunks: Some(vec![
354                chunk("## A\n\nalpha words"),
355                chunk("## B\n\nbeta words [l](https://a.test/x)"),
356            ]),
357            ..Default::default()
358        }
359    }
360
361    #[test]
362    fn chunks_keep_frontmatter_summary_and_limit() {
363        let p = params(&[
364            ("x-respond-with", "frontmatter"),
365            ("x-markdown-chunking", "h2"),
366            ("x-with-links-summary", "true"),
367        ]);
368        let r = render(&chunked_doc(), "https://a.test/", &p).unwrap();
369        assert!(r.text.starts_with("---\ntitle:"), "{}", r.text);
370        assert!(r.text.contains("alpha words\n\n---\n\n## B"), "{}", r.text);
371        assert!(r.text.contains("Links/Buttons"), "{}", r.text);
372        assert_eq!(r.chunks.as_ref().unwrap().len(), 2);
373        let p = params(&[("x-markdown-chunking", "h2"), ("x-max-tokens", "4")]);
374        let r = render(&chunked_doc(), "u", &p).unwrap();
375        assert!(
376            r.warnings[0].starts_with("output_truncated"),
377            "{:?}",
378            r.warnings
379        );
380        let c = r.chunks.unwrap();
381        assert_eq!(c.len(), 1, "{c:?}");
382        assert!(token_count(&r.text) <= 4, "{}", r.text);
383    }
384
385    #[test]
386    fn asserted_status_and_preset_notes() {
387        let mut doc = chunked_doc();
388        doc.metadata.status_code = 404;
389        let e = render(&doc, "u", &params(&[("x-assert-status-code", "200")])).unwrap_err();
390        assert_eq!(e.code, ErrorCode::AssertionFailed);
391        assert!(render(&doc, "u", &params(&[("x-assert-status-code", "404")])).is_ok());
392        let r = render(&doc, "u", &params(&[("x-preset", "reader")])).unwrap();
393        assert!(
394            r.warnings[0].starts_with("preset_partial"),
395            "{:?}",
396            r.warnings
397        );
398    }
399
400    #[test]
401    fn max_tokens_truncates_and_budget_refuses() {
402        let doc = Document {
403            markdown: Some("one two three four five six".into()),
404            ..Default::default()
405        };
406        let r = render(&doc, "u", &params(&[("x-max-tokens", "3")])).unwrap();
407        assert_eq!(r.body, "one two three");
408        assert!(r.warnings[0].starts_with("output_truncated"));
409        let e = render(&doc, "u", &params(&[("x-token-budget", "5")])).unwrap_err();
410        assert_eq!(e.code, ErrorCode::TooLarge);
411        let ok = render(&doc, "u", &params(&[("x-respond-with", "frontmatter")])).unwrap();
412        assert!(ok.body.starts_with("---\ntitle: \"\"\n"));
413    }
414}
pagesift_api/jina/respond.rs

pagesift_api/jina/
respond.rs
