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

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

覆盖：块联合（§2）、SourceItem（§8）、GraphData / GraphNode / GraphEdge /
GraphChain（§9）、SessionOut / MessageOut（§5）、错误信封（§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)


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


# ---------------------------------------------------------------------------
# §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


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


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)


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

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


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