paper-doc/backend/app/mcp/support.py

"""Shared plumbing for the MCP tools: sessions, JSON, lookups, and headings.

Every tool in :mod:`app.mcp.tools` is three steps — open a session, call the
same :mod:`app.crud` helper the REST API calls, serialise the result into a
compact JSON string. This module holds the shared half of that.

Two things here are worth knowing before reading a tool:

**Why the results are hand-shaped rather than schema dumps.** A tool result is
paid for in context tokens, and ``PaperRead.model_dump()`` carries
``created_at``, ``updated_at`` and four count columns through every list row.
The ``*_row`` helpers below pick what a model actually reads, and
:func:`clean` drops every ``None`` on the way out.

**Why a paragraph can be addressed by heading.** ``paper_sentence`` is stored
against a *position* (``paper_template_filed_sort``), which is correct for the
document and useless for a model: nobody writing "1. Introduction" knows that
the template happens to place it at ``sort = 20``. :func:`resolve_position`
accepts either, and matches a heading through the numbering, the level and a
unique substring so ``"Introduction"``, ``"1. Introduction"`` and
``"1 Introduction"`` all land in the same paragraph.

Text helpers
------------
:func:`split_text` is the one place this project ever cuts text, and it exists
only because a model arrives with a block of prose rather than with an array of
sentences. It is opt-in: the write tools default to ``line`` (one line, one
sentence — the project's own rule), and ``sentence`` refuses to cut at a
decimal point or after ``et al.``. The editor never uses it, so no stored text
is ever re-split.
"""

from __future__ import annotations

import json
import re
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from datetime import date, datetime
from decimal import Decimal
from typing import Any, NoReturn

from mcp.server.mcpserver.exceptions import ToolError
from pydantic import BaseModel
from sqlalchemy.orm import Session

from app.crud import paper as paper_crud
from app.crud import template as template_crud
from app.db.session import SessionLocal
from app.models import (
    Paper,
    PaperSentence,
    Template,
    TemplateFieldLibrary,
)
from app.schemas.paper import (
    CitationInput,
    PaperDocumentRead,
    ParagraphRead,
    SentenceInput,
    fold_whitespace,
)

#: Upper bound on a page a tool may ask for. The REST API caps at 200 as well;
#: a tool result is a context window, so it is not raised here.
MAX_PAGE_SIZE = 200


# --- sessions and errors -----------------------------------------------------


@contextmanager
def session() -> Iterator[Session]:
    """One database session per tool call, always closed.

    A tool is a whole request: there is no FastAPI dependency to hang a session
    on, and leaving one open across calls would keep a TiDB connection pinned
    for as long as the client stays alive.
    """
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def fail(message: str) -> NoReturn:
    """Refuse a call with a message the model can act on.

    :class:`~mcp.server.mcpserver.exceptions.ToolError` becomes
    ``isError: true`` with this text as the content, which the model reads and
    retries against — the opposite of an unhandled exception, whose message
    never reaches it.
    """
    raise ToolError(message)


# --- JSON shaping ------------------------------------------------------------


def _jsonable(value: Any) -> Any:
    """Render the two types SQLAlchemy hands back that JSON does not have."""
    if isinstance(value, (datetime, date)):
        # Seconds are enough for a list row, and microseconds are six tokens.
        return value.strftime("%Y-%m-%d %H:%M:%S")
    if isinstance(value, Decimal):
        return float(value)
    raise TypeError(f"不能序列化的类型：{type(value).__name__}")


def clean(value: Any) -> Any:
    """Recursively drop ``None`` from objects, so a result carries only what is set.

    Lists and the empty string survive: an empty ``sentences`` list means "this
    paragraph is empty", which is a fact, while ``author: null`` means "no
    author", which the absence of the key says for free.
    """
    if isinstance(value, dict):
        return {
            key: clean(item) for key, item in value.items() if item is not None
        }
    if isinstance(value, list):
        return [clean(item) for item in value]
    return value


def dumps(value: Any) -> str:
    """Serialise a tool result: compact, unescaped, ``None``-free."""
    return json.dumps(
        clean(value),
        ensure_ascii=False,
        default=_jsonable,
        separators=(",", ":"),
    )


def page_args(page: int, page_size: int) -> tuple[int, int]:
    """Clamp a page request to something a context window can hold."""
    return max(1, int(page)), min(max(1, int(page_size)), MAX_PAGE_SIZE)


def page_payload(*, items: list[Any], total: int, page: int, page_size: int) -> dict:
    """The list envelope every list tool returns, mirroring ``PageResult``."""
    return {
        "items": items,
        "total": total,
        "page": page,
        "page_size": page_size,
        "pages": (total + page_size - 1) // page_size if page_size else 0,
    }


# --- lookups that explain themselves -----------------------------------------


def paper_or_fail(db: Session, paper_id: int) -> Paper:
    """Return one paper, or refuse with the id that was not found."""
    paper = paper_crud.get(db, paper_id)
    if paper is None:
        fail(f"论文 {paper_id} 不存在")
    return paper


def template_or_fail(db: Session, template_id: int) -> Template:
    """Return one template, or refuse with the id that was not found."""
    template = template_crud.get(db, template_id)
    if template is None:
        fail(f"模板 {template_id} 不存在")
    return template


def field_or_fail(db: Session, field_id: int) -> TemplateFieldLibrary:
    """Return one library field, or refuse with the id that was not found."""
    from app.crud import template_field_library as field_crud

    field = field_crud.get(db, field_id)
    if field is None:
        fail(f"字段 {field_id} 不存在")
    return field


def sentence_or_fail(paper: Paper, sentence_id: int) -> PaperSentence:
    """Find a sentence *within* one paper.

    Scoped to the paper on purpose, exactly as the REST route is: an id from
    another paper must not be editable through this one.
    """
    for sentence in paper.sentences:
        if sentence.id == sentence_id:
            return sentence
    fail(f"句子 {sentence_id} 不属于论文 {paper.id}")
    raise AssertionError  # unreachable; keeps type checkers honest


def id_list(value: int | Sequence[int] | None) -> list[int]:
    """Accept ``3`` or ``[3, 4]`` as the same request.

    A model told to "delete these two" should not have to discover a separate
    batch tool, and a batch of one is the common case anyway.
    """
    if value is None:
        return []
    if isinstance(value, bool):
        fail("ids 必须是整数或整数数组")
    if isinstance(value, int):
        return [value]
    return [int(item) for item in value]


# --- rows --------------------------------------------------------------------


def _short_moment(value: datetime | None) -> str | None:
    return value.strftime("%Y-%m-%d %H:%M") if value is not None else None


def paper_row(paper: Paper) -> dict:
    """One paper as a list row or a summary.

    The counts come from :func:`app.crud.paper.read` rather than from a second
    count here, so a tool result and a page of the UI can never disagree about
    how much has been written.
    """
    read = paper_crud.read(paper)
    return {
        "id": read.id,
        "title": read.title,
        "template_id": read.template_id,
        "template": read.template_name,
        "author": read.author,
        "status": read.status,
        "keywords": read.keywords,
        "target_journal": read.target_journal,
        # 已写段落 / 模板定义段落 — the progress read-out the paper list shows.
        "paragraphs": read.paragraph_count,
        "template_paragraphs": read.template_paragraph_count,
        "sentences": read.sentence_count,
        "updated_at": _short_moment(read.updated_at),
    }


def citation_row(citation: Any) -> dict:
    """One citation, as the model reads it."""
    return {"reference_id": citation.reference_id, "quote": citation.quote}


def sentence_row(sentence: Any, *, separator: str | None = None) -> dict:
    """One sentence with its citations.

    ``separator_before`` is passed through when the caller assembled a
    paragraph, because it is derived from the neighbouring sentence and cannot
    be recomputed from this row alone.
    """
    row = {
        "id": sentence.id,
        "position": sentence.paper_template_filed_sort,
        "sort": sentence.sort,
        "content": sentence.content,
        "citations": [citation_row(item) for item in sentence.citations],
    }
    if separator is not None:
        row["separator_before"] = separator
    return row


def paragraph_row(paragraph: ParagraphRead, *, preview: int = 80) -> dict:
    """One paragraph as an outline entry: where it is and how much is in it."""
    text = paragraph_text(paragraph)
    return {
        "position": paragraph.paper_template_filed_sort,
        "heading": paragraph.name,
        "level": paragraph.level,
        "matched": paragraph.matched,
        "sentences": len(paragraph.sentences),
        "chars": len(text),
        "preview": text[:preview] if preview else None,
    }


def paragraph_text(paragraph: ParagraphRead, *, citations: bool = False) -> str:
    """Print one paragraph exactly as the document renders it.

    ``separator_before`` is added, not invented: the server decided what
    belongs at each seam, so a tool result and the paper view agree on where
    the spaces go. Nothing here adds spacing of its own.
    """
    parts: list[str] = []
    for sentence in paragraph.sentences:
        parts.append(sentence.separator_before + sentence.content)
        if citations and sentence.citations:
            markers = ",".join(str(index) for index in range(1, len(sentence.citations) + 1))
            parts.append(f"[{markers}]")
    return "".join(parts)


def document_text(
    document: PaperDocumentRead,
    *,
    include_empty: bool = False,
    citations: bool = True,
) -> str:
    """Render a whole paper as the text a reader would see.

    Headings come from the template, paragraphs from the content, and
    positions the template does not define are printed under 未设定 — the same
    three rules the paper view renders by, applied here so a single tool call
    can hand a model the finished manuscript.
    """
    lines: list[str] = []
    for paragraph in document.paragraphs:
        body = paragraph_text(paragraph, citations=citations)
        if not body and not include_empty:
            continue
        lines.append(paragraph.name if paragraph.matched and paragraph.name else "未设定")
        lines.append(body)
        lines.append("")

    if citations and document.citations:
        lines.append("参考文献")
        for citation in document.citations:
            reference = (
                f"[{citation.reference_id}]" if citation.reference_id else ""
            )
            lines.append(f"[{citation.index}]{reference} {citation.quote}")
        lines.append("")

    text = "\n".join(lines).rstrip("\n")
    if document.warnings:
        text += "\n\n（提示：" + " ".join(document.warnings) + "）"
    return text


# --- headings ----------------------------------------------------------------

#: A leading outline number: ``1``, ``1.``, ``2.3.1``, ``3、``, ``4)``.
_NUMBER_PREFIX = re.compile(r"^\s*\d+(?:\.\d+)*\s*[.、)）]?\s*")


def _heading_key(text: str) -> str:
    """A comparison key for a heading: no markup, no number, no space, no case.

    Removing the numbering is what lets ``"Introduction"`` find
    ``"1. Introduction"``. Removing the space is what lets a line the model
    wrote as ``"1.Introduction"`` still find it.
    """
    stripped = text.strip().lstrip("#").strip()
    stripped = _NUMBER_PREFIX.sub("", stripped)
    return re.sub(r"[\s\u3000]+", "", stripped).casefold()


def heading_index(paper: Paper) -> list[dict]:
    """The paper's outline: every position its template defines.

    Ties are collapsed the way :func:`app.crud.paper.build_document` collapses
    them — the first placement at a position renders the heading — so this list
    and the document it describes cannot disagree about what a position says.
    """
    if paper.template is None:
        return []

    seen: dict[int, dict] = {}
    for placement in paper.template.items:
        if placement.sort in seen:
            continue
        seen[placement.sort] = {
            "position": placement.sort,
            "heading": placement.field.name,
            "level": placement.field.level,
            "field_id": placement.field_id,
        }
    return [seen[position] for position in sorted(seen)]


def resolve_position(
    paper: Paper,
    *,
    position: int | None = None,
    heading: str | None = None,
) -> int:
    """Turn "where the model wants to write" into a paragraph position.

    A position is taken as given — including one the template does not define,
    which is exactly the case the renderer already handles as 未设定. A heading
    is matched against the paper's template and refused with the list of what
    the template actually says, because a model that guessed a heading will
    guess again correctly once it can see the real ones.
    """
    if position is not None:
        return int(position)
    if heading is None or not heading.strip():
        fail("必须给出 position 或 heading，两者至少要有一个")

    outline = heading_index(paper)
    if not outline:
        fail(
            f"论文 {paper.id} 还没有模板，无法按标题定位；"
            "请改用 position，或先用 paper_update 给论文设置 template_id"
        )

    query = _heading_key(heading)
    exact = [item for item in outline if _heading_key(item["heading"]) == query]
    if len(exact) == 1:
        return int(exact[0]["position"])
    if len(exact) > 1:
        fail(_ambiguous(heading, exact))

    partial = [item for item in outline if query and query in _heading_key(item["heading"])]
    if len(partial) == 1:
        return int(partial[0]["position"])
    if len(partial) > 1:
        fail(_ambiguous(heading, partial))

    available = "、".join(f"{item['heading']}({item['position']})" for item in outline)
    fail(f"模板里没有标题“{heading}”。可用标题：{available}")


def _ambiguous(heading: str, matches: list[dict]) -> str:
    listed = "、".join(f"{item['heading']}({item['position']})" for item in matches)
    return f"标题“{heading}”匹配到多个段落（{listed}），请写完整标题或直接用 position"


# --- text --------------------------------------------------------------------

#: Where a sentence may end in CJK text. ``；`` is deliberately absent: a
#: semicolon joins two halves of one argument often enough that cutting there
#: would produce fragments rather than sentences.
_CJK_TERMINATORS = "。！？…"

#: Trailing marks that belong to the sentence they follow.
_TRAILING = "”’」』）)】》…"

#: Latin abbreviations whose full stop is not a sentence end.
_ABBREVIATIONS = frozenset(
    {
        "al.",
        "approx.",
        "ca.",
        "cf.",
        "dr.",
        "e.g.",
        "etc.",
        "fig.",
        "figs.",
        "i.e.",
        "mr.",
        "mrs.",
        "no.",
        "prof.",
        "vs.",
    }
)

SPLIT_MODES = ("line", "sentence", "paragraph")


def split_text(text: str, mode: str = "line") -> list[str]:
    """Cut a block of prose into sentences — the only text cutting in the project.

    Three modes, in increasing order of how much the tool is trusted:

    ``line``
        Each non-blank line is one sentence. The default, and the project's own
        rule ("one line in the editor is one sentence"): nothing is guessed, so
        nothing can be guessed wrong.
    ``paragraph``
        The whole block folds to a single sentence. For a writer who puts a
        paragraph on one line and means it as one unit.
    ``sentence``
        Best-effort splitting after ``。！？`` and after a Latin ``.`` ``!``
        ``?`` that is followed by a space and something that starts a sentence.
        It does not cut at a decimal point or after ``et al.``, but it is still
        a guess — which is why it is never the default and never applied to
        text that is already stored.

    Blank lines are dropped in every mode: the tool is writing, and an empty
    sentence is only meaningful when a person puts one there on purpose.
    """
    if mode not in SPLIT_MODES:
        fail(f"split 只能是 {'/'.join(SPLIT_MODES)}，收到的是“{mode}”")

    if mode == "paragraph":
        folded = fold_whitespace(text)
        return [folded] if folded else []

    if mode == "line":
        return [line.strip() for line in text.splitlines() if line.strip()]

    return _split_sentences(fold_whitespace(text))


def _split_sentences(text: str) -> list[str]:
    """The ``sentence`` mode of :func:`split_text`, on already-folded text."""
    if not text:
        return []

    out: list[str] = []
    start = 0
    index = 0

    while index < len(text):
        char = text[index]

        if char in _CJK_TERMINATORS:
            end = index + 1
            while end < len(text) and text[end] in _TRAILING:
                end += 1
            out.append(text[start:end])
            start = end
            index = end
            continue

        if char in ".!?" and _is_latin_break(text, index):
            out.append(text[start : index + 1])
            start = index + 1

        index += 1

    tail = text[start:]
    if tail:
        out.append(tail)

    return [item.strip() for item in out if item.strip()]


def _is_latin_break(text: str, index: int) -> bool:
    """Whether the ``.``/``!``/``?`` at ``index`` ends a sentence."""
    # A boundary needs a space (or the end) after it, or "3.14" and "www.x.com"
    # would each become two sentences.
    after = index + 1
    if after < len(text) and not text[after].isspace():
        return False

    follower = after
    while follower < len(text) and text[follower].isspace():
        follower += 1
    if follower >= len(text):
        return True

    # The next sentence starts with a capital, a digit, a quote or CJK; a
    # lowercase continuation means this was an abbreviation inside a sentence.
    nxt = text[follower]
    if not (nxt.isupper() or nxt.isdigit() or nxt in "\"'“（([" or ord(nxt) > 0x2E80):
        return False

    token_start = index
    while token_start > 0 and not text[token_start - 1].isspace():
        token_start -= 1
    token = text[token_start : index + 1]

    if token.casefold() in _ABBREVIATIONS:
        return False
    # A lone letter is an initial — "J. Smith" is one name, not two sentences.
    if len(token) == 2 and token[0].isupper():
        return False

    return True


def sentence_inputs(items: Sequence[Any], *, default_split: str = "line") -> list[SentenceInput]:
    """Normalise what a model sent into the project's own sentence payload.

    One item may arrive in any of the shapes a model naturally produces:

    * ``"A sentence."`` — the common case;
    * ``{"content": "...", "citations": [{"quote": "...", "reference_id": 3}]}``;
    * ``{"text": "one\\ntwo", "split": "line"}`` — a block to be cut.

    Every path ends in :class:`~app.schemas.paper.SentenceInput`, so the
    whitespace folding and the "a citation must quote something" rule are the
    ones the REST API enforces, not a second copy of them.
    """
    out: list[SentenceInput] = []

    for item in items:
        if isinstance(item, str):
            out.append(SentenceInput(content=item))
            continue

        # A tool signature declares the union as pydantic models, so the same
        # path has to accept both an already-validated model and the raw dict
        # that a lower-level caller passes.
        if isinstance(item, BaseModel):
            item = item.model_dump()

        if not isinstance(item, dict):
            fail(f"句子只能是字符串或对象，收到的是 {type(item).__name__}")

        citations = [
            CitationInput(
                quote=str(citation.get("quote", "")),
                reference_id=citation.get("reference_id"),
            )
            for citation in (item.get("citations") or [])
        ]

        if item.get("content") is not None:
            out.append(SentenceInput(content=str(item["content"]), citations=citations))
            continue

        if item.get("text") is not None:
            mode = item.get("split") or default_split
            for piece in split_text(str(item["text"]), mode):
                out.append(SentenceInput(content=piece, citations=citations))
            continue

        fail("句子对象需要有 content 或 text 字段")

    return out


def paragraph_inputs(spec: Any) -> list[SentenceInput]:
    """Collect one paragraph's sentences from every shape a model may send.

    A caller may hand over an explicit ``sentences`` list, a single ``text``
    block to be cut, or neither. "Neither" is refused rather than treated as an
    empty paragraph: clearing a paragraph is a different intention with its own
    tool, and silently emptying one because an argument was misspelled is the
    one mistake this layer must not make.
    """
    items: list[Any] = list(spec.sentences or [])
    if not items and spec.text is not None and spec.text.strip():
        items = [{"text": spec.text, "split": spec.split}]

    if not items:
        fail(
            f"段落（position={spec.position} heading={spec.heading}）既没有 sentences "
            "也没有 text。若确实要清空这段，请用 paragraph_delete"
        )

    extra = [
        CitationInput(quote=item.quote, reference_id=item.reference_id)
        for item in (spec.citations or [])
    ]
    inputs = sentence_inputs(items)
    if extra:
        # Paragraph-level citations apply to every sentence that does not carry
        # its own — the shape a model produces when a whole paragraph rests on
        # one source.
        for item in inputs:
            if not item.citations:
                item.citations = list(extra)
    return inputs


__all__ = [
    "MAX_PAGE_SIZE",
    "SPLIT_MODES",
    "clean",
    "citation_row",
    "document_text",
    "dumps",
    "fail",
    "field_or_fail",
    "heading_index",
    "id_list",
    "page_args",
    "page_payload",
    "paper_or_fail",
    "paper_row",
    "paragraph_inputs",
    "paragraph_row",
    "paragraph_text",
    "resolve_position",
    "sentence_inputs",
    "sentence_or_fail",
    "sentence_row",
    "session",
    "split_text",
    "template_or_fail",
]