PlantAI-rag-new/sdk/app/schemas.py

"""schemas.py —— Pydantic v2 模型（API.md 的契约形态）。

覆盖：块联合（§2）、SourceItem（§8）、GraphData / GraphNode / GraphEdge /
GraphChain（§9）、SessionOut / MessageOut（§5）、docgen 文档生成（§12）、
错误信封（§1）。
约定：模型 `extra="ignore"` + 全字段默认值，落库/返参前先经模型验证归一。
"""
from __future__ import annotations

from datetime import datetime, timezone
from typing import Any, Literal, Union

from pydantic import BaseModel, ConfigDict, Field


def utcnow_iso() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


# ---------------------------------------------------------------------------
# §2 块联合（判别字段 type）
# ---------------------------------------------------------------------------

class TextBlock(BaseModel):
    type: Literal["text"] = "text"
    html: str


class TableBlock(BaseModel):
    type: Literal["table"] = "table"
    headers: list[str] = Field(default_factory=list)
    rows: list[list[str]] = Field(default_factory=list)


class CodeBlock(BaseModel):
    type: Literal["code"] = "code"
    text: str


class StreamBlock(BaseModel):
    """瞬态块：只出现在流式过程中，绝不落库。"""

    type: Literal["stream"] = "stream"
    text: str


class PipelineStep(BaseModel):
    name: str = ""
    state: Literal["waiting", "running", "done"] = "waiting"
    desc: str = ""
    phase: str = ""  # planning|retrieval|evidence|generation|done（API.md §4）


class PipelineBlock(BaseModel):
    type: Literal["pipeline"] = "pipeline"
    status: Literal["running", "done", "error"] = "running"
    steps: list[PipelineStep] = Field(default_factory=list)


class ToolLog(BaseModel):
    text: str = ""
    cls: str = "log-ok"  # log-ok|log-run


class ToolBlock(BaseModel):
    type: Literal["tool"] = "tool"
    icon: str = ""
    name: str = ""
    status: Literal["running", "done", "error"] = "running"
    dur: str = ""
    open: bool = True
    params: str = ""
    result: str = ""
    logs: list[ToolLog] = Field(default_factory=list)
    phase: str = "retrieval"


class SourcesBlock(BaseModel):
    type: Literal["sources"] = "sources"
    items: list["SourceItem"] = Field(default_factory=list)


class GraphBlock(BaseModel):
    type: Literal["graph"] = "graph"
    nodes: list["GraphNode"] = Field(default_factory=list)
    edges: list["GraphEdge"] = Field(default_factory=list)
    chains: list["GraphChain"] = Field(default_factory=list)


class ArtifactBlock(BaseModel):
    """docgen 制品块（§2 块联合扩展，P2）：会话消息内挂载文档制品。

    只有 artifact_id 是 BFF 落库键；markdown/source 键保证前端渲染兜底
    （本地 mock 或制品文件缺失时仍可展示）。状态语义见 API.md 块状态约定。
    """

    type: Literal["artifact"] = "artifact"
    artifact_id: str = ""
    title: str = ""
    template_name: str = ""
    status: str = "done"  # running|done|error（生成中/已就绪/失败）


Block = Union[
    TextBlock, TableBlock, CodeBlock, StreamBlock,
    PipelineBlock, ToolBlock, SourcesBlock, GraphBlock, ArtifactBlock,
]


# ---------------------------------------------------------------------------
# §8 SourceItem —— 所有字段键必须存在（缺失为 null）
# ---------------------------------------------------------------------------

class SourceItem(BaseModel):
    model_config = ConfigDict(extra="ignore")

    id: str | None = None
    title: str | None = None
    paper_id: str | None = None
    chunk_id: str | None = None
    pmid: str | None = None
    pmcid: str | None = None
    doi: str | None = None
    url: str | None = None
    snippet: str | None = None
    score: float | None = None
    locator: str | None = None
    section: str | None = None
    support_count: int | None = None
    type: str | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)


# ---------------------------------------------------------------------------
# §9 图谱结构
# ---------------------------------------------------------------------------

class GraphNode(BaseModel):
    model_config = ConfigDict(extra="ignore")

    id: str
    label: str | None = None
    type: str | None = None


class GraphEdge(BaseModel):
    model_config = ConfigDict(extra="ignore")

    source: str | None = None
    target: str | None = None
    relation: str | None = None
    label: str | None = None
    chain_id: str | None = None
    hop: int | None = None
    paper_id: str | None = None
    pmid: str | None = None
    pmcid: str | None = None
    doi: str | None = None


class GraphHop(BaseModel):
    model_config = ConfigDict(extra="ignore")

    source: str | None = None
    target: str | None = None
    relation: str | None = None
    source_type: str | None = None
    target_type: str | None = None
    paper_id: str | None = None
    pmid: str | None = None
    pmcid: str | None = None
    doi: str | None = None


class GraphChain(BaseModel):
    model_config = ConfigDict(extra="ignore")

    chain_id: str | None = None
    hops: list[GraphHop] = Field(default_factory=list)


class GraphData(BaseModel):
    """消息内 graph 块与 /api/kg/neighborhood 共用的 {nodes,edges,chains} 形态。"""

    model_config = ConfigDict(extra="ignore")

    nodes: list[GraphNode] = Field(default_factory=list)
    edges: list[GraphEdge] = Field(default_factory=list)
    chains: list[GraphChain] = Field(default_factory=list)


# ---------------------------------------------------------------------------
# §5 会话与消息
# ---------------------------------------------------------------------------

class SessionMeta(BaseModel):
    title_manual: bool = False
    capabilities: dict[str, bool] = Field(default_factory=dict)


class SessionCreate(BaseModel):
    title: str = "新会话"

    def normalized_title(self) -> str:
        return self.title.strip() or "新会话"


class SessionPatch(BaseModel):
    title: str | None = None
    status: Literal["running", "done", "error", "idle"] | None = None
    meta: dict[str, Any] | None = None  # S4：仅 merge meta.capabilities（白名单）


class SessionOut(BaseModel):
    model_config = ConfigDict(extra="ignore")

    id: str
    title: str
    status: str = "idle"
    created_at: str = Field(default_factory=utcnow_iso)
    updated_at: str = Field(default_factory=utcnow_iso)
    message_count: int = 0
    meta: SessionMeta = Field(default_factory=SessionMeta)


class SessionDetail(SessionOut):
    messages: list["MessageOut"] = Field(default_factory=list)


class MessageCreate(BaseModel):
    role: Literal["user", "assistant"]
    text: str | None = None
    kind: str = "qa"
    status: str = "done"
    blocks: list[Block] = Field(default_factory=list)


class MessagePatch(BaseModel):
    """流式检查点局部更新（PATCH /messages/{mid}）。"""

    status: str | None = None
    blocks: list[Block] | None = None
    raw: dict[str, Any] | None = None


class MessageOut(BaseModel):
    model_config = ConfigDict(extra="ignore")

    id: str
    role: str
    kind: str = "qa"
    status: str = "done"
    blocks: list[Block] = Field(default_factory=list)
    raw: dict[str, Any] | None = None
    text: str | None = None
    created_at: str = Field(default_factory=utcnow_iso)
    updated_at: str = Field(default_factory=utcnow_iso)


# ---------------------------------------------------------------------------
# §12 文档生成（docgen，P2）——BFF ↔ UI 契约模型
# ---------------------------------------------------------------------------

class LibraryItemOut(BaseModel):
    """doc_library 表单行（主库管理返回形态）。"""

    model_config = ConfigDict(extra="ignore")

    id: str
    name: str = ""
    category: str = ""
    tags: str = ""
    description: str = ""
    as_template: bool = False
    active: bool = True
    created_at: str = Field(default_factory=utcnow_iso)


class LibraryPatch(BaseModel):
    """PATCH /api/docgen/library/{id}：改 meta 或 as_template 开关。"""

    name: str | None = None
    category: str | None = None
    tags: str | None = None
    description: str | None = None
    as_template: bool | None = None
    active: bool | None = None


class ArtifactOut(BaseModel):
    """doc_artifacts 表单行（生成/查询返回形态）。"""

    model_config = ConfigDict(extra="ignore")

    id: str
    session_id: str = ""
    status: str = "running"  # running|done|error
    title: str = ""
    template_name: str = ""
    template_key: str = ""          # 2026-09-04：内置模版 key（official/report/...）
    has_html: bool = False          # 模版 HTML 预览是否可用
    file_name: str = ""             # 2026-09-04：导出文件名基名（prompt 核心意思提炼）
    error: str | None = None
    created_at: str = Field(default_factory=utcnow_iso)


class DocgenGenerateRequest(BaseModel):
    """POST /api/docgen/generate 请求体（description 必填，上限 2000 字）。"""

    description: str
    session_id: str | None = None
    library_ids: list[str] = Field(default_factory=list)  # 空=全部主库资料
    template_id: str | None = None


class DocgenEditRequest(BaseModel):
    """POST /api/docgen/artifacts/{id}/edit 请求体。"""

    instruction: str


class ArtifactContentPatch(BaseModel):
    """PATCH /api/docgen/artifacts/{id}/content：保存编辑后的 markdown。

    2026-09-04：letterhead 可选——红头文字就地编辑后一并持久化（None=不改）。
    """

    markdown: str
    letterhead: str | None = None


class DocgenExportRequest(BaseModel):
    """POST /api/docgen/artifacts/{id}/export 请求体。"""

    format: Literal["pdf", "docx"]


class DocgenImportLine(BaseModel):
    """tooling/scripts/docgen_import.py CSV 单行（批量导入工具契约）。"""

    name: str
    category: str = ""
    tags: str = ""
    description: str = ""


# ---------------------------------------------------------------------------
# §1 错误信封
# ---------------------------------------------------------------------------

class ErrorBody(BaseModel):
    code: str
    message: str
    degraded: bool = False


class ErrorEnvelope(BaseModel):
    ok: Literal[False] = False
    error: ErrorBody