315 lines
15 KiB
Python
315 lines
15 KiB
Python
"""services/paper/agent.py —— 论文对话规则状态机（brief→outline→writing→complete/revise）。
|
||
|
||
LLM 仅做 brief 字段抽取（complete_json）；大纲/正文/引用/合并全部由 paper 服务完成。
|
||
不引入 function calling（BFF 无工具循环先例）。助手消息由本模块进程内直写 messages 表。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import uuid
|
||
from datetime import datetime, timezone
|
||
from typing import Any
|
||
|
||
from ...config import get_settings
|
||
from ..docgen.llm import complete_json as llm_complete_json
|
||
from .mcp_client import PaperMcpClient, PaperMcpError
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
BRIEF_FIELDS = [
|
||
{"key": "topic", "label": "论文主题", "options": None},
|
||
{"key": "paper_type", "label": "论文类型", "options": ["综述", "研究论文", "开题报告", "实验报告"]},
|
||
{"key": "language", "label": "语言", "options": ["中文", "英文"]},
|
||
{"key": "length", "label": "篇幅（字数）", "options": ["3000字", "5000字", "8000字"]},
|
||
{"key": "format_req", "label": "格式要求", "options": ["GB/T 7714", "APA", "MLA"]},
|
||
{"key": "style", "label": "写作风格", "options": ["学术严谨", "通俗易懂"]},
|
||
{"key": "requirements", "label": "其他要求", "options": None},
|
||
]
|
||
|
||
REQUIRED_BRIEF_KEYS = ("topic", "paper_type", "language", "length")
|
||
|
||
_BRIEF_SYSTEM = (
|
||
"你是论文撰写需求抽取助手。从用户输入中抽取论文写作 brief。"
|
||
"只输出一个合法 JSON 对象：{\"topic\":\"\",\"paper_type\":\"\",\"language\":\"\","
|
||
"\"length\":\"\",\"format_req\":\"\",\"style\":\"\",\"requirements\":\"\"}。"
|
||
"抽取不到的字段填空字符串。"
|
||
)
|
||
|
||
|
||
def _db_now() -> str:
|
||
return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||
|
||
|
||
def _text_block(html: str) -> dict[str, Any]:
|
||
return {"type": "text", "html": html}
|
||
|
||
|
||
def _new_client() -> PaperMcpClient:
|
||
s = get_settings()
|
||
return PaperMcpClient(s.paper_writer_url, s.paper_writer_read_timeout)
|
||
|
||
|
||
async def _write_message(
|
||
db,
|
||
session_id: str,
|
||
*,
|
||
kind: str = "paper",
|
||
status: str = "done",
|
||
blocks: list[dict[str, Any]],
|
||
text: str | None = None,
|
||
) -> str:
|
||
"""直写助手消息（kind=paper），返回 message_id。"""
|
||
mid = str(uuid.uuid4())
|
||
now = _db_now()
|
||
await db.execute(
|
||
"INSERT INTO messages (id, session_id, role, kind, status, blocks, raw, text, created_at, updated_at) "
|
||
"VALUES (?, ?, 'assistant', ?, ?, ?, NULL, ?, ?, ?)",
|
||
(mid, session_id, kind, status, json.dumps(blocks, ensure_ascii=False), text, now, now),
|
||
)
|
||
await db.execute(
|
||
"UPDATE sessions SET status = CASE WHEN status = 'idle' THEN 'running' ELSE status END, updated_at = ? WHERE id = ?",
|
||
(now, session_id),
|
||
)
|
||
return mid
|
||
|
||
|
||
async def extract_brief(text: str) -> dict[str, str]:
|
||
"""LLM 抽取 brief 字段（失败返回全空，走 paper_brief 卡兜底追问）。"""
|
||
try:
|
||
brief = await llm_complete_json(_BRIEF_SYSTEM, text, max_tokens=1000)
|
||
except Exception as exc: # noqa: BLE001
|
||
logger.warning("paper: brief 抽取失败：%s", exc)
|
||
brief = {}
|
||
return {f["key"]: str((brief or {}).get(f["key"]) or "").strip() for f in BRIEF_FIELDS}
|
||
|
||
|
||
def _missing_brief_keys(brief: dict[str, str]) -> list[str]:
|
||
return [k for k in REQUIRED_BRIEF_KEYS if not brief.get(k)]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 状态机入口
|
||
# ---------------------------------------------------------------------------
|
||
|
||
async def handle_turn(db, session_id: str, text: str, paper_id: str | None = None, template_key: str = "") -> list[dict[str, Any]]:
|
||
"""一轮对话：推进状态机并落库助手消息，返回消息的 blocks 列表。"""
|
||
client = _new_client()
|
||
if not paper_id:
|
||
return await _trigger(db, session_id, text, client, template_key)
|
||
return await _followup(db, session_id, text, paper_id, client, template_key)
|
||
|
||
|
||
async def _trigger(db, session_id: str, text: str, client: PaperMcpClient, template_key: str = "") -> list[dict[str, Any]]:
|
||
brief = await extract_brief(text)
|
||
missing = _missing_brief_keys(brief)
|
||
if missing:
|
||
fields = [{**f, "value": brief.get(f["key"], "")} for f in BRIEF_FIELDS]
|
||
blocks = [
|
||
_text_block("好的，我来帮你写这篇论文。请补充以下信息（也可以直接打字告诉我）："),
|
||
{"type": "paper_brief", "paper_id": "", "status": "pending", "fields": fields},
|
||
]
|
||
await _write_message(db, session_id, blocks=blocks, text="请补充论文信息")
|
||
return blocks
|
||
|
||
outline = await client.call_tool("propose_paper_outline", {**brief, "template_key": template_key})
|
||
if not outline.get("ok"):
|
||
blocks = [_text_block(f"论文大纲生成失败：{outline.get('message', '服务未就绪')}，请重试。")]
|
||
await _write_message(db, session_id, blocks=blocks, text=blocks[0]["html"])
|
||
return blocks
|
||
return await _outline_message(db, session_id, outline)
|
||
|
||
|
||
async def _outline_message(db, session_id: str, outline: dict[str, Any]) -> list[dict[str, Any]]:
|
||
paper_id = outline["paper_id"]
|
||
sections = [
|
||
{"section_id": s.get("section_id"), "title": s.get("title", ""),
|
||
"key_points": s.get("key_points") or [], "selected": True}
|
||
for s in (outline.get("sections") or [])
|
||
]
|
||
blocks = [
|
||
_text_block(
|
||
f"已为你生成论文大纲《{outline.get('title', '')}》。"
|
||
"你可以勾选/编辑章节，或直接说修改意见；确认后我开始逐节撰写。"
|
||
),
|
||
{"type": "paper_outline", "paper_id": paper_id, "title": outline.get("title", ""),
|
||
"status": "pending", "sections": sections},
|
||
]
|
||
await _write_message(db, session_id, blocks=blocks, text=blocks[0]["html"])
|
||
return blocks
|
||
|
||
|
||
_CANCEL_WORDS = ("暂停", "取消", "停止", "不写了", "终止", "中止", "算了")
|
||
|
||
|
||
def _is_cancel_request(text: str) -> bool:
|
||
s = (text or "").strip()
|
||
return bool(s) and any(w in s for w in _CANCEL_WORDS)
|
||
|
||
|
||
async def _followup(db, session_id: str, text: str, paper_id: str, client: PaperMcpClient, template_key: str = "") -> list[dict[str, Any]]:
|
||
prog = await client.call_tool("get_paper_progress", {"paper_id": paper_id})
|
||
if not prog.get("ok"):
|
||
blocks = [_text_block(f"未找到论文（{paper_id}），请重新发起撰写。")]
|
||
await _write_message(db, session_id, blocks=blocks, text=blocks[0]["html"])
|
||
return blocks
|
||
|
||
status = prog.get("status")
|
||
if _is_cancel_request(text):
|
||
res = await client.call_tool("cancel_paper", {"paper_id": paper_id})
|
||
if res.get("ok"):
|
||
blocks = [_text_block("论文已取消。如需重新撰写，请告诉我新主题。")]
|
||
else:
|
||
blocks = [_text_block(f"取消失败：{res.get('message', '')}")]
|
||
await _write_message(db, session_id, blocks=blocks, text=blocks[0]["html"])
|
||
return blocks
|
||
if status == "complete":
|
||
return await _revise(db, session_id, paper_id, text, client)
|
||
if status == "cancelled":
|
||
blocks = [_text_block("论文已取消。如需重新撰写，请告诉我新主题。")]
|
||
await _write_message(db, session_id, blocks=blocks, text=blocks[0]["html"])
|
||
return blocks
|
||
if status in ("outline", "writing"):
|
||
# 大纲/撰写阶段：修改意见 → 用原始 brief + 修改意见重新生成大纲（原大纲作废）
|
||
paper = await client.call_tool("get_paper", {"paper_id": paper_id})
|
||
ob = (paper.get("brief") or {}) if paper.get("ok") else {}
|
||
outline = await client.call_tool("propose_paper_outline", {
|
||
"topic": ob.get("topic") or paper.get("title") or text,
|
||
"requirements": "；".join(x for x in [ob.get("requirements"), text] if x),
|
||
"paper_type": ob.get("paper_type") or "",
|
||
"language": ob.get("language") or "",
|
||
"length": ob.get("length") or "",
|
||
"format_req": ob.get("format_req") or "",
|
||
"template_key": template_key,
|
||
})
|
||
if outline.get("ok"):
|
||
await client.call_tool("cancel_paper", {"paper_id": paper_id})
|
||
return await _outline_message(db, session_id, outline)
|
||
blocks = [_text_block("大纲修改失败，请重试。")]
|
||
await _write_message(db, session_id, blocks=blocks, text=blocks[0]["html"])
|
||
return blocks
|
||
blocks = [_text_block("论文撰写中，请稍候。")]
|
||
await _write_message(db, session_id, blocks=blocks, text=blocks[0]["html"])
|
||
return blocks
|
||
|
||
|
||
async def _revise(db, session_id: str, paper_id: str, text: str, client: PaperMcpClient) -> list[dict[str, Any]]:
|
||
plan = await client.call_tool("revise_paper", {"paper_id": paper_id, "instruction": text})
|
||
if not plan.get("ok"):
|
||
blocks = [_text_block(f"修订计划生成失败：{plan.get('message', '')}")]
|
||
await _write_message(db, session_id, blocks=blocks, text=blocks[0]["html"])
|
||
return blocks
|
||
affected = plan.get("affected_section_ids") or []
|
||
if not affected:
|
||
blocks = [_text_block("没有需要修改的章节。")]
|
||
await _write_message(db, session_id, blocks=blocks, text=blocks[0]["html"])
|
||
return blocks
|
||
|
||
for sid in affected:
|
||
res = await client.call_tool(
|
||
"revise_paper_section", {"paper_id": paper_id, "section_id": sid, "instruction": text}
|
||
)
|
||
if not res.get("ok"):
|
||
blocks = [_text_block(f"章节修订失败：{res.get('message', '')}")]
|
||
await _write_message(db, session_id, blocks=blocks, text="修订失败")
|
||
return blocks
|
||
|
||
comp = await client.call_tool("complete_paper", {"paper_id": paper_id})
|
||
paper = await client.call_tool("get_paper", {"paper_id": paper_id})
|
||
if comp.get("ok") and paper.get("ok"):
|
||
return await _artifact_message(db, session_id, paper, "已按你的意见修改，文献编号已重排。")
|
||
blocks = [_text_block("修改后重新生成失败，请重试。")]
|
||
await _write_message(db, session_id, blocks=blocks, text="修改失败")
|
||
return blocks
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 写节循环（confirm/resume 后台任务）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
async def run_writing_loop(db, paper_id: str, session_id: str) -> None:
|
||
"""后台逐节写节循环：写完→complete→get_paper→写 artifact→artifact 块。"""
|
||
from ...routers.docgen import insert_artifact # 惰性导入，避免服务层→路由层常驻耦合
|
||
|
||
client = _new_client()
|
||
await _write_message(db, session_id, blocks=[
|
||
{"type": "paper_progress", "paper_id": paper_id, "status": "writing", "sections": [], "current_section_id": ""},
|
||
], text="开始撰写")
|
||
try:
|
||
while True:
|
||
prog = await client.call_tool("get_paper_progress", {"paper_id": paper_id})
|
||
if not prog.get("ok"):
|
||
await _write_message(db, session_id, blocks=[_text_block("撰写进度查询失败。")], text="撰写失败")
|
||
return
|
||
if prog.get("status") == "cancelled":
|
||
await _write_message(db, session_id, blocks=[_text_block("论文已取消。")], text="已取消")
|
||
return
|
||
sections = prog.get("sections") or []
|
||
pending = [s for s in sections if s.get("status") != "done"]
|
||
if not pending:
|
||
break
|
||
sec = pending[0]
|
||
res = await client.call_tool("write_paper_section", {"paper_id": paper_id, "section_id": sec["section_id"]})
|
||
if not res.get("ok"):
|
||
await _write_message(db, session_id, blocks=[_text_block(
|
||
f"章节「{sec.get('title', '')}」撰写失败：{res.get('message', '')}"
|
||
)], text="撰写失败")
|
||
return
|
||
wc = len(res.get("markdown") or "")
|
||
await _write_message(db, session_id, blocks=[_text_block(
|
||
f"「{sec.get('title', '')}」已完成，约 {wc} 字。"
|
||
)], text=f"「{sec.get('title', '')}」已完成")
|
||
|
||
comp = await client.call_tool("complete_paper", {"paper_id": paper_id})
|
||
if not comp.get("ok"):
|
||
await _write_message(db, session_id, blocks=[_text_block(f"论文合并失败：{comp.get('message', '')}")], text="合并失败")
|
||
return
|
||
paper = await client.call_tool("get_paper", {"paper_id": paper_id})
|
||
if not paper.get("ok"):
|
||
await _write_message(db, session_id, blocks=[_text_block("论文拉取失败。")], text="拉取失败")
|
||
return
|
||
await _artifact_message(db, session_id, paper, "论文撰写完成！你可以在右侧文档页查看全文，或用自然语言提出修改意见。")
|
||
except PaperMcpError as exc:
|
||
logger.warning("paper: 写节循环失败：%s", exc)
|
||
await _write_message(db, session_id, blocks=[_text_block(f"论文服务异常：{exc}")], text="论文服务异常")
|
||
except Exception as exc: # noqa: BLE001
|
||
logger.exception("paper: 写节循环异常")
|
||
await _write_message(db, session_id, blocks=[_text_block(f"撰写失败：{exc}")], text="撰写失败")
|
||
|
||
|
||
def _build_full_markdown(paper: dict[str, Any]) -> str:
|
||
"""把 full_markdown + 文末参考文献列表 + disclaimer 拼成完整 markdown（供渲染 HTML 与落盘）。"""
|
||
md = paper.get("full_markdown", "")
|
||
refs = paper.get("references") or []
|
||
if refs:
|
||
md += "\n\n## 参考文献\n\n" + "\n".join(
|
||
f"[{r.get('n', '')}] {r.get('text', '')}".strip() for r in refs
|
||
)
|
||
disclaimer = (paper.get("disclaimer") or "").strip()
|
||
if disclaimer:
|
||
md += f"\n\n> {disclaimer}"
|
||
return md
|
||
|
||
|
||
async def _artifact_message(db, session_id: str, paper: dict[str, Any], nl: str) -> list[dict[str, Any]]:
|
||
from ...routers.docgen import insert_artifact # 惰性导入
|
||
from ..docgen import templating # 惰性导入：复用 docgen markdown→HTML 渲染链
|
||
|
||
refs = paper.get("references") or []
|
||
title = paper.get("title", "")
|
||
full_md = _build_full_markdown(paper)
|
||
html, css = templating.render_html(full_md, title=title, key=paper.get("template_key") or "document")
|
||
artifact_id = await insert_artifact(
|
||
db, session_id=session_id, title=title,
|
||
markdown=full_md, template_name="论文", template_key=paper.get("template_key") or "document",
|
||
template_css=css, html=html,
|
||
)
|
||
blocks = [
|
||
_text_block(nl),
|
||
{"type": "artifact", "artifact_id": artifact_id, "title": title,
|
||
"template_name": "论文", "status": "done", "paper_id": paper.get("paper_id", ""),
|
||
"citation_count": len(refs)},
|
||
]
|
||
await _write_message(db, session_id, blocks=blocks, text=nl)
|
||
return blocks
|