Skip to main content

pagesift_api/types/
parse.rs

1//! `/v2/parse`(上传文件解析)的输入类型:上传的文件、文件类型判断、`parsers` 里的 PDF 选项。
2//! 字段和取值照 Firecrawl 文档 <https://docs.firecrawl.dev/api-reference/endpoint/parse>
3//! 与源码 `apps/api/src/controllers/v2/parse.ts`。
4
5use serde_json::Value;
6
7use crate::error::ApiError;
8
9/// 调用方上传的文件。字节只有这一份,交给引擎后不再复制。
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub struct Upload {
12    pub bytes: Vec<u8>,
13    /// multipart 里的文件名(没有时为 `upload`),作为文档的 `metadata.sourceURL`。
14    pub filename: String,
15    /// multipart 里这个文件的 Content-Type。
16    pub content_type: Option<String>,
17    pub kind: UploadKind,
18    /// 调用方给的页面原地址(`url` 字段):HTML 的相对链接按它解析,也作为 `metadata.sourceURL`。
19    /// 只用来解析地址,不会去抓。
20    pub url: Option<url::Url>,
21}
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum UploadKind {
25    Pdf,
26    Html,
27    /// Word、Excel、PowerPoint、EPUB、CSV、邮件、压缩包、文本类(见 pagesift-office)。
28    Office(pagesift_office::Format),
29    /// 图片(要服务端打开 OCR)。参数是按内容认出的 MIME 类型。
30    Image(&'static str),
31}
32
33impl UploadKind {
34    /// 判断上传的文件类型:内容以 `%PDF-` 开头就是 PDF;ZIP 容器按里面的部件认 Word、Excel、
35    /// PowerPoint、EPUB(不管文件名怎么写);否则看扩展名,再看 Content-Type。
36    /// 和 Firecrawl 一样,认不出的类型不猜,直接拒绝。
37    pub fn detect(bytes: &[u8], filename: &str, content_type: Option<&str>) -> Option<Self> {
38        if bytes.starts_with(b"%PDF-") {
39            return Some(UploadKind::Pdf);
40        }
41        if let Some(mime) = image_mime(bytes) {
42            return Some(UploadKind::Image(mime));
43        }
44        if let Some(f) = pagesift_office::Format::detect(bytes, Some(filename), content_type) {
45            return Some(UploadKind::Office(f));
46        }
47        let ext = filename
48            .rsplit_once('.')
49            .map(|(_, e)| e.to_ascii_lowercase())
50            .unwrap_or_default();
51        let mime = content_type
52            .and_then(|c| c.split(';').next())
53            .map(|c| c.trim().to_ascii_lowercase())
54            .unwrap_or_default();
55        if ext == "pdf" || matches!(mime.as_str(), "application/pdf" | "application/x-pdf") {
56            return Some(UploadKind::Pdf);
57        }
58        if matches!(ext.as_str(), "html" | "htm" | "xhtml")
59            || matches!(mime.as_str(), "text/html" | "application/xhtml+xml")
60        {
61            return Some(UploadKind::Html);
62        }
63        None
64    }
65
66    pub fn mime(self) -> &'static str {
67        match self {
68            UploadKind::Pdf => "application/pdf",
69            UploadKind::Html => "text/html",
70            UploadKind::Office(f) => f.mime(),
71            UploadKind::Image(mime) => mime,
72        }
73    }
74
75    /// 网页存档里也按这种类型转换吗?和抓取时一致:文本类和图片不转。
76    pub fn fetchable(self) -> bool {
77        match self {
78            UploadKind::Pdf | UploadKind::Html => true,
79            UploadKind::Office(f) => f.fetchable(),
80            UploadKind::Image(_) => false,
81        }
82    }
83}
84
85/// 按签名认图片。pagesift-api 不依赖 PDF 解码,这里只看开头几个字节。
86fn image_mime(b: &[u8]) -> Option<&'static str> {
87    if b.starts_with(&[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A]) {
88        Some("image/png")
89    } else if b.starts_with(&[0xFF, 0xD8, 0xFF]) {
90        Some("image/jpeg")
91    } else if b.starts_with(b"GIF87a") || b.starts_with(b"GIF89a") {
92        Some("image/gif")
93    } else if b.starts_with(b"II*\0") || b.starts_with(b"MM\0*") {
94        Some("image/tiff")
95    } else if b.len() >= 12 && &b[..4] == b"RIFF" && &b[8..12] == b"WEBP" {
96        Some("image/webp")
97    } else if b.starts_with(b"BM")
98        // `BM` 太短,文字也可能这样开头:再看信息头的长度字段
99        && b.get(14..18).is_some_and(|s| {
100            matches!(
101                u32::from_le_bytes([s[0], s[1], s[2], s[3]]),
102                12 | 40 | 52 | 56 | 64 | 108 | 124
103            )
104        })
105    {
106        Some("image/bmp")
107    } else {
108        None
109    }
110}
111
112/// `parsers` 里 `{type: "pdf"}` 的选项。
113#[derive(Debug, Clone, Default, PartialEq, Eq)]
114pub struct PdfParser {
115    /// `fast`、`auto`(默认)或 `ocr`。
116    pub mode: Option<String>,
117    /// 只输出前这么多页(1–10000)。
118    pub max_pages: Option<usize>,
119    /// 另给每页的 Markdown(`data.pages`)。
120    pub pages: bool,
121    /// 给了 `parsers` 却没有 `pdf`(比如 `[]`):和 Firecrawl 一样不解析,把原文件按 base64
122    /// 放进 `markdown`。
123    pub raw: bool,
124}
125
126impl PdfParser {
127    /// 从 `parsers` 取 PDF 选项。接受 Firecrawl 的写法:`["pdf"]`、`[{type: "pdf", ...}]`,
128    /// 以及 `"image"`(图片 OCR,由服务端的 OCR 开关决定,这里照收)。
129    /// 列表里没有 `pdf` 时 PDF 不解析([`Self::raw`])。
130    /// 做不到的选项(版面块、分页标记)明确报 UNSUPPORTED_OPTION,不悄悄忽略;`mode: "ocr"` 要
131    /// 服务端打开了 OCR,这里不知道,由解析 PDF 时检查。
132    pub fn from_parsers(parsers: Option<&Value>) -> Result<Self, ApiError> {
133        let items = match parsers {
134            None | Some(Value::Null) => return Ok(Self::default()),
135            Some(Value::Array(items)) => items,
136            Some(_) => return Err(ApiError::invalid("parsers must be an array")),
137        };
138        let mut out = Self {
139            raw: true,
140            ..Self::default()
141        };
142        for item in items {
143            let kind = match item {
144                Value::String(t) => t.as_str(),
145                Value::Object(o) => o.get("type").and_then(Value::as_str).unwrap_or_default(),
146                _ => "",
147            };
148            let obj = match (kind, item) {
149                ("image", _) => continue,
150                ("pdf", Value::String(_)) => {
151                    out.raw = false;
152                    continue;
153                }
154                ("pdf", Value::Object(o)) => {
155                    out.raw = false;
156                    o
157                }
158                _ => {
159                    return Err(ApiError::invalid(
160                        "parsers items must be \"pdf\", \"image\" or {\"type\": \"pdf\" | \"image\", ...}",
161                    ));
162                }
163            };
164            for (key, v) in obj {
165                match key.as_str() {
166                    "type" => {}
167                    "mode" => match v.as_str() {
168                        Some(m @ ("fast" | "auto" | "ocr")) => out.mode = Some(m.into()),
169                        _ => {
170                            return Err(ApiError::invalid(
171                                "parsers[].mode must be fast, auto or ocr",
172                            ));
173                        }
174                    },
175                    "maxPages" => match v.as_u64() {
176                        Some(n @ 1..=10_000) => out.max_pages = Some(n as usize),
177                        _ => {
178                            return Err(ApiError::invalid(
179                                "parsers[].maxPages must be an integer from 1 to 10000",
180                            ));
181                        }
182                    },
183                    "pages" => out.pages = flag(key, v)?,
184                    "blocks" | "pageMarkers" => {
185                        if flag(key, v)? {
186                            return Err(ApiError::unsupported(format!(
187                                "parsers[].{key} is not supported by this build"
188                            )));
189                        }
190                    }
191                    _ => {
192                        return Err(ApiError::invalid(format!(
193                            "unknown PDF parser option {key:?}"
194                        )));
195                    }
196                }
197            }
198        }
199        Ok(out)
200    }
201}
202
203fn flag(key: &str, v: &Value) -> Result<bool, ApiError> {
204    v.as_bool()
205        .ok_or_else(|| ApiError::invalid(format!("parsers[].{key} must be a boolean")))
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211    use crate::error::ErrorCode;
212    use serde_json::json;
213
214    #[test]
215    fn kind_prefers_content_then_name_then_mime() {
216        let d = UploadKind::detect;
217        assert_eq!(
218            d(b"%PDF-1.7", "a.html", Some("text/html")),
219            Some(UploadKind::Pdf)
220        );
221        assert_eq!(d(b"x", "Report.PDF", None), Some(UploadKind::Pdf));
222        assert_eq!(
223            d(b"x", "upload", Some("application/pdf; q=1")),
224            Some(UploadKind::Pdf)
225        );
226        assert_eq!(d(b"<p>", "a.htm", None), Some(UploadKind::Html));
227        assert_eq!(
228            d(b"<p>", "blob", Some("Text/HTML; charset=utf-8")),
229            Some(UploadKind::Html)
230        );
231        // 内容不是 ZIP 的 .docx 认不出;ZIP 按里面的部件认
232        assert_eq!(d(b"PK", "a.docx", None), None);
233        assert_eq!(
234            d(b"a,b\n1,2", "t.csv", None),
235            Some(UploadKind::Office(pagesift_office::Format::Csv))
236        );
237        assert_eq!(
238            d(b"a,b", "upload", Some("text/csv")),
239            Some(UploadKind::Office(pagesift_office::Format::Csv))
240        );
241        assert_eq!(d(b"x", "upload", Some("application/octet-stream")), None);
242        assert_eq!(
243            d(b"\x89PNG\r\n\x1a\nrest", "a.txt", None),
244            Some(UploadKind::Image("image/png"))
245        );
246        assert!(!UploadKind::Image("image/png").fetchable());
247        assert!(!UploadKind::Office(pagesift_office::Format::Text).fetchable());
248        assert!(UploadKind::Office(pagesift_office::Format::Eml).fetchable());
249        // a.html 以 text/plain 上传仍然是 HTML
250        assert_eq!(
251            d(b"<p>", "a.html", Some("text/plain")),
252            Some(UploadKind::Html)
253        );
254    }
255
256    #[test]
257    fn pdf_parser_options() {
258        let p = |v| PdfParser::from_parsers(Some(&v));
259        assert_eq!(p(json!(["pdf"])).unwrap(), PdfParser::default());
260        assert_eq!(p(json!(["pdf", "image"])).unwrap(), PdfParser::default());
261        assert_eq!(
262            p(json!([{"type": "pdf", "mode": "fast", "maxPages": 3, "pages": true}])).unwrap(),
263            PdfParser {
264                mode: Some("fast".into()),
265                max_pages: Some(3),
266                pages: true,
267                raw: false,
268            }
269        );
270        // 列表里没有 pdf:不解析,原样返回
271        for v in [json!([]), json!(["image"]), json!([{"type": "image"}])] {
272            assert!(p(v.clone()).unwrap().raw, "{v}");
273        }
274        assert!(!PdfParser::from_parsers(None).unwrap().raw);
275        // 关着的 blocks / pageMarkers 不算请求
276        assert!(p(json!([{"type": "pdf", "blocks": false, "pageMarkers": false}])).is_ok());
277        assert_eq!(
278            p(json!([{"type": "pdf", "mode": "ocr"}]))
279                .unwrap()
280                .mode
281                .as_deref(),
282            Some("ocr")
283        );
284        for (v, code) in [
285            (
286                json!([{"type": "pdf", "blocks": true}]),
287                ErrorCode::UnsupportedOption,
288            ),
289            (
290                json!([{"type": "pdf", "maxPages": 0}]),
291                ErrorCode::InvalidRequest,
292            ),
293            (
294                json!([{"type": "pdf", "maxPages": 10001}]),
295                ErrorCode::InvalidRequest,
296            ),
297            (
298                json!([{"type": "pdf", "mode": "x"}]),
299                ErrorCode::InvalidRequest,
300            ),
301            (json!([{"type": "docx"}]), ErrorCode::InvalidRequest),
302            (json!({"type": "pdf"}), ErrorCode::InvalidRequest),
303        ] {
304            assert_eq!(p(v.clone()).unwrap_err().code, code, "{v}");
305        }
306    }
307}