"""深度研究流水线(净室实现):拆解 → 检索 → 阅读摘要 → 合成带引用报告。 本模块只依赖标准库与 untrusted。所有 IO(LLM 调用、检索、抓取、状态更新、取消检查) 均通过依赖注入传入,因此可在不安装 httpx/psycopg/redis 的环境下用 unittest 测试。 注意:这是对"计划-检索-阅读-合成"这一通用研究方法的独立实现,未复制任何第三方代码。 """ import json import re from untrusted import untrusted_message class CanceledError(Exception): """任务在执行中被取消。""" _PLAN_SYS = ( "你是政务研究助理。把用户的研究题目拆解为 3-5 个互补的检索式子问题," "覆盖背景、现行政策、数据现状、影响与对策等角度。" "只输出一个 JSON 字符串数组(如 [\"...\", \"...\"]),不要任何其它文字。" ) def parse_json_list(text: str) -> list[str]: """从 LLM 输出中稳健地解析出字符串数组(容忍代码围栏与多余文字)。""" if not text: return [] t = text.strip() t = re.sub(r"^```(?:json)?", "", t).strip() t = re.sub(r"```$", "", t).strip() m = re.search(r"\[.*\]", t, re.S) if m: t = m.group(0) try: data = json.loads(t) if isinstance(data, list): return [str(x).strip() for x in data if str(x).strip()] except Exception: pass # 兜底:按行拆分并去掉列表符号 out = [] for line in text.splitlines(): s = line.strip().lstrip("-*0123456789. ").strip() if s: out.append(s) return out[:5] def dedupe_sources(raw: list[dict], limit: int) -> list[dict]: """按 URL 去重并归一化检索结果,截断到 limit 条。""" seen = set() out = [] for r in raw or []: url = (r.get("url") or "").strip() if not url or url in seen: continue seen.add(url) out.append({ "title": (r.get("title") or url).strip(), "url": url, "snippet": r.get("snippet", "") or "", "content": r.get("content", "") or "", }) if len(out) >= limit: break return out class ResearchPipeline: def __init__( self, task_id: str, task: dict, *, llm_chat, # callable(messages: list[dict]) -> str search_fn, # callable(query: str, k: int) -> list[dict] fetch_fn=None, # callable(url: str) -> str(可选,抓取正文) update_fn=None, # callable(**fields)(落库,可选) progress_cb=None, # callable(status, progress, message)(可选) is_canceled=None, # callable() -> bool(可选) max_subqueries: int = 4, max_sources: int = 6, ): self.task_id = task_id self.task = task self.llm_chat = llm_chat self.search_fn = search_fn self.fetch_fn = fetch_fn self.update_fn = update_fn self.progress_cb = progress_cb self.is_canceled = is_canceled cfg = task.get("config") or {} self.max_subqueries = int(cfg.get("max_steps", max_subqueries)) self.max_sources = int(cfg.get("max_sources", max_sources)) self.tokens_used = 0 # ---------------- 主流程 ---------------- def run(self) -> dict: topic = (self.task.get("topic") or "").strip() if not topic: raise ValueError("研究题目为空") self._progress("planning", 10, "拆解研究问题…") self._check_cancel() subqueries = self._plan(topic) self._progress("searching", 30, "检索资料…") self._check_cancel() sources = self._search(subqueries) self._progress("reading", 55, "阅读与摘要…") self._check_cancel() findings = self._read(topic, sources) self._progress("synthesizing", 80, "合成报告…") self._check_cancel() report = self._synthesize(topic, findings) result = { "report": report, "sources": [ {"title": s["title"], "url": s["url"], "snippet": s.get("snippet", "")} for s in sources ], "tokens_used": self.tokens_used, } if self.update_fn: self.update_fn( status="completed", progress=100, report=report, sources=result["sources"], tokens_used=self.tokens_used, ) self._progress("completed", 100, "完成") return result # ---------------- 各阶段 ---------------- def _plan(self, topic: str) -> list[str]: out = self._chat([ {"role": "system", "content": _PLAN_SYS}, {"role": "user", "content": f"研究题目:{topic}"}, ]) qs = parse_json_list(out) if not qs: qs = [topic] return qs[: self.max_subqueries] def _search(self, subqueries: list[str]) -> list[dict]: raw: list[dict] = [] for q in subqueries: try: raw.extend(self.search_fn(q, self.max_sources) or []) except Exception: continue return dedupe_sources(raw, self.max_sources) def _read(self, topic: str, sources: list[dict]) -> list[dict]: findings = [] for i, s in enumerate(sources, 1): self._check_cancel() content = s.get("content") or s.get("snippet") or "" if not content and self.fetch_fn and s.get("url"): try: content = self.fetch_fn(s["url"]) or "" except Exception: content = "" if not content: continue summary = self._chat([ {"role": "system", "content": ( f"你在为研究题目「{topic}」整理资料。请从下面这条外部资料中提取与题目相关的" "关键事实,用要点列出;与题目无关则只回复『无相关内容』。不要编造。" )}, untrusted_message(f"来源[{i}] {s['title']}", content), ]) if summary and "无相关内容" not in summary: findings.append({"idx": i, "title": s["title"], "url": s["url"], "summary": summary}) return findings def _synthesize(self, topic: str, findings: list[dict]) -> str: if not findings: return self._chat([ {"role": "system", "content": ( "你是政务研究助理。本次未检索到可用的外部资料,请基于既有知识撰写结构化报告," "并在报告开头明确声明『本报告未使用外部检索资料,仅供参考』。不要编造引用来源。" )}, {"role": "user", "content": f"研究题目:{topic}"}, ]) sys = ( "你是政务研究助理,撰写结构化 Markdown 研究报告。要求:" "1) 只依据下面提供的『发现』素材,不得编造;" "2) 正文用行内角标 [n] 标注信息来源(n 对应发现编号);" "3) 结构包含:摘要、背景、关键发现、分析、结论与建议;" "4) 末尾输出『## 参考来源』,按 [n] 列出标题与链接。" ) messages = [ {"role": "system", "content": sys}, {"role": "user", "content": f"研究题目:{topic}"}, ] for f in findings: messages.append(untrusted_message(f"发现[{f['idx']}] {f['title']}", f["summary"])) report = self._chat(messages) # 兜底:若模型未输出来源清单,则补一份,保证可溯源 if "参考来源" not in report: lines = ["", "", "## 参考来源"] for f in findings: lines.append(f"[{f['idx']}] {f['title']} - {f['url']}") report += "\n".join(lines) return report # ---------------- 工具 ---------------- def _chat(self, messages: list[dict]) -> str: text = self.llm_chat(messages) or "" # 无 usage 信息时的粗略 token 估算 self.tokens_used += sum(len(m.get("content", "")) for m in messages) // 4 + len(text) // 4 return text def _progress(self, status: str, progress: int, message: str): if self.progress_cb: self.progress_cb(status, progress, message) def _check_cancel(self): if self.is_canceled and self.is_canceled(): raise CanceledError("任务已取消")