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

"""complete.py —— get_paper / get_paper_progress / complete_paper / cancel_paper。"""
from __future__ import annotations

from typing import Any

from ...db import db, dumps_json, loads_json, now_iso
from .cite import render_disclaimer, renumber_map, replace_numbers


async def _load(paper_id: str) -> dict[str, Any]:
    paper = await db.fetchone("SELECT * FROM papers WHERE id = ?", (paper_id,))
    if not paper:
        return {"ok": False, "code": "NOT_FOUND", "message": "论文不存在"}
    sections = [dict(r) for r in await db.fetchall(
        "SELECT * FROM paper_sections WHERE paper_id = ? ORDER BY rowid", (paper_id,))]
    citations = [dict(r) for r in await db.fetchall(
        "SELECT * FROM citations WHERE paper_id = ? ORDER BY temp_n", (paper_id,))]
    return {"ok": True, "paper": dict(paper), "sections": sections, "citations": citations}


def _merge_sections(paper: dict[str, Any], sections: list[dict[str, Any]]) -> str:
    parts = [f"# {paper.get('title', '')}"]
    for s in sections:
        parts.append(f"\n## {s['title']}\n\n{s.get('markdown', '')}")
    return "\n".join(parts)


async def get_paper(paper_id: str, format: str = "md") -> dict[str, Any]:
    p = await _load(paper_id)
    if not p.get("ok"):
        return p
    paper, sections, citations = p["paper"], p["sections"], p["citations"]
    full_markdown = paper.get("final_markdown") or _merge_sections(paper, sections)
    references = loads_json(paper.get("references_json"), []) or [
        {"n": c.get("final_n") or c.get("temp_n"), "text": c["ref_text"], "kind": c["kind"],
         "source_doc_id": c["source_doc_id"]} for c in citations
    ]
    return {
        "ok": True, "paper_id": paper_id, "title": paper["title"], "status": paper["status"],
        "sections": [
            {"section_id": s["section_id"], "title": s["title"], "status": s["status"],
             "word_count": s["word_count"]} for s in sections
        ],
        "full_markdown": full_markdown, "references": references,
        "disclaimer": paper.get("disclaimer", ""),
    }


async def get_paper_progress(paper_id: str) -> dict[str, Any]:
    paper = await db.fetchone("SELECT * FROM papers WHERE id = ?", (paper_id,))
    if not paper:
        return {"ok": False, "code": "NOT_FOUND", "message": "论文不存在"}
    paper = dict(paper)
    sections = [dict(r) for r in await db.fetchall(
        "SELECT * FROM paper_sections WHERE paper_id = ? ORDER BY rowid", (paper_id,))]
    cit_rows = [dict(r) for r in await db.fetchall(
        "SELECT section_id, COUNT(*) AS c FROM citations WHERE paper_id = ? GROUP BY section_id", (paper_id,))]
    cit_by_sec = {r["section_id"]: r["c"] for r in cit_rows}
    current = next((s["section_id"] for s in sections if s["status"] != "done"), None)
    return {
        "ok": True, "paper_id": paper_id, "status": paper["status"], "current_section_id": current,
        "sections": [
            {"section_id": s["section_id"], "title": s["title"], "status": s["status"],
             "word_count": s["word_count"], "citation_count": cit_by_sec.get(s["section_id"], 0)}
            for s in sections
        ],
    }


async def complete_paper(paper_id: str) -> dict[str, Any]:
    p = await _load(paper_id)
    if not p.get("ok"):
        return p
    paper, sections, citations = p["paper"], p["sections"], p["citations"]
    full = _merge_sections(paper, sections)
    temp_numbers = [c.get("temp_n") for c in citations if c.get("temp_n")]
    mapping = renumber_map(temp_numbers, full)

    final_full = replace_numbers(full, mapping)
    for c in citations:
        await db.execute("UPDATE citations SET final_n = ? WHERE id = ?", (mapping.get(c.get("temp_n")), c["id"]))
    for s in sections:
        await db.execute("UPDATE paper_sections SET markdown = ? WHERE id = ?",
                         (replace_numbers(s.get("markdown"), mapping), s["id"]))

    ordered = sorted(citations, key=lambda c: mapping.get(c.get("temp_n"), 0))
    references = [
        {"n": mapping.get(c.get("temp_n"), c.get("temp_n")), "text": c["ref_text"],
         "kind": c["kind"], "source_doc_id": c["source_doc_id"]} for c in ordered
    ]
    has_degraded = any(c["kind"] == "title" for c in citations)
    disclaimer = render_disclaimer() if has_degraded else ""
    await db.execute(
        "UPDATE papers SET status = 'complete', final_markdown = ?, references_json = ?, disclaimer = ?, updated_at = ? WHERE id = ?",
        (final_full, dumps_json(references), disclaimer, now_iso(), paper_id),
    )
    return {"ok": True, "full_markdown": final_full, "references": references,
            "renumber_map": mapping, "disclaimer": disclaimer}


async def cancel_paper(paper_id: str) -> dict[str, Any]:
    paper = await db.fetchone("SELECT id FROM papers WHERE id = ?", (paper_id,))
    if not paper:
        return {"ok": False, "code": "NOT_FOUND", "message": "论文不存在"}
    await db.execute("UPDATE papers SET status = 'cancelled', updated_at = ? WHERE id = ?", (now_iso(), paper_id))
    return {"ok": True}