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        self.search_with_budget(req, None).await
16    }
17
18    pub async fn search_with_budget(
19        &self,
20        req: &pagesift_api::types::SearchRequest,
21        budget: Option<&dyn pagesift_api::PageBudget>,
22    ) -> Result<serde_json::Value, ApiError> {
23        let Some(backend) = &self.search else {
24            return Err(ApiError::unsupported(
25                "search is not configured: set [search] in pagesift.toml",
26            ));
27        };
28        let mut q = pagesift_search::SearchQuery::new(req.query.clone(), req.limit);
29        q.lang = req.lang.clone();
30        q.country = req.country.clone();
31        q.time_range = match req.tbs.as_deref() {
32            Some("qdr:d") => Some(pagesift_search::TimeRange::Day),
33            Some("qdr:w") => Some(pagesift_search::TimeRange::Week),
34            Some("qdr:m") => Some(pagesift_search::TimeRange::Month),
35            Some("qdr:y") => Some(pagesift_search::TimeRange::Year),
36            _ => None,
37        };
38        let search = backend.search(&q);
39        let results = match req.timeout {
40            Some(ms) => tokio::time::timeout(Duration::from_millis(ms), search)
41                .await
42                .map_err(|_| ApiError::new(ErrorCode::Timeout, "search timed out"))?,
43            None => search.await,
44        }
45        .map_err(|e| {
46            ApiError::new(
47                ErrorCode::UpstreamError,
48                format!("search backend failed: {e}"),
49            )
50        })?;
51
52        // 给了 scrapeOptions.formats 才抓取结果页面，并把文档字段合并进对应的结果
53        let scrape = req
54            .scrape_options
55            .as_ref()
56            .filter(|o| !o.formats.is_empty());
57        let docs: Vec<Option<Document>> = match scrape {
58            Some(opts) => {
59                if let Some(budget) = budget {
60                    let mut docs = Vec::with_capacity(results.len());
61                    for result in &results {
62                        budget.acquire().await?;
63                        docs.push(self.scrape(&result.url, opts).await.ok());
64                    }
65                    docs
66                } else {
67                    let futs = results
68                        .iter()
69                        .map(|r| async move { self.scrape(&r.url, opts).await.ok() });
70                    futures_util::future::join_all(futs).await
71                }
72            }
73            None => vec![None; results.len()],
74        };
75        let web: Vec<serde_json::Value> = results
76            .into_iter()
77            .zip(docs)
78            .map(|(r, doc)| {
79                let mut item = serde_json::json!({
80                    "url": r.url,
81                    "title": r.title,
82                    "description": r.description,
83                    "position": r.position,
84                });
85                if let Some(doc) = doc
86                    && let (Some(obj), Ok(serde_json::Value::Object(d))) =
87                        (item.as_object_mut(), serde_json::to_value(&doc))
88                {
89                    for (k, v) in d {
90                        obj.entry(k).or_insert(v);
91                    }
92                }
93                item
94            })
95            .collect();
96        Ok(serde_json::json!({ "web": web }))
97    }
98
99    pub async fn extract(
100        &self,
101        req: &pagesift_api::types::ExtractRequest,
102    ) -> Result<serde_json::Value, ApiError> {
103        self.extract_with_budget(req, None).await
104    }
105
106    pub async fn extract_with_budget(
107        &self,
108        req: &pagesift_api::types::ExtractRequest,
109        budget: Option<&dyn pagesift_api::PageBudget>,
110    ) -> Result<serde_json::Value, ApiError> {
111        let Some(llm) = &self.llm else {
112            return Err(ApiError::unsupported(
113                "extract needs an LLM: set [llm] in pagesift.toml",
114            ));
115        };
116        let mut urls: Vec<String> = Vec::new();
117        for u in &req.urls {
118            if let Some(site) = u.strip_suffix("/*") {
119                let map: pagesift_api::types::MapRequest = serde_json::from_value(
120                    serde_json::json!({ "url": site, "limit": EXTRACT_MAX_PAGES }),
121                )
122                .map_err(|e| ApiError::invalid(e.to_string()))?;
123                urls.extend(self.map(&map).await?.into_iter().map(|l| l.url));
124            } else {
125                urls.push(u.clone());
126            }
127        }
128        urls.truncate(EXTRACT_MAX_PAGES);
129        let opts = req.scrape_options.clone().unwrap_or_default();
130        let opts = &opts;
131        let results = if let Some(budget) = budget {
132            let mut results = Vec::with_capacity(urls.len());
133            for u in &urls {
134                budget.acquire().await?;
135                results.push((u, self.scrape(u, opts).await));
136            }
137            results
138        } else {
139            let futs = urls
140                .iter()
141                .map(|u| async move { (u, self.scrape(u, opts).await) });
142            futures_util::future::join_all(futs).await
143        };
144        let mut corpus = String::new();
145        let mut errors = Vec::new();
146        for (u, r) in results {
147            match r {
148                Ok(doc) => {
149                    let md: String = doc
150                        .markdown
151                        .unwrap_or_default()
152                        .chars()
153                        .take(EXTRACT_MAX_CHARS_PER_PAGE)
154                        .collect();
155                    corpus.push_str(&format!("\n\n## Source: {u}\n\n{md}"));
156                }
157                Err(e) => errors.push(format!("{u}: {}", e.code.as_str())),
158            }
159        }
160        if corpus.is_empty() {
161            return Err(ApiError::new(
162                ErrorCode::UpstreamError,
163                format!("none of the pages could be scraped ({})", errors.join(", ")),
164            ));
165        }
166        let schema = req
167            .schema
168            .clone()
169            .unwrap_or_else(|| serde_json::json!({ "type": "object" }));
170        let eo = pagesift_llm::ExtractOptions {
171            prompt: req.prompt.clone(),
172        };
173        pagesift_llm::extract_json(llm, &corpus, &schema, &eo)
174            .await
175            .map_err(llm_error)
176    }
177}
178
179pub(super) fn llm_error(e: pagesift_llm::LlmError) -> ApiError {
180    ApiError::new(ErrorCode::UpstreamError, format!("LLM request failed: {e}"))
181}
pagesift_engine/engine/extract.rs

pagesift_engine/engine/
extract.rs
