PlantAI-rag-new/sdk/app/routers/sessions.py

"""routers/sessions.py —— 会话与消息 CRUD + 会话级来源聚合（API.md §5）。

- GET/POST /api/sessions —— 列表（updated_at 倒序）/ 创建（title 可空，默认「新会话」）
- GET/PATCH/DELETE /api/sessions/{id} —— 详情（含 messages）/ 更新（title|status）/ 删除（204，级联消息）
- POST /api/sessions/{id}/messages —— 用户消息（{role:"user", text}）或助手流式骨架
  （{role:"assistant", kind, status:"streaming", blocks:[…]}）
- PATCH /api/sessions/{id}/messages/{mid} —— 局部更新 {status, blocks, raw}（流式检查点）
- GET /api/sessions/{id}/sources —— 全会话来源聚合：按 paper_id 去重（无 paper_id 按 id），
  带 count（出现次数）。

404 → NOT_FOUND 信封；时间 ISO 8601 UTC。
"""
from __future__ import annotations

import logging
import uuid
from datetime import datetime, timezone
from typing import Any

from fastapi import APIRouter, Depends, Response, status as http_status
from fastapi.responses import JSONResponse

from ..db import Database, dumps_json, get_db, loads_json
from ..schemas import (
    MessageCreate,
    MessageOut,
    MessagePatch,
    SessionCreate,
    SessionDetail,
    SessionOut,
    SessionPatch,
)
from ..upstream import envelope

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/sessions", tags=["sessions"])

_SESSION_SELECT = (
    "SELECT s.id, s.title, s.status, s.meta, s.created_at, s.updated_at, "
    "(SELECT COUNT(*) FROM messages m WHERE m.session_id = s.id) AS message_count "
    "FROM sessions s"
)

_MESSAGE_SELECT = (
    "SELECT id, session_id, role, kind, status, blocks, raw, text, created_at, updated_at "
    "FROM messages"
)


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


def _row_to_session(row: Any) -> SessionOut:
    return SessionOut(
        id=row["id"],
        title=row["title"],
        status=row["status"],
        created_at=row["created_at"],
        updated_at=row["updated_at"],
        message_count=row["message_count"],
        meta=loads_json(row["meta"], default={"title_manual": False}),
    )


def _row_to_message(row: Any) -> MessageOut:
    return MessageOut(
        id=row["id"],
        role=row["role"],
        kind=row["kind"],
        status=row["status"],
        blocks=loads_json(row["blocks"], default=[]),
        raw=loads_json(row["raw"], default=None),
        text=row["text"],
        created_at=row["created_at"],
        updated_at=row["updated_at"],
    )


async def _get_session_or_none(
    db: Database, session_id: str, *, with_count: bool = True
) -> SessionOut | None:
    sql = _SESSION_SELECT + " WHERE s.id = ?"
    row = await db.fetchone(sql, (session_id,))
    return _row_to_session(row) if row else None


async def _get_message_or_none(
    db: Database, session_id: str, message_id: str
) -> MessageOut | None:
    row = await db.fetchone(
        _MESSAGE_SELECT + " WHERE id = ? AND session_id = ?", (message_id, session_id)
    )
    return _row_to_message(row) if row else None


# ---------------------------------------------------------------------------
# 会话
# ---------------------------------------------------------------------------

@router.get("", response_model=list[SessionOut])
async def list_sessions(db: Database = Depends(get_db)) -> list[SessionOut]:
    rows = await db.fetchall(_SESSION_SELECT + " ORDER BY s.updated_at DESC")
    return [_row_to_session(r) for r in rows]


@router.post("", response_model=SessionOut, status_code=http_status.HTTP_201_CREATED)
async def create_session(
    payload: SessionCreate, db: Database = Depends(get_db)
) -> SessionOut:
    now = _db_now()
    session_id = str(uuid.uuid4())
    title = payload.normalized_title()
    await db.execute(
        "INSERT INTO sessions (id, title, status, meta, created_at, updated_at) "
        "VALUES (?, ?, 'idle', ?, ?, ?)",
        (session_id, title, dumps_json({"title_manual": False}), now, now),
    )
    out = await _get_session_or_none(db, session_id)
    assert out is not None  # 刚插入必存在
    return out


@router.get("/{session_id}", response_model=SessionDetail)
async def get_session(session_id: str, db: Database = Depends(get_db)):
    session = await _get_session_or_none(db, session_id)
    if session is None:
        return _not_found(session_id)
    rows = await db.fetchall(
        _MESSAGE_SELECT + " WHERE session_id = ? ORDER BY created_at ASC", (session_id,)
    )
    return SessionDetail(**session.model_dump(), messages=[_row_to_message(r) for r in rows])


@router.patch("/{session_id}", response_model=SessionOut)
async def patch_session(
    session_id: str, payload: SessionPatch, db: Database = Depends(get_db)
):
    session = await _get_session_or_none(db, session_id)
    if session is None:
        return _not_found(session_id)

    title = payload.title.strip() if payload.title is not None else None
    status = payload.status

    if title:
        meta = session.meta.model_dump()
        meta["title_manual"] = True
        await db.execute(
            "UPDATE sessions SET title = ?, meta = ?, updated_at = ? WHERE id = ?",
            (title, dumps_json(meta), _db_now(), session_id),
        )
    if status is not None:
        await db.execute(
            "UPDATE sessions SET status = ?, updated_at = ? WHERE id = ?",
            (status, _db_now(), session_id),
        )

    out = await _get_session_or_none(db, session_id)
    assert out is not None
    return out


@router.delete("/{session_id}", status_code=http_status.HTTP_204_NO_CONTENT)
async def delete_session(session_id: str, db: Database = Depends(get_db)) -> Response:
    session = await _get_session_or_none(db, session_id)
    if session is None:
        return _not_found_response(session_id)
    await db.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
    return Response(status_code=http_status.HTTP_204_NO_CONTENT)


# ---------------------------------------------------------------------------
# 消息
# ---------------------------------------------------------------------------

@router.post("/{session_id}/messages", response_model=MessageOut, status_code=201)
async def create_message(
    session_id: str, payload: MessageCreate, db: Database = Depends(get_db)
):
    session = await _get_session_or_none(db, session_id)
    if session is None:
        return _not_found(session_id)

    now = _db_now()
    message_id = str(uuid.uuid4())

    if payload.role == "user":
        text = (payload.text or "").strip()
        await db.execute(
            "INSERT INTO messages (id, session_id, role, kind, status, blocks, raw, text, "
            "created_at, updated_at) VALUES (?, ?, 'user', ?, 'done', '[]', NULL, ?, ?, ?)",
            (message_id, session_id, payload.kind, text, now, now),
        )
    else:
        # 助手流式骨架：text 置空，blocks 落库初始块
        blocks = [b.model_dump() for b in payload.blocks]
        await db.execute(
            "INSERT INTO messages (id, session_id, role, kind, status, blocks, raw, text, "
            "created_at, updated_at) VALUES (?, ?, 'assistant', ?, ?, ?, NULL, NULL, ?, ?)",
            (message_id, session_id, payload.kind, payload.status or "streaming",
             dumps_json(blocks), now, now),
        )

    # 消息活动 → 会话 updated_at 前移、idle → running
    await db.execute(
        "UPDATE sessions SET status = CASE WHEN status = 'idle' THEN 'running' ELSE status END, "
        "updated_at = ? WHERE id = ?",
        (now, session_id),
    )

    out = await _get_message_or_none(db, session_id, message_id)
    assert out is not None
    return out


@router.patch("/{session_id}/messages/{message_id}", response_model=MessageOut)
async def patch_message(
    session_id: str,
    message_id: str,
    payload: MessagePatch,
    db: Database = Depends(get_db),
):
    message = await _get_message_or_none(db, session_id, message_id)
    if message is None:
        return _not_found(session_id, message_id)

    sets: list[str] = []
    params: list[Any] = []

    if payload.status is not None:
        sets.append("status = ?")
        params.append(payload.status)
    if payload.blocks is not None:
        sets.append("blocks = ?")
        params.append(dumps_json([b.model_dump() for b in payload.blocks]))
    if payload.raw is not None:
        sets.append("raw = ?")
        params.append(dumps_json(payload.raw))

    if sets:
        sets.append("updated_at = ?")
        params.append(_db_now())
        params.append(message_id)
        await db.execute(
            f"UPDATE messages SET {', '.join(sets)} WHERE id = ?", params
        )
        # 同步前移会话 updated_at
        await db.execute(
            "UPDATE sessions SET updated_at = ? WHERE id = ?",
            (_db_now(), session_id),
        )

    out = await _get_message_or_none(db, session_id, message_id)
    assert out is not None
    return out


# ---------------------------------------------------------------------------
# 会话来源聚合
# ---------------------------------------------------------------------------

@router.get("/{session_id}/sources")
async def get_session_sources(session_id: str, db: Database = Depends(get_db)):
    """按 paper_id 去重（无 paper_id 时按 id），带 count（出现次数）。

    来源来自该会话全部助手消息 blocks 里的 sources 块（§2 块联合）。
    """
    session = await _get_session_or_none(db, session_id)
    if session is None:
        return _not_found(session_id)

    rows = await db.fetchall(
        "SELECT blocks FROM messages WHERE session_id = ? AND role = 'assistant'",
        (session_id,),
    )

    def is_sources_block(b: Any) -> bool:
        return isinstance(b, dict) and b.get("type") == "sources" and isinstance(b.get("items"), list)

    def key_of(item: dict[str, Any]) -> str:
        return str(item.get("paper_id") or item.get("id") or "")

    agg: dict[str, dict[str, Any]] = {}
    order: list[str] = []
    for row in rows:
        blocks = loads_json(row["blocks"], default=[])
        for block in blocks:
            if not is_sources_block(block):
                continue
            for item in block["items"]:
                if not isinstance(item, dict):
                    continue
                key = key_of(item)
                if not key:
                    continue
                if key not in agg:
                    agg[key] = dict(item)
                    agg[key]["count"] = 0
                    order.append(key)
                agg[key]["count"] += 1

    sources = [agg[k] for k in order]
    return {"session_id": session_id, "total": sum(s["count"] for s in sources), "sources": sources}


# ---------------------------------------------------------------------------
# 404 帮助
# ---------------------------------------------------------------------------

def _not_found(session_id: str, message_id: str | None = None) -> JSONResponse:
    detail = f"会话 {session_id} 不存在" if message_id is None else f"消息 {message_id} 不存在"
    body, status = envelope("NOT_FOUND", message=detail, status=404)
    return JSONResponse(status_code=status, content=body)


def _not_found_response(session_id: str) -> JSONResponse:
    return _not_found(session_id)