PlantAI-rag-new/paper_writer_service/app/services/paper/write.py

"""write.py —— 章节撰写与修订（write_paper_section / revise_paper_section / revise_paper）。"""
from __future__ import annotations

from typing import Any

from ...db import db, dumps_json, loads_json, now_iso
from . import llm
from .cite import find_cite_numbers, render_gbt7714

_WRITE_SYSTEM = (
    "你是植物学论文撰写助手，正在撰写论文的某一章节。"
    "引用文献一律用 [n] 指代，n 必须取自下方提供的全局编号列表；如需引用新文献，编号按当前最大编号依次 +1。"
    "禁止杜撰 DOI/PMID/作者/期刊——只能使用编号列表或知识库线索中真实给出的元数据。"
    "只输出一个合法 JSON 对象：{\"markdown\": 章节正文(不含章标题), \"citations_used\": "
    "[{\"temp_n\": 编号, \"title\": \"...\", \"authors\": \"...\", \"journal\": \"...\", "
    "\"year\": \"...\", \"doi\": \"...\", \"pmid\": \"...\", \"source_doc_id\": \"...\", \"snippet\": \"...\"}]}。"
    "citations_used 必须列出本节正文中出现的每一个 [n] 引用（既有编号也列出其已有元数据）。"
)

_REVISE_SYSTEM = (
    "你是植物学论文撰写助手，正在按用户意见修订论文的某一章节。"
    "保持引用协议不变：[n] 指代，n 只能取自全局编号列表或最大编号 +1；禁止杜撰元数据。"
    "只输出一个合法 JSON 对象：{\"markdown\": 修订后的章节正文, \"citations_used\": [...]}。"
)


# ---------------------------------------------------------------------------
# 读取辅助
# ---------------------------------------------------------------------------

async def _get_paper(paper_id: str) -> dict[str, Any] | None:
    row = await db.fetchone("SELECT * FROM papers WHERE id = ?", (paper_id,))
    return dict(row) if row else None


async def _get_sections(paper_id: str) -> list[dict[str, Any]]:
    rows = await db.fetchall(
        "SELECT * FROM paper_sections WHERE paper_id = ? ORDER BY rowid", (paper_id,)
    )
    return [dict(r) for r in rows]


async def _get_citations(paper_id: str) -> list[dict[str, Any]]:
    rows = await db.fetchall(
        "SELECT * FROM citations WHERE paper_id = ? ORDER BY temp_n", (paper_id,)
    )
    return [dict(r) for r in rows]


def _global_citations_text(citations: list[dict[str, Any]]) -> str:
    if not citations:
        return "（当前无已编号文献；新引用从 [1] 开始）"
    return "\n".join(f"[{c['temp_n']}] {c['ref_text']}" for c in citations)


def _max_temp_n(citations: list[dict[str, Any]]) -> int:
    return max((c.get("temp_n") or 0 for c in citations), default=0)


def _build_write_user(paper: dict, outline: dict, sec: dict, cit_text: str, style: str, cite_mode: str) -> str:
    parts: list[str] = []
    parts.append(f"论文标题：{paper.get('title', '')}")
    parts.append(f"论文类型：{paper.get('paper_type') or '综述'}")
    parts.append(f"语言：{paper.get('language') or '中文'}")
    parts.append(f"格式要求：{paper.get('format_req') or '无'}")
    parts.append(f"写作风格：{style or paper.get('style') or '学术严谨'}")
    parts.append("")
    parts.append("论文大纲：")
    for s in outline.get("sections") or []:
        mark = "（当前撰写节）" if s.get("section_id") == sec.get("section_id") else ""
        parts.append(f"- {s.get('section_id')}: {s.get('title')}{mark} 要点：{'；'.join(s.get('key_points') or [])}")
    parts.append("")
    parts.append(f"现在撰写章节「{sec.get('title', '')}」（section_id={sec.get('section_id')}）。")
    if sec.get("key_points"):
        parts.append("本节要点：" + "；".join(loads_json(sec.get("key_points"), [])))
    parts.append("")
    parts.append("全局已编号文献列表（[n] 引用只能从这里取值，或从最大编号 +1 起新增）：")
    parts.append(cit_text)
    return "\n".join(parts)


def _validate_numbers(markdown: str, valid: set[int]) -> list[int]:
    return [n for n in find_cite_numbers(markdown) if n not in valid]


def _kind_of(meta: dict[str, Any]) -> str:
    if meta.get("doi"):
        return "doi"
    if meta.get("pmid"):
        return "pmid"
    return "title"


async def _persist_new_citations(paper_id: str, section_id: str, used: list[Any], existing_nums: set[int]) -> list[dict[str, Any]]:
    """把 citations_used 中的新引用（temp_n 不在既有集合）落库，返回新增摘要。"""
    new_citations: list[dict[str, Any]] = []
    seen_new: set[int] = set()
    for c in used:
        if not isinstance(c, dict):
            continue
        tn = int(c.get("temp_n") or 0)
        if not tn or tn in existing_nums or tn in seen_new:
            continue
        seen_new.add(tn)
        ref_text = render_gbt7714(c)
        await db.execute(
            "INSERT INTO citations (paper_id, section_id, temp_n, kind, ref_text, source_doc_id, snippet, created_at) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            (
                paper_id, section_id, tn, _kind_of(c), ref_text,
                c.get("source_doc_id") or "", c.get("snippet") or "", now_iso(),
            ),
        )
        new_citations.append({"temp_n": tn, "ref_text": ref_text})
    return new_citations


# ---------------------------------------------------------------------------
# 写节
# ---------------------------------------------------------------------------

async def write_paper_section(paper_id: str, section_id: str, style: str = "", cite_mode: str = "") -> dict[str, Any]:
    paper = await _get_paper(paper_id)
    if not paper:
        return {"ok": False, "code": "NOT_FOUND", "message": "论文不存在"}
    sections = await _get_sections(paper_id)
    sec = next((s for s in sections if s["section_id"] == section_id), None)
    if not sec:
        return {"ok": False, "code": "NOT_FOUND", "message": "章节不存在"}
    citations = await _get_citations(paper_id)
    outline = loads_json(paper.get("outline_json"), {})
    cit_text = _global_citations_text(citations)
    existing_nums = {c.get("temp_n") for c in citations}

    user = _build_write_user(paper, outline, sec, cit_text, style, cite_mode)
    markdown = ""
    used: list[Any] = []
    bad: list[int] = []
    for attempt in range(2):
        try:
            res = await llm.complete_json(_WRITE_SYSTEM, user, max_tokens=8000)
        except llm.LLMError as exc:
            return {"ok": False, "code": "LLM_ERROR", "message": f"写节 LLM 失败：{exc}"}
        markdown = res.get("markdown") or ""
        used = res.get("citations_used") or []
        used_nums = {int(c.get("temp_n") or 0) for c in used if isinstance(c, dict)}
        valid = existing_nums | used_nums
        bad = _validate_numbers(markdown, valid)
        if not bad:
            break
        user += f"\n（上次正文出现的引用编号 {sorted(bad)} 非法；只能使用已有编号或从 {_max_temp_n(citations) + 1} 起的新编号，重试。）"
    else:
        return {"ok": False, "code": "CITE_VALIDATION_FAILED", "message": f"引用编号非法：{sorted(bad)}"}

    new_citations = await _persist_new_citations(paper_id, section_id, used, existing_nums)
    await db.execute(
        "UPDATE paper_sections SET markdown = ?, word_count = ?, status = 'done', updated_at = ? WHERE paper_id = ? AND section_id = ?",
        (markdown, len(markdown), now_iso(), paper_id, section_id),
    )
    await db.execute("UPDATE papers SET status = 'writing', updated_at = ? WHERE id = ?", (now_iso(), paper_id))
    return {"ok": True, "section_id": section_id, "status": "done", "markdown": markdown,
            "citations_used": used, "new_citations": new_citations}


# ---------------------------------------------------------------------------
# 修订
# ---------------------------------------------------------------------------

async def revise_paper_section(paper_id: str, section_id: str, instruction: str) -> dict[str, Any]:
    paper = await _get_paper(paper_id)
    if not paper:
        return {"ok": False, "code": "NOT_FOUND", "message": "论文不存在"}
    sections = await _get_sections(paper_id)
    sec = next((s for s in sections if s["section_id"] == section_id), None)
    if not sec:
        return {"ok": False, "code": "NOT_FOUND", "message": "章节不存在"}
    citations = await _get_citations(paper_id)
    outline = loads_json(paper.get("outline_json"), {})
    existing_nums = {c.get("temp_n") for c in citations}

    user = _build_write_user(paper, outline, sec, _global_citations_text(citations), "", "")
    user += "\n\n原章节正文：\n" + (sec.get("markdown") or "")
    user += f"\n\n修订要求：{instruction}"

    markdown = ""
    used: list[Any] = []
    bad: list[int] = []
    for attempt in range(2):
        try:
            res = await llm.complete_json(_REVISE_SYSTEM, user, max_tokens=8000)
        except llm.LLMError as exc:
            return {"ok": False, "code": "LLM_ERROR", "message": f"修订 LLM 失败：{exc}"}
        markdown = res.get("markdown") or ""
        used = res.get("citations_used") or []
        used_nums = {int(c.get("temp_n") or 0) for c in used if isinstance(c, dict)}
        valid = existing_nums | used_nums
        bad = _validate_numbers(markdown, valid)
        if not bad:
            break
        user += f"\n（引用编号 {sorted(bad)} 非法，重试。）"
    else:
        return {"ok": False, "code": "CITE_VALIDATION_FAILED", "message": f"引用编号非法：{sorted(bad)}"}

    new_citations = await _persist_new_citations(paper_id, section_id, used, existing_nums)
    await db.execute(
        "UPDATE paper_sections SET markdown = ?, word_count = ?, status = 'done', updated_at = ? WHERE paper_id = ? AND section_id = ?",
        (markdown, len(markdown), now_iso(), paper_id, section_id),
    )
    return {"ok": True, "section_id": section_id, "markdown": markdown,
            "citations_used": used, "citation_delta": len(new_citations), "new_citations": new_citations}


async def revise_paper(paper_id: str, instruction: str) -> dict[str, Any]:
    """LLM 判定受影响章节，返回受影响节清单 + 修订计划（由调用方逐节 revise_paper_section）。"""
    paper = await _get_paper(paper_id)
    if not paper:
        return {"ok": False, "code": "NOT_FOUND", "message": "论文不存在"}
    sections = await _get_sections(paper_id)
    outline = loads_json(paper.get("outline_json"), {})

    system = (
        "你是论文修订计划助手。根据用户修订意见与论文大纲，判定需要修改哪些章节。"
        "只输出一个合法 JSON 对象：{\"affected_section_ids\": [\"s1\", ...], \"plan\": 修订计划(一句话)}。"
    )
    sec_list = "\n".join(f"{s['section_id']}: {s['title']}" for s in sections)
    user = f"论文标题：{paper.get('title')}\n大纲章节：\n{sec_list}\n\n修订意见：{instruction}"
    try:
        res = await llm.complete_json(system, user, max_tokens=2000)
    except llm.LLMError as exc:
        return {"ok": False, "code": "LLM_ERROR", "message": f"修订计划 LLM 失败：{exc}"}
    affected = [s for s in (res.get("affected_section_ids") or []) if any(x["section_id"] == s for x in sections)]
    return {"ok": True, "affected_section_ids": affected, "plan": res.get("plan", "")}