pagesift_engine/engine/
extract.rs1use super::*;
2
3pub(super) const EXTRACT_MAX_PAGES: usize = 20;
7pub(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 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}