Skip to main content

pagesift_engine/engine/
extract.rs

1use super::*;
2
3/// extract 最多处理多少个页面,每页最多取多少字符的正文。
4/// [策略] 用户策略:一次 extract 交给 LLM 的输入上限(调用成本和上下文窗口由用户的模型决定),
5/// 先验包 v0:0.81 的现值
6pub(super) const EXTRACT_MAX_PAGES: usize = 20;
7/// [策略] 同上,每页的字符上限
8pub(super) const EXTRACT_MAX_CHARS_PER_PAGE: usize = 20_000;
9
10impl Scraper {
11    pub async fn search(
12        &self,
13        req: &pagesift_api::types::SearchRequest,
14    ) -> Result<serde_json::Value, ApiError> {
15        let Some(backend) = &self.search else {
16            return Err(ApiError::unsupported(
17                "search is not configured: set [search] in pagesift.toml",
18            ));
19        };
20        let mut q = pagesift_search::SearchQuery::new(req.query.clone(), req.limit);
21        q.lang = req.lang.clone();
22        q.country = req.country.clone();
23        q.time_range = match req.tbs.as_deref() {
24            Some("qdr:d") => Some(pagesift_search::TimeRange::Day),
25            Some("qdr:w") => Some(pagesift_search::TimeRange::Week),
26            Some("qdr:m") => Some(pagesift_search::TimeRange::Month),
27            Some("qdr:y") => Some(pagesift_search::TimeRange::Year),
28            _ => None,
29        };
30        let search = backend.search(&q);
31        let results = match req.timeout {
32            Some(ms) => tokio::time::timeout(Duration::from_millis(ms), search)
33                .await
34                .map_err(|_| ApiError::new(ErrorCode::Timeout, "search timed out"))?,
35            None => search.await,
36        }
37        .map_err(|e| {
38            ApiError::new(
39                ErrorCode::UpstreamError,
40                format!("search backend failed: {e}"),
41            )
42        })?;
43
44        // 给了 scrapeOptions.formats 才抓取结果页面,并把文档字段合并进对应的结果
45        let scrape = req
46            .scrape_options
47            .as_ref()
48            .filter(|o| !o.formats.is_empty());
49        let docs: Vec<Option<Document>> = match scrape {
50            Some(opts) => {
51                let futs = results
52                    .iter()
53                    .map(|r| async move { self.scrape(&r.url, opts).await.ok() });
54                futures_util::future::join_all(futs).await
55            }
56            None => vec![None; results.len()],
57        };
58        let web: Vec<serde_json::Value> = results
59            .into_iter()
60            .zip(docs)
61            .map(|(r, doc)| {
62                let mut item = serde_json::json!({
63                    "url": r.url,
64                    "title": r.title,
65                    "description": r.description,
66                    "position": r.position,
67                });
68                if let Some(doc) = doc
69                    && let (Some(obj), Ok(serde_json::Value::Object(d))) =
70                        (item.as_object_mut(), serde_json::to_value(&doc))
71                {
72                    for (k, v) in d {
73                        obj.entry(k).or_insert(v);
74                    }
75                }
76                item
77            })
78            .collect();
79        Ok(serde_json::json!({ "web": web }))
80    }
81
82    pub async fn extract(
83        &self,
84        req: &pagesift_api::types::ExtractRequest,
85    ) -> Result<serde_json::Value, ApiError> {
86        let Some(llm) = &self.llm else {
87            return Err(ApiError::unsupported(
88                "extract needs an LLM: set [llm] in pagesift.toml",
89            ));
90        };
91        let mut urls: Vec<String> = Vec::new();
92        for u in &req.urls {
93            if let Some(site) = u.strip_suffix("/*") {
94                let map: pagesift_api::types::MapRequest = serde_json::from_value(
95                    serde_json::json!({ "url": site, "limit": EXTRACT_MAX_PAGES }),
96                )
97                .map_err(|e| ApiError::invalid(e.to_string()))?;
98                urls.extend(self.map(&map).await?.into_iter().map(|l| l.url));
99            } else {
100                urls.push(u.clone());
101            }
102        }
103        urls.truncate(EXTRACT_MAX_PAGES);
104        let opts = req.scrape_options.clone().unwrap_or_default();
105        let opts = &opts;
106        let futs = urls
107            .iter()
108            .map(|u| async move { (u, self.scrape(u, opts).await) });
109        let mut corpus = String::new();
110        let mut errors = Vec::new();
111        for (u, r) in futures_util::future::join_all(futs).await {
112            match r {
113                Ok(doc) => {
114                    let md: String = doc
115                        .markdown
116                        .unwrap_or_default()
117                        .chars()
118                        .take(EXTRACT_MAX_CHARS_PER_PAGE)
119                        .collect();
120                    corpus.push_str(&format!("\n\n## Source: {u}\n\n{md}"));
121                }
122                Err(e) => errors.push(format!("{u}: {}", e.code.as_str())),
123            }
124        }
125        if corpus.is_empty() {
126            return Err(ApiError::new(
127                ErrorCode::UpstreamError,
128                format!("none of the pages could be scraped ({})", errors.join(", ")),
129            ));
130        }
131        let schema = req
132            .schema
133            .clone()
134            .unwrap_or_else(|| serde_json::json!({ "type": "object" }));
135        let eo = pagesift_llm::ExtractOptions {
136            prompt: req.prompt.clone(),
137        };
138        pagesift_llm::extract_json(llm, &corpus, &schema, &eo)
139            .await
140            .map_err(llm_error)
141    }
142}
143
144pub(super) fn llm_error(e: pagesift_llm::LlmError) -> ApiError {
145    ApiError::new(ErrorCode::UpstreamError, format!("LLM request failed: {e}"))
146}