351 lines
13 KiB
Python
351 lines
13 KiB
Python
"""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（出现次数）。
|
||
|
||
用户级数据隔离：auth_enabled 开启时按 sessions.owner_id 过滤/校验（get_current_user 注入当前用户；
|
||
匿名模式下 user=None 则不过滤，兼容旧形态）。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 ..auth import get_current_user
|
||
from ..db import Database, dumps_json, get_db, loads_json
|
||
from ..schemas import (
|
||
MessageCreate,
|
||
MessageOut,
|
||
MessagePatch,
|
||
SessionCreate,
|
||
SessionDetail,
|
||
SessionOut,
|
||
SessionPatch,
|
||
)
|
||
from ..services.capabilities import sanitize_capabilities
|
||
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, owner_id: str | None = None
|
||
) -> SessionOut | None:
|
||
sql = _SESSION_SELECT + " WHERE s.id = ?"
|
||
params: list[Any] = [session_id]
|
||
if owner_id is not None:
|
||
sql += " AND s.owner_id = ?"
|
||
params.append(owner_id)
|
||
row = await db.fetchone(sql, tuple(params))
|
||
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), user: dict | None = Depends(get_current_user)) -> list[SessionOut]:
|
||
owner_id = user["id"] if user else None
|
||
sql = _SESSION_SELECT
|
||
params: list[Any] = []
|
||
if owner_id is not None:
|
||
sql += " WHERE s.owner_id = ?"
|
||
params.append(owner_id)
|
||
sql += " ORDER BY s.updated_at DESC"
|
||
rows = await db.fetchall(sql, tuple(params))
|
||
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), user: dict | None = Depends(get_current_user)
|
||
) -> SessionOut:
|
||
now = _db_now()
|
||
session_id = str(uuid.uuid4())
|
||
title = payload.normalized_title()
|
||
owner_id = user["id"] if user else None
|
||
await db.execute(
|
||
"INSERT INTO sessions (id, title, status, meta, owner_id, created_at, updated_at) "
|
||
"VALUES (?, ?, 'idle', ?, ?, ?, ?)",
|
||
(session_id, title, dumps_json({"title_manual": False}), owner_id, now, now),
|
||
)
|
||
out = await _get_session_or_none(db, session_id, owner_id=owner_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), user: dict | None = Depends(get_current_user)):
|
||
owner_id = user["id"] if user else None
|
||
session = await _get_session_or_none(db, session_id, owner_id=owner_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), user: dict | None = Depends(get_current_user)
|
||
):
|
||
owner_id = user["id"] if user else None
|
||
session = await _get_session_or_none(db, session_id, owner_id=owner_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),
|
||
)
|
||
if payload.meta is not None:
|
||
caps = sanitize_capabilities(payload.meta.get("capabilities"))
|
||
if caps:
|
||
meta = session.meta.model_dump()
|
||
merged_caps = dict(meta.get("capabilities") or {})
|
||
merged_caps.update(caps)
|
||
meta["capabilities"] = merged_caps
|
||
await db.execute(
|
||
"UPDATE sessions SET meta = ?, updated_at = ? WHERE id = ?",
|
||
(dumps_json(meta), _db_now(), session_id),
|
||
)
|
||
|
||
out = await _get_session_or_none(db, session_id, owner_id=owner_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), user: dict | None = Depends(get_current_user)) -> Response:
|
||
owner_id = user["id"] if user else None
|
||
session = await _get_session_or_none(db, session_id, owner_id=owner_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), user: dict | None = Depends(get_current_user)
|
||
):
|
||
owner_id = user["id"] if user else None
|
||
session = await _get_session_or_none(db, session_id, owner_id=owner_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:
|
||
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),
|
||
)
|
||
|
||
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),
|
||
user: dict | None = Depends(get_current_user),
|
||
):
|
||
owner_id = user["id"] if user else None
|
||
session = await _get_session_or_none(db, session_id, owner_id=owner_id)
|
||
if session is None:
|
||
return _not_found(session_id)
|
||
|
||
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
|
||
)
|
||
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), user: dict | None = Depends(get_current_user)):
|
||
"""按 paper_id 去重（无 paper_id 时按 id），带 count（出现次数）。"""
|
||
owner_id = user["id"] if user else None
|
||
session = await _get_session_or_none(db, session_id, owner_id=owner_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)
|