110 lines
5.1 KiB
Python
110 lines
5.1 KiB
Python
"""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
|
|
]
|
|
brief = loads_json(paper.get("brief_json"), {})
|
|
return {
|
|
"ok": True, "paper_id": paper_id, "title": paper["title"], "status": paper["status"],
|
|
"template_key": brief.get("template_key") or "document",
|
|
"brief": brief,
|
|
"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}
|