"""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", ]