PlantAI-rag-new 首次提交：文献知识工作台
- BFF 薄适配层（FastAPI，8090）：QA 代理/SSE 中继/Neo4j 直连/抽取代理/会话持久化 - Vue3 前端：五相位检索流水线、文献来源与证据链、知识图谱、提取任务四 tab - 四档响应式（桌面/笔记本/平板/手机）、深浅主题 - 81 个 FR 全部实现，测试追溯矩阵 81/81，后端 121 测/前端 156 测 - 真实凭据已脱敏，运行时配置在 backend/.env（不入 git） Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
jerry
@@ -0,0 +1,65 @@
|
||||
"""config.py —— 环境变量驱动的 BFF 配置（pydantic-settings）。
|
||||
|
||||
全部键见 backend/.env.example。凭据类键（Neo4j）只存在于本进程，
|
||||
**绝不下发到浏览器**（不进入任何 API 响应）。
|
||||
|
||||
上游拓扑与默认值依据 docs/DISCOVERY.md：
|
||||
- QA 7920：`/api/v1/qa-v1-1-73`（DISCOVERY §4）
|
||||
- RAG 8000（Methodology API，DISCOVERY §3）/ ragAlt 7900（protected 容器，§3.7）
|
||||
- Neo4j 候选顺序：18788（direct，首选）→ 19788（normalized）→ 17687（protected，空库）
|
||||
凭据见 DISCOVERY §2.1，可由 BFF_NEO4J_* 环境变量覆盖（测试/部署）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from functools import lru_cache
|
||||
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
model_config = SettingsConfigDict(
|
||||
env_file=".env",
|
||||
env_file_encoding="utf-8",
|
||||
extra="ignore",
|
||||
)
|
||||
|
||||
# ---- 应用 ----
|
||||
app_name: str = "plantai-rag-new-bff"
|
||||
version: str = "0.1.0"
|
||||
|
||||
# ---- 上游地址（只读消费，见 ARCHITECTURE.md §2）----
|
||||
qa_base_url: str = "http://127.0.0.1:7920"
|
||||
qa_api_path: str = "/api/v1/qa-v1-1-73"
|
||||
rag_base_url: str = "http://127.0.0.1:8000"
|
||||
rag_alt_base_url: str = "http://127.0.0.1:7900"
|
||||
|
||||
# ---- Neo4j 候选（顺序即回退顺序，首选 18788；凭据见 DISCOVERY §2.1）----
|
||||
neo4j_bolt_uris: str = (
|
||||
"bolt://127.0.0.1:18788,bolt://127.0.0.1:19788,bolt://127.0.0.1:17687"
|
||||
)
|
||||
neo4j_user: str = "neo4j"
|
||||
neo4j_passwords: str = (
|
||||
"" # 真实凭据在 backend/.env（不入 git），见 backend/.env.example
|
||||
)
|
||||
|
||||
# ---- 超时阶梯（秒）----
|
||||
timeout_health: float = 1.0 # 健康探测
|
||||
timeout_qa_chat: float = 300.0 # qa chat（读超时）
|
||||
timeout_qa_search: float = 180.0 # qa search（读超时）
|
||||
timeout_upload: float = 600.0 # 抽取 upload（读超时）
|
||||
timeout_extraction: float = 15.0 # 抽取其余接口
|
||||
|
||||
# ---- 本地 ----
|
||||
db_path: str = "backend/data/plantai.db"
|
||||
health_cache_ttl: float = 5.0
|
||||
|
||||
# ---- 生产形态：BFF 静态托管 frontend/dist（ARCHITECTURE.md §3，M6 挂载）----
|
||||
frontend_dist_dir: str = "frontend/dist"
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_settings() -> Settings:
|
||||
return Settings()
|
||||
|
||||
|
||||
settings = get_settings()
|
||||
@@ -0,0 +1,163 @@
|
||||
"""db.py —— SQLite 持久层（aiosqlite + WAL，无 ORM）。
|
||||
|
||||
Schema 与 API.md §5 对齐：
|
||||
- sessions: id / title / status / meta / created_at / updated_at
|
||||
- messages: id / session_id(FK CASCADE) / role / kind / status / blocks(JSON)
|
||||
/ raw(JSON) / text / created_at / updated_at
|
||||
|
||||
启动时幂等迁移（CREATE TABLE IF NOT EXISTS）。时间一律 ISO 8601 UTC
|
||||
（`2026-08-25T12:34:56Z`），由 BFF 侧统一生成（SQLite 不写 DEFAULT 时间戳）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from collections.abc import AsyncIterator, Iterable
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import aiosqlite
|
||||
|
||||
from .config import get_settings
|
||||
|
||||
_SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS sessions (
|
||||
id TEXT PRIMARY KEY,
|
||||
title TEXT NOT NULL DEFAULT '新会话',
|
||||
status TEXT NOT NULL DEFAULT 'idle',
|
||||
meta TEXT NOT NULL DEFAULT '{}',
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
id TEXT PRIMARY KEY,
|
||||
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
|
||||
role TEXT NOT NULL,
|
||||
kind TEXT NOT NULL DEFAULT 'qa',
|
||||
status TEXT NOT NULL DEFAULT 'done',
|
||||
blocks TEXT NOT NULL DEFAULT '[]',
|
||||
raw TEXT,
|
||||
text TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_sessions_updated ON sessions(updated_at);
|
||||
"""
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
|
||||
|
||||
def _db_path() -> Path:
|
||||
path = Path(get_settings().db_path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
class Database:
|
||||
"""aiosqlite 单连接封装。
|
||||
|
||||
SQLite 写并发不高，BFF 用单写连接（WAL 模式下读写不互斥）；
|
||||
依赖注入经 get_db() 提供本实例。
|
||||
"""
|
||||
|
||||
def __init__(self, path: str | Path | None = None) -> None:
|
||||
self._path: Path | None = Path(path) if path is not None else None
|
||||
self._conn: aiosqlite.Connection | None = None
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
@property
|
||||
def path(self) -> Path:
|
||||
if self._path is None:
|
||||
self._path = _db_path()
|
||||
return self._path
|
||||
|
||||
async def connect(self) -> None:
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._conn = await aiosqlite.connect(self.path)
|
||||
self._conn.row_factory = aiosqlite.Row
|
||||
await self._conn.execute("PRAGMA journal_mode=WAL")
|
||||
await self._conn.execute("PRAGMA foreign_keys=ON")
|
||||
await self._conn.execute("PRAGMA busy_timeout=5000")
|
||||
await self._conn.executescript(_SCHEMA)
|
||||
await self._conn.commit()
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._conn is not None:
|
||||
await self._conn.close()
|
||||
self._conn = None
|
||||
|
||||
@property
|
||||
def conn(self) -> aiosqlite.Connection:
|
||||
if self._conn is None:
|
||||
raise RuntimeError("Database 未初始化：先 await init_db()")
|
||||
return self._conn
|
||||
|
||||
# ---- 通用查询（只读）----
|
||||
|
||||
async def fetchall(self, sql: str, params: Iterable[Any] = ()) -> list[aiosqlite.Row]:
|
||||
async with self._conn.execute(sql, tuple(params)) as cursor:
|
||||
return list(await cursor.fetchall())
|
||||
|
||||
async def fetchone(self, sql: str, params: Iterable[Any] = ()) -> aiosqlite.Row | None:
|
||||
async with self._conn.execute(sql, tuple(params)) as cursor:
|
||||
return await cursor.fetchone()
|
||||
|
||||
# ---- 写操作（串行化，提交后返回 lastrowid）----
|
||||
|
||||
async def execute(self, sql: str, params: Iterable[Any] = ()) -> int:
|
||||
async with self._lock:
|
||||
cur = await self._conn.execute(sql, tuple(params))
|
||||
await self._conn.commit()
|
||||
return cur.lastrowid
|
||||
|
||||
|
||||
db = Database()
|
||||
|
||||
|
||||
async def init_db(path: str | Path | None = None) -> Database:
|
||||
"""app 启动钩子：创建并迁移数据库（幂等）。返回全局单例（或指定路径的实例）。"""
|
||||
global db
|
||||
if path is not None:
|
||||
db = Database(path)
|
||||
await db.connect()
|
||||
return db
|
||||
|
||||
|
||||
async def get_db() -> AsyncIterator[Database]:
|
||||
"""FastAPI 依赖：每个请求交出同一数据库实例。"""
|
||||
yield db
|
||||
|
||||
|
||||
async def close_db() -> None:
|
||||
"""app shutdown 钩子。"""
|
||||
await db.close()
|
||||
|
||||
|
||||
# ---- JSON 编解码（落库用；失败时按契约置 null/[] 而不吞请求）----
|
||||
|
||||
def dumps_json(value: Any, default: Any = None) -> str:
|
||||
try:
|
||||
return json.dumps(value, ensure_ascii=False)
|
||||
except (TypeError, ValueError):
|
||||
return json.dumps(default, ensure_ascii=False)
|
||||
|
||||
|
||||
def loads_json(text: str | None, default: Any = None) -> Any:
|
||||
if text is None:
|
||||
return default
|
||||
try:
|
||||
return json.loads(text)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
def monotonic() -> float:
|
||||
"""测试/内部计时用（time.monotonic 别名，避免各处直接 import time）。"""
|
||||
return time.monotonic()
|
||||
@@ -0,0 +1,91 @@
|
||||
"""main.py —— BFF 应用入口（ARCHITECTURE.md §2/§3）。
|
||||
|
||||
- 挂全部 5 个 router（qa/kg/extraction 为 A2/A3 实现，import 路径不变）。
|
||||
- CORS：允许 http://localhost:5180 与 http://127.0.0.1:5180（Vite dev）。
|
||||
- lifespan：init_db（幂等迁移）；关闭回收 httpx 客户端与 Neo4j 连接。
|
||||
- 全局 RequestValidationError → VALIDATION 信封（API.md §1，2026-08-25 PM 集成）。
|
||||
- 生产形态（M6）：frontend/dist 存在时静态托管（ARCHITECTURE.md §3）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
||||
from .config import get_settings
|
||||
from .db import close_db, init_db
|
||||
from .routers import extraction, health, kg, qa, sessions
|
||||
from .upstream import close_client
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
_ALLOWED_ORIGINS = [
|
||||
"http://localhost:5180",
|
||||
"http://127.0.0.1:5180",
|
||||
]
|
||||
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
settings = get_settings()
|
||||
app = FastAPI(title=settings.app_name, version=settings.version)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=_ALLOWED_ORIGINS,
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
app.include_router(health.router)
|
||||
app.include_router(sessions.router)
|
||||
app.include_router(qa.router)
|
||||
app.include_router(kg.router)
|
||||
app.include_router(extraction.router)
|
||||
|
||||
@app.exception_handler(RequestValidationError)
|
||||
async def _validation_handler(_: Request, exc: RequestValidationError) -> JSONResponse:
|
||||
# API.md §1 已裁决：参数校验错误收敛为统一 VALIDATION 信封（422）
|
||||
return JSONResponse(
|
||||
status_code=422,
|
||||
content={
|
||||
"ok": False,
|
||||
"error": {"code": "VALIDATION", "message": "请求参数不合法", "degraded": False},
|
||||
},
|
||||
)
|
||||
|
||||
@asynccontextmanager
|
||||
async def _lifespan(_app: FastAPI):
|
||||
await init_db()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await close_client()
|
||||
try:
|
||||
from .services.neo4j import close_neo4j # A3 提供；缺失则跳过
|
||||
|
||||
await close_neo4j()
|
||||
except ImportError:
|
||||
pass
|
||||
await close_db()
|
||||
|
||||
app.router.lifespan_context = _lifespan
|
||||
|
||||
# 生产形态：静态托管 frontend/dist（存在时挂载；本地开发期可无）。
|
||||
# 2026-08-25 PM 修复：dist 路径必须基于项目根解析（BFF 工作目录是 backend/，
|
||||
# 相对路径会解析成 backend/frontend/dist）；sync.sh 排除 dist，需单独同步。
|
||||
project_root = Path(__file__).resolve().parent.parent.parent
|
||||
dist = project_root / settings.frontend_dist_dir if not Path(settings.frontend_dist_dir).is_absolute() else Path(settings.frontend_dist_dir)
|
||||
if dist.is_dir():
|
||||
app.mount("/", StaticFiles(directory=str(dist), html=True), name="frontend")
|
||||
|
||||
return app
|
||||
|
||||
|
||||
app = create_app()
|
||||
@@ -0,0 +1,150 @@
|
||||
"""routers/extraction.py —— 抽取任务代理（8000 主 + 7900 回落，API.md §10）。
|
||||
|
||||
- 上传仅限 PDF ≤50MB，否则 422 VALIDATION 信封。
|
||||
- 双上游全不可达 → 503 UPSTREAM_DOWN 降级信封。
|
||||
- 上游 5xx（如 MinerU 不可用）→ 502 UPSTREAM_ERROR。
|
||||
错误信封与 mock_server 一致：直接返回 JSONResponse，信封在顶层。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any, AsyncIterator
|
||||
|
||||
from fastapi import APIRouter, File, Form, Query, UploadFile
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from ..services import extraction as _extraction
|
||||
from ..services.extraction import (
|
||||
MAX_PDF_BYTES,
|
||||
ExtractionUnavailable,
|
||||
ExtractionUpstreamError,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/api/extraction", tags=["extraction"])
|
||||
|
||||
_DOWN_MSG = "抽取服务当前不可用，请稍后重试"
|
||||
_MINERU_MSG = "抽取任务提交失败：上游 MinerU 服务不可用，请检查抽取环境后重试"
|
||||
|
||||
|
||||
def _envelope(code: str, message: str, degraded: bool, status: int) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
status_code=status,
|
||||
content={
|
||||
"ok": False,
|
||||
"error": {"code": code, "message": message, "degraded": degraded},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _down(exc: BaseException | None = None) -> JSONResponse:
|
||||
logger.warning("extraction: 上游不可用：%s", exc)
|
||||
return _envelope("UPSTREAM_DOWN", _DOWN_MSG, True, 503)
|
||||
|
||||
|
||||
def _upstream_error(message: str) -> JSONResponse:
|
||||
logger.warning("extraction: 上游错误：%s", message)
|
||||
return _envelope("UPSTREAM_ERROR", message, True, 502)
|
||||
|
||||
|
||||
def _validation(message: str) -> JSONResponse:
|
||||
return _envelope("VALIDATION", message, False, 422)
|
||||
|
||||
|
||||
def _check_pdf_upload(filename: str, size: int) -> JSONResponse | None:
|
||||
"""FR-EXTRACT-004：仅 PDF，单文件 ≤50MB。合法返回 None。"""
|
||||
if size <= 0:
|
||||
return _validation("上传文件为空")
|
||||
if size > MAX_PDF_BYTES:
|
||||
return _validation("文件超过 50MB 限制")
|
||||
if not filename.lower().endswith(".pdf"):
|
||||
return _validation("仅支持 PDF 文件")
|
||||
return None
|
||||
|
||||
|
||||
async def _stream_upload(file: UploadFile) -> AsyncIterator[bytes]:
|
||||
try:
|
||||
while chunk := await file.read(1024 * 1024): # 1MB
|
||||
yield chunk
|
||||
finally:
|
||||
await file.close()
|
||||
|
||||
|
||||
@router.get("/jobs")
|
||||
async def jobs(limit: int = Query(20, ge=1, le=200)) -> dict[str, Any]:
|
||||
try:
|
||||
return await _extraction.service.list_jobs(limit=limit)
|
||||
except ExtractionUnavailable as exc:
|
||||
return _down(exc)
|
||||
|
||||
|
||||
@router.get("/jobs/{job_id}")
|
||||
async def job_detail(job_id: str) -> dict[str, Any]:
|
||||
try:
|
||||
return await _extraction.service.get_job(job_id)
|
||||
except ExtractionUnavailable as exc:
|
||||
return _down(exc)
|
||||
|
||||
|
||||
@router.get("/processing-status")
|
||||
async def processing_status() -> dict[str, Any]:
|
||||
try:
|
||||
return await _extraction.service._passthrough("/processing-status")
|
||||
except ExtractionUnavailable as exc:
|
||||
return _down(exc)
|
||||
|
||||
|
||||
@router.get("/progress")
|
||||
async def progress() -> dict[str, Any]:
|
||||
try:
|
||||
return await _extraction.service._passthrough("/progress")
|
||||
except ExtractionUnavailable as exc:
|
||||
return _down(exc)
|
||||
|
||||
|
||||
@router.get("/summary")
|
||||
async def summary() -> dict[str, Any]:
|
||||
try:
|
||||
return await _extraction.service._passthrough("/summary")
|
||||
except ExtractionUnavailable as exc:
|
||||
return _down(exc)
|
||||
|
||||
|
||||
@router.get("/run-log")
|
||||
async def run_log() -> dict[str, Any]:
|
||||
try:
|
||||
return await _extraction.service._passthrough("/run-log")
|
||||
except ExtractionUnavailable as exc:
|
||||
return _down(exc)
|
||||
|
||||
|
||||
@router.post("/upload")
|
||||
async def upload(
|
||||
file: UploadFile = File(...),
|
||||
paper_id: str = Form("uploaded_paper"),
|
||||
family: str = Form(""),
|
||||
output_name: str = Form(""),
|
||||
) -> dict[str, Any]:
|
||||
filename = file.filename or ""
|
||||
size = file.size or 0
|
||||
invalid = _check_pdf_upload(filename, size)
|
||||
if invalid is not None:
|
||||
return invalid
|
||||
try:
|
||||
return await _extraction.service.upload(
|
||||
filename=Path(filename).name,
|
||||
stream=_stream_upload(file),
|
||||
content_type=file.content_type or "application/pdf",
|
||||
paper_id=paper_id,
|
||||
family=family,
|
||||
output_name=output_name,
|
||||
)
|
||||
except ExtractionUnavailable as exc:
|
||||
return _down(exc)
|
||||
except ExtractionUpstreamError as exc:
|
||||
return _upstream_error(f"{_MINERU_MSG}（{exc}）")
|
||||
except Exception as exc:
|
||||
logger.exception("extraction: 上传失败")
|
||||
return _upstream_error(f"提交抽取任务失败：{exc}")
|
||||
@@ -0,0 +1,35 @@
|
||||
"""routers/health.py —— GET /api/health（API.md §3）。
|
||||
|
||||
并发探测 7920、8000、7900 与 Neo4j（1s 超时，5s 缓存），外加 db 状态。
|
||||
返回 `{status, version, upstreams, db}`；健康接口不包错误信封（API.md §1 例外）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import APIRouter
|
||||
|
||||
from ..config import get_settings
|
||||
from ..db import db
|
||||
from ..upstream import probe_all
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["health"])
|
||||
|
||||
|
||||
async def _db_status() -> dict:
|
||||
try:
|
||||
rows = await db.fetchall("SELECT COUNT(*) AS n FROM sessions")
|
||||
return {"ok": True, "sessions": rows[0]["n"] if rows else 0}
|
||||
except Exception:
|
||||
return {"ok": False, "sessions": 0}
|
||||
|
||||
|
||||
@router.get("/health")
|
||||
async def health() -> dict:
|
||||
upstreams = await probe_all()
|
||||
db_status = await _db_status()
|
||||
degraded = not all(u.get("ok") for u in upstreams.values()) or not db_status.get("ok")
|
||||
return {
|
||||
"status": "degraded" if degraded else "ok",
|
||||
"version": get_settings().version,
|
||||
"upstreams": upstreams,
|
||||
"db": db_status,
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
"""routers/kg.py —— 知识图谱浏览（BFF Bolt 直连 Neo4j，API.md §6）。
|
||||
|
||||
Neo4j 不可用 → 503 UPSTREAM_DOWN 降级信封；未命中实体 → 404 NOT_FOUND。
|
||||
错误信封与 mock_server 一致：直接返回 JSONResponse，信封在顶层
|
||||
（不用 HTTPException，避免被 FastAPI 包成 {"detail": {...}}）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Query
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from ..services import neo4j as _neo4j
|
||||
from ..services.neo4j import Neo4jUnavailable
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/api/kg", tags=["kg"])
|
||||
|
||||
_UPSTREAM_DOWN_MSG = "知识图谱服务当前不可用，请稍后重试"
|
||||
|
||||
|
||||
def _down(exc: BaseException | None = None) -> JSONResponse:
|
||||
logger.warning("kg: Neo4j 不可用：%s", exc)
|
||||
return JSONResponse(
|
||||
status_code=503,
|
||||
content={
|
||||
"ok": False,
|
||||
"error": {
|
||||
"code": "UPSTREAM_DOWN",
|
||||
"message": _UPSTREAM_DOWN_MSG,
|
||||
"degraded": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _not_found() -> JSONResponse:
|
||||
return JSONResponse(
|
||||
status_code=404,
|
||||
content={
|
||||
"ok": False,
|
||||
"error": {
|
||||
"code": "NOT_FOUND",
|
||||
"message": "未找到该实体",
|
||||
"degraded": False,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.get("/status")
|
||||
async def kg_status():
|
||||
try:
|
||||
return await _neo4j.service.kg_status()
|
||||
except Neo4jUnavailable as exc:
|
||||
return _down(exc)
|
||||
|
||||
|
||||
@router.get("/overview")
|
||||
async def kg_overview():
|
||||
try:
|
||||
return await _neo4j.service.kg_overview()
|
||||
except Neo4jUnavailable as exc:
|
||||
return _down(exc)
|
||||
|
||||
|
||||
@router.get("/entities")
|
||||
async def kg_entities(
|
||||
q: str = Query("", description="实体名（name_norm 精确或前缀匹配）"),
|
||||
limit: int = Query(20, ge=1, le=20),
|
||||
) -> dict[str, Any]:
|
||||
try:
|
||||
results = await _neo4j.service.search_entities(q, limit=limit)
|
||||
except Neo4jUnavailable as exc:
|
||||
return _down(exc)
|
||||
return {"query": q, "results": results}
|
||||
|
||||
|
||||
@router.get("/entities/{entity_id}")
|
||||
async def kg_entity(entity_id: str) -> dict[str, Any]:
|
||||
try:
|
||||
entity = await _neo4j.service.entity(entity_id)
|
||||
except Neo4jUnavailable as exc:
|
||||
return _down(exc)
|
||||
if entity is None:
|
||||
return _not_found()
|
||||
return entity
|
||||
|
||||
|
||||
@router.get("/entities/{entity_id}/neighborhood")
|
||||
async def kg_neighborhood(
|
||||
entity_id: str,
|
||||
depth: int = Query(1, ge=1, le=2),
|
||||
limit: int = Query(50, ge=1, le=200),
|
||||
) -> dict[str, Any]:
|
||||
try:
|
||||
result = await _neo4j.service.neighborhood(entity_id, depth=depth, limit=limit)
|
||||
except Neo4jUnavailable as exc:
|
||||
return _down(exc)
|
||||
if result.get("center") is None:
|
||||
return _not_found()
|
||||
return result
|
||||
@@ -0,0 +1,211 @@
|
||||
"""routers/qa.py —— /api/qa/* ：7920 代理 + 归一化 + SSE 中继（A2 归属，API.md §3/§4）。
|
||||
|
||||
端点：
|
||||
- GET /api/qa/info 上游 /info 精简（name/version/pipeline/config）
|
||||
- POST /api/qa/chat 上游 /chat，300s 读超时，invalid_llm_citation 自动重试一次
|
||||
- POST /api/qa/chat/stream SSE 中继（app.services.sse_relay.relay_stream）
|
||||
- POST /api/qa/search 上游 /search，强制 use_llm:false，180s 读超时，不重试，
|
||||
响应 mode:"search"、answer:""
|
||||
|
||||
请求体一律透传（前端字段集见 API.md §4，BFF 不增删字段）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Optional
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, Request
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
|
||||
from ..services import sse_relay
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/api/qa", tags=["qa"])
|
||||
|
||||
# 超时阶梯（TODO(A2→A1 对齐点)：A1 的 app.config.settings 就绪后改为配置读取）
|
||||
CHAT_READ_TIMEOUT = 300.0
|
||||
SEARCH_READ_TIMEOUT = 180.0
|
||||
INFO_READ_TIMEOUT = 5.0
|
||||
CONNECT_TIMEOUT = 10.0
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 错误信封（TODO(A2→A1 对齐点)：HTTP JSON 信封换用 app.upstream.envelope()）
|
||||
# ---------------------------------------------------------------------------
|
||||
MESSAGES = {
|
||||
"UPSTREAM_DOWN": "知识问答服务当前不可用，请稍后重试",
|
||||
"UPSTREAM_ERROR": "知识问答服务返回错误",
|
||||
"TIMEOUT": "检索超时",
|
||||
"INVALID_CITATION": "引用校验失败，请重试",
|
||||
"VALIDATION": "请求参数不合法",
|
||||
"INTERNAL": "内部错误",
|
||||
}
|
||||
|
||||
|
||||
def _envelope(code: str, status: int, message: Optional[str] = None, degraded: Optional[bool] = None) -> JSONResponse:
|
||||
if degraded is None:
|
||||
degraded = code == "UPSTREAM_DOWN"
|
||||
return JSONResponse(
|
||||
status_code=status,
|
||||
content={"ok": False, "error": {"code": code, "message": message or MESSAGES[code], "degraded": degraded}},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 上游调用
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _upstream_path(endpoint: str) -> str:
|
||||
return f"{sse_relay.QA_UPSTREAM_PREFIX}/{endpoint}"
|
||||
|
||||
|
||||
def _read_timeout(seconds: float) -> httpx.Timeout:
|
||||
return httpx.Timeout(connect=CONNECT_TIMEOUT, read=seconds, write=30.0, pool=CONNECT_TIMEOUT)
|
||||
|
||||
|
||||
async def _parse_json(text: Optional[bytes]) -> Any:
|
||||
if not text:
|
||||
return None
|
||||
try:
|
||||
return json.loads(text.decode("utf-8", errors="replace"))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _error_from_status(status: int, payload: Any) -> tuple[str, int, Optional[str]]:
|
||||
"""上游 HTTP≥400 → (code, bff_http, message)。HTTP 按 API.md §1 固定映射，不透传上游状态码。"""
|
||||
if sse_relay.is_invalid_citation(payload):
|
||||
return "INVALID_CITATION", 502, None
|
||||
if status in (408, 504):
|
||||
return "TIMEOUT", 504, None
|
||||
detail = payload.get("detail") if isinstance(payload, dict) else None
|
||||
if isinstance(detail, dict):
|
||||
message = detail.get("message") or detail.get("error")
|
||||
return "UPSTREAM_ERROR", 502, message
|
||||
if isinstance(payload, dict):
|
||||
message = payload.get("message") or payload.get("error")
|
||||
return "UPSTREAM_ERROR", 502, message
|
||||
return "UPSTREAM_ERROR", 502, None
|
||||
|
||||
|
||||
async def _proxy(
|
||||
endpoint: str,
|
||||
body: Optional[dict] = None,
|
||||
*,
|
||||
read_timeout: float,
|
||||
retry_invalid_citation: bool,
|
||||
mode: str,
|
||||
normalizer=None,
|
||||
) -> Any:
|
||||
"""非流式上游代理：重试、错误信封、归一化。返回 (JSONResponse)。"""
|
||||
client = sse_relay.get_qa_client()
|
||||
normalizer = normalizer or sse_relay.get_qa_normalizer()
|
||||
url = f"{sse_relay.qa_base_url()}{_upstream_path(endpoint)}"
|
||||
timeout = _read_timeout(read_timeout)
|
||||
json_body = body if body is not None else {}
|
||||
|
||||
for attempt in (1, 2):
|
||||
try:
|
||||
resp = await client.request("POST", url, json=json_body, timeout=timeout)
|
||||
except httpx.ReadTimeout:
|
||||
return _envelope("TIMEOUT", 504)
|
||||
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout) as exc:
|
||||
logger.warning("上游 %s 连接失败: %s", endpoint, exc)
|
||||
return _envelope("UPSTREAM_DOWN", 503)
|
||||
except httpx.HTTPError as exc:
|
||||
logger.warning("上游 %s 请求失败: %s", endpoint, exc)
|
||||
return _envelope("UPSTREAM_ERROR", 502)
|
||||
payload = await _parse_json(resp.content)
|
||||
if resp.status_code >= 400:
|
||||
if (
|
||||
retry_invalid_citation
|
||||
and attempt == 1
|
||||
and sse_relay.is_invalid_citation(payload)
|
||||
):
|
||||
logger.warning(
|
||||
"上游 %s 返回 %d invalid_llm_citation，自动重试（第 1 次）", endpoint, resp.status_code
|
||||
)
|
||||
continue
|
||||
code, bff_status, message = _error_from_status(resp.status_code, payload)
|
||||
return _envelope(code, bff_status, message)
|
||||
try:
|
||||
normalized = normalizer(payload if isinstance(payload, dict) else {})
|
||||
except Exception:
|
||||
logger.exception("归一化上游 %s 响应失败", endpoint)
|
||||
return _envelope("INTERNAL", 500)
|
||||
normalized["mode"] = mode
|
||||
if mode == "search":
|
||||
normalized["answer"] = ""
|
||||
return JSONResponse(content=normalized)
|
||||
# 不可达（重试循环内必然 return）；防御性兜底
|
||||
return _envelope("UPSTREAM_ERROR", 502)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 端点
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.get("/info")
|
||||
async def qa_info():
|
||||
"""上游 /info 精简（API.md §3）。"""
|
||||
client = sse_relay.get_qa_client()
|
||||
url = f"{sse_relay.qa_base_url()}{_upstream_path('info')}"
|
||||
try:
|
||||
resp = await client.get(url, timeout=_read_timeout(INFO_READ_TIMEOUT))
|
||||
except httpx.ReadTimeout:
|
||||
return _envelope("TIMEOUT", 504)
|
||||
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout) as exc:
|
||||
logger.warning("上游 /info 连接失败: %s", exc)
|
||||
return _envelope("UPSTREAM_DOWN", 503)
|
||||
except httpx.HTTPError as exc:
|
||||
logger.warning("上游 /info 请求失败: %s", exc)
|
||||
return _envelope("UPSTREAM_ERROR", 502)
|
||||
payload = await _parse_json(resp.content)
|
||||
if resp.status_code >= 400 or not isinstance(payload, dict):
|
||||
return _envelope("UPSTREAM_ERROR", 502)
|
||||
info = payload.get("info") if isinstance(payload.get("info"), dict) else payload
|
||||
return JSONResponse(
|
||||
content={
|
||||
"name": info.get("name") or payload.get("name"),
|
||||
"version": info.get("version") or payload.get("version"),
|
||||
"pipeline": info.get("pipeline") or payload.get("pipeline") or [],
|
||||
"config": info.get("config") or payload.get("config") or {},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@router.post("/chat")
|
||||
async def qa_chat(request: Request):
|
||||
"""非流式问答：上游 /chat + 归一化 + invalid_llm_citation 重试一次。"""
|
||||
body = await request.json()
|
||||
return await _proxy(
|
||||
"chat", body, read_timeout=CHAT_READ_TIMEOUT, retry_invalid_citation=True, mode="chat"
|
||||
)
|
||||
|
||||
|
||||
@router.post("/chat/stream")
|
||||
async def qa_chat_stream(request: Request):
|
||||
"""SSE 中继（API.md §4）。"""
|
||||
try:
|
||||
body = await request.json()
|
||||
except Exception:
|
||||
return _envelope("VALIDATION", 422)
|
||||
return StreamingResponse(
|
||||
sse_relay.relay_stream(body),
|
||||
media_type="text/event-stream",
|
||||
headers=sse_relay.sse_headers,
|
||||
)
|
||||
|
||||
|
||||
@router.post("/search")
|
||||
async def qa_search(request: Request):
|
||||
"""检索模式：强制 use_llm:false，180s 读超时，不重试。"""
|
||||
body = await request.json()
|
||||
body = {**body, "use_llm": False}
|
||||
return await _proxy(
|
||||
"search", body, read_timeout=SEARCH_READ_TIMEOUT, retry_invalid_citation=False, mode="search"
|
||||
)
|
||||
@@ -0,0 +1,319 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,207 @@
|
||||
"""routers/test_kg.py —— kg 路由单元测试。
|
||||
|
||||
用假 Neo4jService 替换 services.neo4j.service（monkeypatch），覆盖：
|
||||
五个端点成功形态、Neo4j 不可用 → 503 UPSTREAM_DOWN、未命中实体 → 404
|
||||
NOT_FOUND、参数校验（limit ≤ 20、depth 1-2）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
|
||||
from app.routers import kg as kg_router
|
||||
from app.services import neo4j as neo4j_mod
|
||||
from app.services.neo4j import Neo4jUnavailable
|
||||
|
||||
|
||||
class FakeService:
|
||||
"""方法逐项可配的假服务。"""
|
||||
|
||||
def __init__(self, **overrides):
|
||||
self.calls: dict[str, list] = {}
|
||||
self.overrides = overrides
|
||||
|
||||
async def kg_status(self):
|
||||
return self.overrides.get("kg_status", {"ok": True, "active": "x"})
|
||||
|
||||
async def kg_overview(self):
|
||||
return self.overrides.get("kg_overview", {"entities_by_type": []})
|
||||
|
||||
async def search_entities(self, q, limit=20):
|
||||
self.calls.setdefault("search_entities", []).append((q, limit))
|
||||
return self.overrides.get(
|
||||
"search_entities",
|
||||
[{"id": "cent_ospl14", "name": "OsSPL14", "type": "gene", "support_count": 18}],
|
||||
)
|
||||
|
||||
async def entity(self, entity_id):
|
||||
self.calls.setdefault("entity", []).append(entity_id)
|
||||
return self.overrides.get(
|
||||
"entity",
|
||||
{"id": "cent_ospl14", "name": "OsSPL14", "type": "gene", "properties": {}},
|
||||
)
|
||||
|
||||
async def neighborhood(self, entity_id, depth=1, limit=200):
|
||||
self.calls.setdefault("neighborhood", []).append((entity_id, depth, limit))
|
||||
return self.overrides.get(
|
||||
"neighborhood",
|
||||
{
|
||||
"center": "cent_ospl14",
|
||||
"center_name": "OsSPL14",
|
||||
"depth": depth,
|
||||
"nodes": [{"id": "cent_ospl14", "label": "OsSPL14", "type": "gene"}],
|
||||
"edges": [],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_service(monkeypatch):
|
||||
fake = FakeService()
|
||||
monkeypatch.setattr(neo4j_mod, "service", fake)
|
||||
return fake
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
# 用真实 FastAPI() 包装：自动获得 ExceptionMiddleware + AsyncExitStackMiddleware，
|
||||
# 与生产挂载形态一致（生产 main.py 把 kg.router include 进 FastAPI 实例）。
|
||||
app = FastAPI()
|
||||
app.include_router(kg_router.router)
|
||||
return httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=app, raise_app_exceptions=False),
|
||||
base_url="http://test",
|
||||
)
|
||||
|
||||
async def test_status_ok(fake_service, client):
|
||||
fake_service.overrides["kg_status"] = {
|
||||
"ok": True,
|
||||
"active": "neo4j-plant-srt-direct-v1-full",
|
||||
"schema": "canonical",
|
||||
"labels": ["SRT2CanonicalEntity", "SRT2CanonicalAssertion", "SRT2Paper"],
|
||||
"relationship_types": ["SRT2_CANONICAL_ASSERTS", "SRT2_SUPPORTED_BY"],
|
||||
"entity_count": 3820996,
|
||||
"assertion_count": 8246875,
|
||||
"name_index": "srt2_canonical_entity_name_norm",
|
||||
"fulltext_index": None,
|
||||
}
|
||||
resp = await client.get("/api/kg/status")
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["entity_count"] == 3820996
|
||||
|
||||
async def test_status_down_returns_503_envelope(fake_service, client):
|
||||
async def boom():
|
||||
raise Neo4jUnavailable("all down")
|
||||
|
||||
fake_service.kg_status = boom
|
||||
resp = await client.get("/api/kg/status")
|
||||
assert resp.status_code == 503
|
||||
body = resp.json()
|
||||
assert body["ok"] is False
|
||||
assert body["error"]["code"] == "UPSTREAM_DOWN"
|
||||
assert body["error"]["degraded"] is True
|
||||
|
||||
async def test_overview_ok(fake_service, client):
|
||||
fake_service.overrides["kg_overview"] = {
|
||||
"entities_by_type": [{"type": "gene", "count": 43210}],
|
||||
"relations_by_category": [{"type": "characterized_in", "count": 12033}],
|
||||
}
|
||||
resp = await client.get("/api/kg/overview")
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["entities_by_type"][0]["type"] == "gene"
|
||||
assert body["relations_by_category"][0]["count"] == 12033
|
||||
|
||||
async def test_overview_down_returns_503(fake_service, client):
|
||||
async def boom():
|
||||
raise Neo4jUnavailable("down")
|
||||
|
||||
fake_service.kg_overview = boom
|
||||
resp = await client.get("/api/kg/overview")
|
||||
assert resp.status_code == 503
|
||||
assert resp.json()["error"]["code"] == "UPSTREAM_DOWN"
|
||||
|
||||
async def test_entities_ok_and_params_passed(fake_service, client):
|
||||
resp = await client.get("/api/kg/entities", params={"q": "OsSPL14", "limit": 5})
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["query"] == "OsSPL14"
|
||||
assert body["results"][0]["id"] == "cent_ospl14"
|
||||
assert fake_service.calls["search_entities"] == [("OsSPL14", 5)]
|
||||
|
||||
async def test_entities_limit_over_20_rejected(client):
|
||||
resp = await client.get("/api/kg/entities", params={"q": "OsSPL14", "limit": 21})
|
||||
assert resp.status_code == 422 # FastAPI 参数校验（VALIDATION 语义由 A1 统一处理）
|
||||
|
||||
async def test_entities_down_returns_503(fake_service, client):
|
||||
async def boom(q, limit=20):
|
||||
raise Neo4jUnavailable("down")
|
||||
|
||||
fake_service.search_entities = boom
|
||||
resp = await client.get("/api/kg/entities", params={"q": "x"})
|
||||
assert resp.status_code == 503
|
||||
assert resp.json()["error"]["code"] == "UPSTREAM_DOWN"
|
||||
|
||||
async def test_entity_ok(fake_service, client):
|
||||
resp = await client.get("/api/kg/entities/cent_ospl14")
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["id"] == "cent_ospl14"
|
||||
assert body["name"] == "OsSPL14"
|
||||
assert fake_service.calls["entity"] == ["cent_ospl14"]
|
||||
|
||||
async def test_entity_miss_returns_404(fake_service, client):
|
||||
fake_service.overrides["entity"] = None
|
||||
resp = await client.get("/api/kg/entities/NoSuch")
|
||||
assert resp.status_code == 404
|
||||
body = resp.json()
|
||||
assert body["ok"] is False
|
||||
assert body["error"]["code"] == "NOT_FOUND"
|
||||
assert body["error"]["degraded"] is False
|
||||
|
||||
async def test_entity_down_returns_503(fake_service, client):
|
||||
async def boom(entity_id):
|
||||
raise Neo4jUnavailable("down")
|
||||
|
||||
fake_service.entity = boom
|
||||
resp = await client.get("/api/kg/entities/cent_ospl14")
|
||||
assert resp.status_code == 503
|
||||
assert resp.json()["error"]["code"] == "UPSTREAM_DOWN"
|
||||
|
||||
async def test_neighborhood_ok_and_params(fake_service, client):
|
||||
resp = await client.get(
|
||||
"/api/kg/entities/cent_ospl14/neighborhood", params={"depth": 2, "limit": 100}
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["center"] == "cent_ospl14"
|
||||
assert body["center_name"] == "OsSPL14"
|
||||
assert body["depth"] == 2
|
||||
assert fake_service.calls["neighborhood"] == [("cent_ospl14", 2, 100)]
|
||||
|
||||
async def test_neighborhood_defaults_depth1_limit50(fake_service, client):
|
||||
resp = await client.get("/api/kg/entities/cent_ospl14/neighborhood")
|
||||
assert resp.status_code == 200
|
||||
assert fake_service.calls["neighborhood"] == [("cent_ospl14", 1, 50)]
|
||||
|
||||
async def test_neighborhood_depth_out_of_range_rejected(client):
|
||||
resp = await client.get("/api/kg/entities/cent_ospl14/neighborhood", params={"depth": 3})
|
||||
assert resp.status_code == 422
|
||||
|
||||
async def test_neighborhood_miss_returns_404(fake_service, client):
|
||||
fake_service.overrides["neighborhood"] = {
|
||||
"center": None, "center_name": None, "depth": 1, "nodes": [], "edges": []
|
||||
}
|
||||
resp = await client.get("/api/kg/entities/NoSuch/neighborhood")
|
||||
assert resp.status_code == 404
|
||||
assert resp.json()["error"]["code"] == "NOT_FOUND"
|
||||
|
||||
async def test_neighborhood_down_returns_503(fake_service, client):
|
||||
async def boom(entity_id, depth=1, limit=200):
|
||||
raise Neo4jUnavailable("down")
|
||||
|
||||
fake_service.neighborhood = boom
|
||||
resp = await client.get("/api/kg/entities/cent_ospl14/neighborhood")
|
||||
assert resp.status_code == 503
|
||||
assert resp.json()["error"]["code"] == "UPSTREAM_DOWN"
|
||||
@@ -0,0 +1,260 @@
|
||||
"""test_qa.py —— /api/qa/* 路由单元测试（A2 归属）。
|
||||
|
||||
通过 httpx.MockTransport 注入假上游（把 sse_relay.get_qa_client 替换为指向
|
||||
MockTransport 的 client），不依赖真实 7920。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from app.main import app
|
||||
from app.services import sse_relay
|
||||
|
||||
|
||||
def _transport_handler(handler) -> httpx.AsyncClient:
|
||||
"""以 MockTransport 构造一个 client。"""
|
||||
return httpx.AsyncClient(transport=httpx.MockTransport(handler))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bff_client():
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
UPSTREAM_INFO = {
|
||||
"name": "QA V1.1.73 Assertion-Anchored Full Graph",
|
||||
"version": "1.1.73-assertion-anchored",
|
||||
"pipeline": ["question", "provisional answer"],
|
||||
"config": {"max_evidence": 40, "bge_rerank": True},
|
||||
}
|
||||
|
||||
|
||||
def _json_response(payload, status=200) -> httpx.Response:
|
||||
return httpx.Response(status, json=payload)
|
||||
|
||||
|
||||
def test_info_passthrough_trimmed(bff_client, monkeypatch):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == "GET"
|
||||
assert request.url.path.endswith("/api/v1/qa-v1-1-73/info")
|
||||
return _json_response(UPSTREAM_INFO)
|
||||
|
||||
client = _transport_handler(handler)
|
||||
monkeypatch.setattr(sse_relay, "get_qa_client", lambda: client)
|
||||
|
||||
resp = bff_client.get("/api/qa/info")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data == UPSTREAM_INFO
|
||||
assert set(data.keys()) == {"name", "version", "pipeline", "config"}
|
||||
|
||||
|
||||
def test_info_upstream_down_503(bff_client, monkeypatch):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
raise httpx.ConnectError("down", request=request)
|
||||
|
||||
client = _transport_handler(handler)
|
||||
monkeypatch.setattr(sse_relay, "get_qa_client", lambda: client)
|
||||
|
||||
resp = bff_client.get("/api/qa/info")
|
||||
assert resp.status_code == 503
|
||||
data = resp.json()
|
||||
assert data["ok"] is False
|
||||
assert data["error"]["code"] == "UPSTREAM_DOWN"
|
||||
assert data["error"]["degraded"] is True
|
||||
|
||||
|
||||
def test_info_wrapped_info_field_tolerated(bff_client, monkeypatch):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return _json_response({"info": UPSTREAM_INFO})
|
||||
|
||||
client = _transport_handler(handler)
|
||||
monkeypatch.setattr(sse_relay, "get_qa_client", lambda: client)
|
||||
|
||||
resp = bff_client.get("/api/qa/info")
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["name"] == UPSTREAM_INFO["name"]
|
||||
|
||||
|
||||
def test_chat_ok_normalized(bff_client, monkeypatch):
|
||||
calls = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
calls.append(request)
|
||||
assert request.url.path.endswith("/api/v1/qa-v1-1-73/chat")
|
||||
return _json_response({"answer": "some answer", "sources": []})
|
||||
|
||||
client = _transport_handler(handler)
|
||||
monkeypatch.setattr(sse_relay, "get_qa_client", lambda: client)
|
||||
|
||||
body = {"question": "水稻SPL？", "use_llm": True, "chain_top_k": 3}
|
||||
resp = bff_client.post("/api/qa/chat", json=body)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["mode"] == "chat"
|
||||
assert data["answer"] == "some answer"
|
||||
# 请求体透传，不增删字段
|
||||
sent = json.loads(calls[0].content)
|
||||
assert sent == body
|
||||
|
||||
|
||||
def test_chat_invalid_citation_retries_once(bff_client, monkeypatch):
|
||||
calls = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
calls.append(request)
|
||||
if len(calls) == 1:
|
||||
return _json_response({"detail": {"error_code": "invalid_llm_citation"}}, status=402)
|
||||
return _json_response({"answer": "retry ok", "sources": []})
|
||||
|
||||
client = _transport_handler(handler)
|
||||
monkeypatch.setattr(sse_relay, "get_qa_client", lambda: client)
|
||||
|
||||
resp = bff_client.post("/api/qa/chat", json={"question": "q"})
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["answer"] == "retry ok"
|
||||
assert len(calls) == 2
|
||||
|
||||
|
||||
def test_chat_invalid_citation_retry_still_fails_502(bff_client, monkeypatch):
|
||||
calls = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
calls.append(request)
|
||||
return _json_response({"detail": {"error_code": "invalid_llm_citation"}}, status=402)
|
||||
|
||||
client = _transport_handler(handler)
|
||||
monkeypatch.setattr(sse_relay, "get_qa_client", lambda: client)
|
||||
|
||||
resp = bff_client.post("/api/qa/chat", json={"question": "q"})
|
||||
assert resp.status_code == 502
|
||||
data = resp.json()
|
||||
assert data["error"]["code"] == "INVALID_CITATION"
|
||||
assert data["error"]["message"] == "引用校验失败，请重试"
|
||||
assert len(calls) == 2 # 重试仅一次
|
||||
|
||||
|
||||
def test_chat_upstream_500_no_retry(bff_client, monkeypatch):
|
||||
calls = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
calls.append(request)
|
||||
return _json_response({"detail": "boom"}, status=500)
|
||||
|
||||
client = _transport_handler(handler)
|
||||
monkeypatch.setattr(sse_relay, "get_qa_client", lambda: client)
|
||||
|
||||
resp = bff_client.post("/api/qa/chat", json={"question": "q"})
|
||||
assert resp.status_code == 502
|
||||
assert resp.json()["error"]["code"] == "UPSTREAM_ERROR"
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
def test_search_forces_use_llm_false_and_mode_search(bff_client, monkeypatch):
|
||||
calls = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
calls.append(request)
|
||||
assert request.url.path.endswith("/api/v1/qa-v1-1-73/search")
|
||||
return _json_response({"answer": "SHOULD BE WIPED", "sources": [{"title": "T"}]})
|
||||
|
||||
client = _transport_handler(handler)
|
||||
monkeypatch.setattr(sse_relay, "get_qa_client", lambda: client)
|
||||
|
||||
resp = bff_client.post("/api/qa/search", json={"question": "q", "use_llm": True})
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["mode"] == "search"
|
||||
assert data["answer"] == ""
|
||||
sent = json.loads(calls[0].content)
|
||||
assert sent["use_llm"] is False # 强制
|
||||
assert sent["question"] == "q" # 其余透传
|
||||
|
||||
|
||||
def test_search_no_retry_on_invalid_citation(bff_client, monkeypatch):
|
||||
calls = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
calls.append(request)
|
||||
return _json_response({"detail": {"error_code": "invalid_llm_citation"}}, status=402)
|
||||
|
||||
client = _transport_handler(handler)
|
||||
monkeypatch.setattr(sse_relay, "get_qa_client", lambda: client)
|
||||
|
||||
resp = bff_client.post("/api/qa/search", json={"question": "q"})
|
||||
assert resp.status_code == 502
|
||||
assert len(calls) == 1 # 不重试
|
||||
|
||||
|
||||
def test_chat_upstream_down_503(bff_client, monkeypatch):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
raise httpx.ConnectError("down", request=request)
|
||||
|
||||
client = _transport_handler(handler)
|
||||
monkeypatch.setattr(sse_relay, "get_qa_client", lambda: client)
|
||||
|
||||
resp = bff_client.post("/api/qa/chat", json={"question": "q"})
|
||||
assert resp.status_code == 503
|
||||
assert resp.json()["error"]["code"] == "UPSTREAM_DOWN"
|
||||
|
||||
|
||||
def test_chat_upstream_read_timeout_504(bff_client, monkeypatch):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
raise httpx.ReadTimeout("slow", request=request)
|
||||
|
||||
client = _transport_handler(handler)
|
||||
monkeypatch.setattr(sse_relay, "get_qa_client", lambda: client)
|
||||
|
||||
resp = bff_client.post("/api/qa/chat", json={"question": "q"})
|
||||
assert resp.status_code == 504
|
||||
assert resp.json()["error"]["code"] == "TIMEOUT"
|
||||
|
||||
|
||||
def test_chat_stream_headers_and_relay(bff_client, monkeypatch):
|
||||
"""SSE 端到端：MockTransport 吐 SSE 字节，断言契约帧与响应头。"""
|
||||
upstream_bytes = (
|
||||
b'data: {"stage":"entities_identified","details":{"entities":["OsSPL14"]}}\n\n'
|
||||
b'event: complete\ndata: {"answer":"A","sources":[]}\n\n'
|
||||
)
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.url.path.endswith("/api/v1/qa-v1-1-73/chat/stream")
|
||||
return httpx.Response(200, stream=httpx.ByteStream(upstream_bytes))
|
||||
|
||||
client = _transport_handler(handler)
|
||||
monkeypatch.setattr(sse_relay, "get_qa_client", lambda: client)
|
||||
|
||||
with bff_client.stream("POST", "/api/qa/chat/stream", json={"question": "q"}) as resp:
|
||||
assert resp.status_code == 200
|
||||
assert resp.headers["content-type"].startswith("text/event-stream")
|
||||
assert resp.headers["cache-control"] == "no-cache"
|
||||
assert resp.headers["x-accel-buffering"] == "no"
|
||||
body = "".join(resp.iter_text())
|
||||
|
||||
assert "event: progress" in body
|
||||
assert '"phase": "planning"' in body
|
||||
assert "event: complete" in body
|
||||
assert "event: error" not in body
|
||||
|
||||
|
||||
def test_chat_stream_invalid_citation_retry_then_error(bff_client, monkeypatch):
|
||||
"""上游头阶段 402 → 重试一次 → 仍 402 → event: error。"""
|
||||
calls = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
calls.append(request)
|
||||
return _json_response({"detail": {"error_code": "invalid_llm_citation"}}, status=402)
|
||||
|
||||
client = _transport_handler(handler)
|
||||
monkeypatch.setattr(sse_relay, "get_qa_client", lambda: client)
|
||||
|
||||
with bff_client.stream("POST", "/api/qa/chat/stream", json={"question": "q"}) as resp:
|
||||
body = "".join(resp.iter_text())
|
||||
|
||||
assert "event: error" in body
|
||||
assert '"code": "INVALID_CITATION"' in body
|
||||
assert len(calls) == 2
|
||||
@@ -0,0 +1,255 @@
|
||||
"""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
|
||||
@@ -0,0 +1,325 @@
|
||||
"""services/extraction.py —— 抽取任务上游代理（8000 主 + 7900 回落，API.md §10）。
|
||||
|
||||
- 默认 `RAG_BASE_URL`（http://127.0.0.1:8000），每次请求失败回落
|
||||
`RAG_ALT_BASE_URL`（http://127.0.0.1:7900）。
|
||||
- 回落失败状态不缓存超过 30s。
|
||||
- progress_pct 由 BFF 按 stage 派生：queued=5 / mineru=35 / llm=75 /
|
||||
completed=100 / failed=null（上游没有进度数值字段）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
from typing import Any, AsyncIterator
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_UPSTREAM_PREFIX = "/api/v1/methodology"
|
||||
_DEFAULT_PRIMARY = "http://127.0.0.1:8000"
|
||||
_DEFAULT_ALT = "http://127.0.0.1:7900"
|
||||
|
||||
PRIMARY_BASE_URL = os.environ.get("RAG_BASE_URL", _DEFAULT_PRIMARY)
|
||||
ALT_BASE_URL = os.environ.get("RAG_ALT_BASE_URL", _DEFAULT_ALT)
|
||||
|
||||
# 回落失败状态缓存上限（API.md §10：不缓存失败状态超过 30s）
|
||||
ALT_FAILURE_TTL = 30.0
|
||||
|
||||
# 超时阶梯（API.md §10）
|
||||
READ_TIMEOUT = 30.0
|
||||
UPLOAD_READ_TIMEOUT = 600.0
|
||||
|
||||
# 上传限制（FR-EXTRACT-004）
|
||||
MAX_PDF_BYTES = 50 * 1024 * 1024 # 50MB
|
||||
_SPOOL_MAX_MEMORY = 5 * 1024 * 1024 # spool 驻内存上限，超出落临时文件
|
||||
|
||||
_UPLOAD_ROUTE = "/extract-upload"
|
||||
|
||||
# 上游成功上传形态：{success, job_id, status_url}；message 由 BFF 附加
|
||||
_SUCCESS_UPLOAD_MESSAGE = "已提交抽取任务"
|
||||
|
||||
_PROGRESS_BY_STAGE: dict[str, int | None] = {
|
||||
"queued": 5,
|
||||
"mineru": 35,
|
||||
"llm": 75,
|
||||
"completed": 100,
|
||||
"failed": None,
|
||||
}
|
||||
|
||||
|
||||
def derive_progress_pct(stage: str | None) -> int | None:
|
||||
"""按 stage 派生 progress_pct（API.md §10）。未知 stage 返回 None。"""
|
||||
return _PROGRESS_BY_STAGE.get(stage if stage is not None else "")
|
||||
|
||||
|
||||
class ExtractionUnavailable(Exception):
|
||||
"""两个抽取上游（8000 + 7900）均不可达。"""
|
||||
|
||||
def __init__(self, message: str, *, cause: BaseException | None = None) -> None:
|
||||
super().__init__(message)
|
||||
self.cause = cause
|
||||
|
||||
|
||||
class ExtractionUpstreamError(Exception):
|
||||
"""两个上游均可连通但都返回了错误（5xx，如 MinerU 不可用）。"""
|
||||
|
||||
def __init__(self, message: str) -> None:
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class ExtractionService:
|
||||
"""每次请求先打 primary（8000）；连接失败时回落 alt（7900）。
|
||||
|
||||
失败回落状态记 `_alt_dead_until` 时间戳（30s 内不再试 alt），
|
||||
任何 alt 成功都会清掉该状态。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
primary_base_url: str = PRIMARY_BASE_URL,
|
||||
alt_base_url: str = ALT_BASE_URL,
|
||||
) -> None:
|
||||
self._primary = primary_base_url.rstrip("/")
|
||||
self._alt = alt_base_url.rstrip("/")
|
||||
self._alt_dead_until: float = 0.0
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
@property
|
||||
def client(self) -> httpx.AsyncClient:
|
||||
if self._client is None:
|
||||
self._client = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(READ_TIMEOUT, connect=3.0)
|
||||
)
|
||||
return self._client
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client is not None:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
# ---------- 回落机制 ----------
|
||||
|
||||
def _alt_usable(self) -> bool:
|
||||
return time.monotonic() >= self._alt_dead_until
|
||||
|
||||
def _mark_alt_dead(self) -> None:
|
||||
self._alt_dead_until = time.monotonic() + ALT_FAILURE_TTL
|
||||
|
||||
def _mark_alt_alive(self) -> None:
|
||||
self._alt_dead_until = 0.0
|
||||
|
||||
async def _request(
|
||||
self,
|
||||
path: str,
|
||||
*,
|
||||
method: str = "GET",
|
||||
read_timeout: float = READ_TIMEOUT,
|
||||
**kwargs: Any,
|
||||
) -> httpx.Response:
|
||||
"""primary → alt 回落。连接失败回落 alt；5xx 时也回落 alt，
|
||||
两个上游都 5xx 则抛 ExtractionUpstreamError（含最后一个 5xx 响应）。
|
||||
"""
|
||||
client = self.client
|
||||
primary_failed: Exception | None = None
|
||||
primary_status: int | None = None
|
||||
|
||||
try:
|
||||
resp = await client.request(
|
||||
method,
|
||||
f"{self._primary}{_UPSTREAM_PREFIX}{path}",
|
||||
timeout=httpx.Timeout(read_timeout, connect=3.0),
|
||||
**kwargs,
|
||||
)
|
||||
if resp.status_code < 500:
|
||||
return resp
|
||||
primary_status = resp.status_code
|
||||
primary_failed = RuntimeError(f"primary 返回 {resp.status_code}")
|
||||
except (httpx.HTTPError, httpx.TimeoutException) as exc:
|
||||
primary_failed = exc
|
||||
|
||||
# ---- 回落 alt ----
|
||||
if not self._alt_usable():
|
||||
raise ExtractionUnavailable(
|
||||
"抽取服务（8000）不可用，且 7900 最近不可用，稍后自动重试",
|
||||
cause=primary_failed,
|
||||
)
|
||||
try:
|
||||
resp = await client.request(
|
||||
method,
|
||||
f"{self._alt}{_UPSTREAM_PREFIX}{path}",
|
||||
timeout=httpx.Timeout(read_timeout, connect=3.0),
|
||||
**kwargs,
|
||||
)
|
||||
if resp.status_code < 500:
|
||||
self._mark_alt_alive()
|
||||
return resp
|
||||
self._mark_alt_dead()
|
||||
raise ExtractionUpstreamError(
|
||||
f"两个抽取上游均返回错误：8000 状态 {primary_status}，"
|
||||
f"7900 状态 {resp.status_code}"
|
||||
)
|
||||
except (httpx.HTTPError, httpx.TimeoutException) as exc:
|
||||
self._mark_alt_dead()
|
||||
raise ExtractionUnavailable(
|
||||
"抽取服务不可用（8000 与 7900 均不可达）",
|
||||
cause=exc,
|
||||
) from exc
|
||||
|
||||
# ---------- 归一化 ----------
|
||||
|
||||
@staticmethod
|
||||
def normalize_job(item: dict[str, Any]) -> dict[str, Any]:
|
||||
"""列表 item 归一化（API.md §10 列表形态）。
|
||||
|
||||
title = source_filename 优先 → paper_id → job_id。
|
||||
"""
|
||||
job_id = item.get("job_id")
|
||||
title = item.get("source_filename") or item.get("paper_id") or job_id
|
||||
stage = item.get("stage")
|
||||
return {
|
||||
"job_id": job_id,
|
||||
"title": title,
|
||||
"status": item.get("status"),
|
||||
"stage": stage,
|
||||
"substage": item.get("substage"),
|
||||
"progress_pct": derive_progress_pct(stage),
|
||||
"message": item.get("message"),
|
||||
"paper_id": item.get("paper_id"),
|
||||
"family": item.get("family"),
|
||||
"source_filename": item.get("source_filename"),
|
||||
"stage_count": item.get("stage_count"),
|
||||
"species_count": item.get("species_count"),
|
||||
"figure_table_caption_count": item.get("figure_table_caption_count"),
|
||||
"figure_table_stage_count": item.get("figure_table_stage_count"),
|
||||
"updated_at": item.get("updated_at"),
|
||||
"config_url": item.get("config_url"),
|
||||
"methodology_plan_url": item.get("methodology_plan_url"),
|
||||
}
|
||||
|
||||
# ---------- 业务操作 ----------
|
||||
|
||||
async def list_jobs(self, limit: int = 20) -> dict[str, Any]:
|
||||
"""GET /extract-jobs（可带 ?limit=，默认 20）。"""
|
||||
resp = await self._request(f"/extract-jobs?limit={limit}")
|
||||
payload = resp.json()
|
||||
return {
|
||||
"jobs": [self.normalize_job(item) for item in payload.get("jobs", [])],
|
||||
"count": len(payload.get("jobs", [])),
|
||||
}
|
||||
|
||||
async def get_job(self, job_id: str) -> dict[str, Any]:
|
||||
"""GET /extract-jobs/{id}：扁平详情归一化（API.md §10）。
|
||||
|
||||
plan：completed 时请求 /plan；未完成时上游 409 → 置 null 不报错。
|
||||
"""
|
||||
resp = await self._request(f"/extract-jobs/{job_id}")
|
||||
item = resp.json()
|
||||
|
||||
status = item.get("status")
|
||||
stage = item.get("stage")
|
||||
job_id_val = item.get("job_id") or job_id
|
||||
title = (
|
||||
item.get("source_filename")
|
||||
or item.get("paper_id")
|
||||
or job_id_val
|
||||
)
|
||||
|
||||
plan: dict[str, Any] | None = None
|
||||
if status == "completed":
|
||||
try:
|
||||
plan_resp = await self._request(f"/extract-jobs/{job_id}/plan")
|
||||
except ExtractionUnavailable as exc:
|
||||
logger.warning(
|
||||
"extraction: job %s 的 plan 请求失败（上游不可用）：%s", job_id, exc
|
||||
)
|
||||
else:
|
||||
if plan_resp.status_code == 409:
|
||||
# 上游认为未完成（与扁平详情不一致），契约要求置 null 不报错
|
||||
plan = None
|
||||
else:
|
||||
plan = plan_resp.json()
|
||||
|
||||
return {
|
||||
"job_id": job_id_val,
|
||||
"title": title,
|
||||
"status": status,
|
||||
"stage": stage,
|
||||
"progress_pct": derive_progress_pct(stage),
|
||||
"message": item.get("message"),
|
||||
"paper_id": item.get("paper_id"),
|
||||
"family": item.get("family"),
|
||||
"stats": {
|
||||
"stage_count": item.get("stage_count"),
|
||||
"species_count": item.get("species_count"),
|
||||
"figure_table_caption_count": item.get("figure_table_caption_count"),
|
||||
"figure_table_stage_count": item.get("figure_table_stage_count"),
|
||||
},
|
||||
"detail": item,
|
||||
"plan": plan,
|
||||
}
|
||||
|
||||
async def _passthrough(self, path: str) -> dict[str, Any]:
|
||||
"""processing-status / progress / summary / run-log：JSON 透传。
|
||||
|
||||
上游 404 视为空数据（rows 为空属正常现象，前端显示空态）。
|
||||
"""
|
||||
resp = await self._request(path)
|
||||
return resp.json()
|
||||
|
||||
async def upload(
|
||||
self,
|
||||
*,
|
||||
filename: str,
|
||||
stream: AsyncIterator[bytes] | bytes,
|
||||
content_type: str | None,
|
||||
paper_id: str = "uploaded_paper",
|
||||
family: str = "",
|
||||
output_name: str = "",
|
||||
) -> dict[str, Any]:
|
||||
"""POST /extract-upload：multipart 转发（内存有界）。
|
||||
|
||||
httpx 的 multipart 编码只接受同步 file-like/bytes，异步流会先
|
||||
spool 进 SpooledTemporaryFile（≤5MB 驻内存，超出落磁盘），
|
||||
因此不会整文件读进内存。
|
||||
返回上游形态 + message:"已提交抽取任务"。
|
||||
"""
|
||||
created = not isinstance(stream, (bytes, bytearray))
|
||||
if created:
|
||||
spooled: Any = tempfile.SpooledTemporaryFile(max_size=_SPOOL_MAX_MEMORY)
|
||||
else:
|
||||
spooled = stream
|
||||
try:
|
||||
if created:
|
||||
async for chunk in stream:
|
||||
await anyio.to_thread.run_sync(spooled.write, chunk)
|
||||
spooled.seek(0)
|
||||
resp = await self._request(
|
||||
_UPLOAD_ROUTE,
|
||||
method="POST",
|
||||
read_timeout=UPLOAD_READ_TIMEOUT,
|
||||
data={
|
||||
"paper_id": paper_id,
|
||||
"family": family,
|
||||
"output_name": output_name,
|
||||
},
|
||||
files={"file": (filename, spooled, content_type)},
|
||||
)
|
||||
finally:
|
||||
if created:
|
||||
spooled.close()
|
||||
payload = resp.json()
|
||||
payload["message"] = _SUCCESS_UPLOAD_MESSAGE
|
||||
return payload
|
||||
|
||||
|
||||
# 模块级单例：路由层共享（app 关闭时 close）。
|
||||
service = ExtractionService()
|
||||
|
||||
|
||||
async def close_extraction() -> None:
|
||||
"""app shutdown 钩子。"""
|
||||
await service.close()
|
||||
@@ -0,0 +1,646 @@
|
||||
"""services/neo4j.py —— Neo4j Bolt 直连（官方 neo4j 驱动，全异步）。
|
||||
|
||||
选型与凭据依据 DISCOVERY.md §2：
|
||||
- 首选 neo4j-plant-srt-direct-v1-full（bolt 127.0.0.1:18788，db `neo4j`，
|
||||
用户 neo4j / 密码在 backend/.env），schema canonical，QA 7920 生产使用。
|
||||
- 回退 neo4j-plant-srt-pc2025-v3-20260817-normalized（bolt 127.0.0.1:19788），
|
||||
schema 不同（Entity/ASSERTS），本波未做 Cypher 适配（TODO）。
|
||||
- 再次 plantai-protected-neo4j-1（bolt 127.0.0.1:17687），空库。
|
||||
|
||||
所有本模块异常统一抛 :class:`Neo4jUnavailable`，由路由层转 503 降级信封。
|
||||
凭据只存在于本文件（或 env 覆盖），绝不下发浏览器。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from neo4j import AsyncDriver, AsyncGraphDatabase, exceptions as neo4j_exc
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 可选 env 覆盖（测试/部署）；未设置时使用 DISCOVERY.md 实测凭据。
|
||||
_ENV_PREFIX = "BFF_NEO4J_"
|
||||
|
||||
PROBE_CACHE_TTL = 5.0
|
||||
|
||||
# 连接/查询超时（秒）。entity 详情与一跳邻域走全量属性，超时给得宽些。
|
||||
CONNECT_TIMEOUT = 3.0
|
||||
QUERY_TIMEOUT = 10.0
|
||||
NEIGHBORHOOD_QUERY_TIMEOUT = 30.0
|
||||
OVERVIEW_QUERY_TIMEOUT = 60.0 # 2026-08-25 PM 修复：1644 万边 COUNT 聚合冷缓存 >10s
|
||||
|
||||
NAME_INDEX = "srt2_canonical_entity_name_norm"
|
||||
FULLTEXT_INDEX = None # DISCOVERY.md §2：任何实例都不存在全文索引，且不得创建
|
||||
|
||||
# status 展示名（须与 API.md §6 fixture 一致）
|
||||
DIRECT_ACTIVE_NAME = "neo4j-plant-srt-direct-v1-full"
|
||||
NORMALIZED_ACTIVE_NAME = "neo4j-plant-srt-pc2025-v3-20260817-normalized"
|
||||
PROTECTED_ACTIVE_NAME = "plantai-protected-neo4j-1"
|
||||
|
||||
_CANONICAL_LABELS = ["SRT2CanonicalEntity", "SRT2CanonicalAssertion", "SRT2Paper"]
|
||||
_CANONICAL_REL_TYPES = ["SRT2_CANONICAL_ASSERTS", "SRT2_SUPPORTED_BY"]
|
||||
|
||||
|
||||
class Neo4jUnavailable(Exception):
|
||||
"""Neo4j 上游不可用（连接失败/认证失败/超时/Cypher 错误统一转此异常）。"""
|
||||
|
||||
def __init__(self, message: str, *, cause: BaseException | None = None) -> None:
|
||||
super().__init__(message)
|
||||
self.cause = cause
|
||||
|
||||
|
||||
class _Instance:
|
||||
"""一个候选实例。本波只实现 canonical schema（direct 实例）的 Cypher。"""
|
||||
|
||||
__slots__ = ("name", "uri", "auth", "database", "schema", "label", "rel")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
uri: str,
|
||||
user: str,
|
||||
password: str,
|
||||
database: str,
|
||||
schema: str,
|
||||
label: str | None,
|
||||
rel: str | None,
|
||||
) -> None:
|
||||
self.name = name
|
||||
self.uri = uri
|
||||
self.auth = (user, password)
|
||||
self.database = database
|
||||
self.schema = schema
|
||||
self.label = label
|
||||
self.rel = rel
|
||||
|
||||
|
||||
def _env(name: str, default: str) -> str:
|
||||
return os.environ.get(f"{_ENV_PREFIX}{name}", default)
|
||||
|
||||
|
||||
_INSTANCES: list[_Instance] = [
|
||||
_Instance(
|
||||
DIRECT_ACTIVE_NAME,
|
||||
_env("URI_DIRECT", "bolt://127.0.0.1:18788"),
|
||||
"neo4j",
|
||||
_env("PASSWORD_DIRECT", ""),
|
||||
"neo4j",
|
||||
"canonical",
|
||||
"SRT2CanonicalEntity",
|
||||
"SRT2_CANONICAL_ASSERTS",
|
||||
),
|
||||
_Instance(
|
||||
NORMALIZED_ACTIVE_NAME,
|
||||
_env("URI_NORMALIZED", "bolt://127.0.0.1:19788"),
|
||||
"neo4j",
|
||||
_env("PASSWORD_NORMALIZED", ""),
|
||||
"neo4j",
|
||||
"normalized",
|
||||
None,
|
||||
None,
|
||||
),
|
||||
_Instance(
|
||||
PROTECTED_ACTIVE_NAME,
|
||||
_env("URI_PROTECTED", "bolt://127.0.0.1:17687"),
|
||||
"neo4j",
|
||||
_env("PASSWORD_PROTECTED", ""),
|
||||
"neo4j",
|
||||
"empty",
|
||||
None,
|
||||
None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class Neo4jService:
|
||||
"""生命周期：app 启动时 create，关闭时 close()。health 路由与 kg 路由共用。
|
||||
|
||||
内部为每个候选实例缓存一个 AsyncDriver（懒建）。实例选择只发生在
|
||||
显式调用 select()（probe/kg_status）或模块加载时：每个请求的处理函数
|
||||
首先调用 select()，得到当前实例（未探测时默认首选，惰性回退）。
|
||||
切换由下一次显式 probe 触发，探测结果缓存 5s。
|
||||
"""
|
||||
|
||||
def __init__(self, instances: list[_Instance] | None = None) -> None:
|
||||
self._instances = instances if instances is not None else _INSTANCES
|
||||
self._drivers: dict[str, AsyncDriver] = {}
|
||||
self._active: _Instance = self._instances[0]
|
||||
self._selected = False # 是否已成功选出一个实例
|
||||
self._last_probe: float = 0.0
|
||||
self._last_probe_ok: bool = False
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
# ---------- 生命周期 ----------
|
||||
|
||||
async def close(self) -> None:
|
||||
for driver in self._drivers.values():
|
||||
await driver.close()
|
||||
self._drivers.clear()
|
||||
|
||||
def _driver_for(self, instance: _Instance) -> AsyncDriver:
|
||||
driver = self._drivers.get(instance.name)
|
||||
if driver is None:
|
||||
driver = AsyncGraphDatabase.driver(
|
||||
instance.uri,
|
||||
auth=instance.auth,
|
||||
connection_timeout=CONNECT_TIMEOUT,
|
||||
connection_acquisition_timeout=CONNECT_TIMEOUT,
|
||||
)
|
||||
self._drivers[instance.name] = driver
|
||||
return driver
|
||||
|
||||
# ---------- 实例选择 ----------
|
||||
|
||||
async def select(self) -> _Instance:
|
||||
"""返回当前可用实例。
|
||||
|
||||
顺序探测（5s 结果缓存）：direct → normalized → protected。
|
||||
首次调用若未探测成功则依次尝试探测，全失败抛 Neo4jUnavailable。
|
||||
"""
|
||||
if self._selected:
|
||||
return self._active
|
||||
now = time.monotonic()
|
||||
if now - self._last_probe < PROBE_CACHE_TTL:
|
||||
if self._last_probe_ok:
|
||||
return self._active
|
||||
# 探测失败结果同样缓存 5s，避免每个请求都打 3 个实例
|
||||
raise Neo4jUnavailable("Neo4j 不可用（最近探测失败，结果缓存 5s）")
|
||||
await self._probe()
|
||||
if self._selected:
|
||||
return self._active
|
||||
raise Neo4jUnavailable("所有 Neo4j 候选实例均不可达")
|
||||
|
||||
async def _probe(self) -> None:
|
||||
async with self._lock:
|
||||
now = time.monotonic()
|
||||
if now - self._last_probe < PROBE_CACHE_TTL:
|
||||
if not self._last_probe_ok:
|
||||
raise Neo4jUnavailable("Neo4j 不可用（最近探测失败，结果缓存 5s）")
|
||||
return
|
||||
self._selected = False
|
||||
last_error: BaseException | None = None
|
||||
for instance in self._instances:
|
||||
try:
|
||||
await self._check_instance(instance)
|
||||
except Neo4jUnavailable as exc:
|
||||
last_error = exc
|
||||
continue
|
||||
self._active = instance
|
||||
self._selected = True
|
||||
break
|
||||
self._last_probe = now
|
||||
self._last_probe_ok = self._selected
|
||||
if not self._selected:
|
||||
raise Neo4jUnavailable(
|
||||
f"所有 Neo4j 候选实例均不可达：{last_error}", cause=last_error
|
||||
)
|
||||
|
||||
async def _check_instance(self, instance: _Instance) -> None:
|
||||
"""连通性检查：TCP 预检 + `RETURN 1`，整体限时 CONNECT_TIMEOUT。
|
||||
|
||||
先做 TCP 连通预检（毫秒级）：端口没开时直接判死，避免官方驱动
|
||||
的指数退避重试拖慢健康探测。TCP 通但 Bolt 协议失败时才走
|
||||
wait_for + 驱动（仍限时 3s）。
|
||||
"""
|
||||
if not await self._tcp_reachable(instance.uri):
|
||||
raise Neo4jUnavailable(f"{instance.name} 连接失败（TCP 不可达）")
|
||||
driver = self._driver_for(instance)
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
driver.execute_query(
|
||||
"RETURN 1 AS x", database_=instance.database, routing_="r"
|
||||
),
|
||||
timeout=CONNECT_TIMEOUT,
|
||||
)
|
||||
except asyncio.TimeoutError as exc:
|
||||
raise Neo4jUnavailable(f"{instance.name} 连接超时", cause=exc) from exc
|
||||
except neo4j_exc.DriverError as exc:
|
||||
raise Neo4jUnavailable(
|
||||
f"{instance.name} 连接失败：{exc.__class__.__name__}", cause=exc
|
||||
) from exc
|
||||
|
||||
@staticmethod
|
||||
async def _tcp_reachable(uri: str) -> bool:
|
||||
"""bolt://host:port 的 TCP 连通预检（超时 0.5s）。"""
|
||||
try:
|
||||
host, port = uri.split("://", 1)[1].rsplit(":", 1)
|
||||
except (ValueError, IndexError):
|
||||
return False
|
||||
try:
|
||||
reader, writer = await asyncio.wait_for(
|
||||
asyncio.open_connection(host, int(port)), timeout=0.5
|
||||
)
|
||||
except (OSError, asyncio.TimeoutError, ValueError):
|
||||
return False
|
||||
writer.close()
|
||||
try:
|
||||
await writer.wait_closed()
|
||||
except (OSError, AttributeError):
|
||||
pass
|
||||
return True
|
||||
|
||||
# ---------- 通用查询 ----------
|
||||
|
||||
async def _query(
|
||||
self,
|
||||
cypher: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
*,
|
||||
timeout: float = QUERY_TIMEOUT,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""按当前实例执行查询，返回 records 列表（每条为 dict）。"""
|
||||
instance = await self.select()
|
||||
driver = self._driver_for(instance)
|
||||
try:
|
||||
records, _, _ = await asyncio.wait_for(
|
||||
driver.execute_query(
|
||||
cypher,
|
||||
parameters_=params or {},
|
||||
database_=instance.database,
|
||||
routing_="r",
|
||||
),
|
||||
timeout=timeout,
|
||||
)
|
||||
except asyncio.TimeoutError as exc:
|
||||
raise Neo4jUnavailable(
|
||||
f"Neo4j 查询超时（>{timeout}s）", cause=exc
|
||||
) from exc
|
||||
except neo4j_exc.DriverError as exc:
|
||||
raise Neo4jUnavailable(
|
||||
f"Neo4j 查询失败：{exc.__class__.__name__}", cause=exc
|
||||
) from exc
|
||||
return [dict(record) for record in records]
|
||||
|
||||
@staticmethod
|
||||
def _canonical(instance: _Instance) -> _Instance:
|
||||
"""当前仅 canonical schema 的 Cypher 就绪。"""
|
||||
if instance.schema != "canonical":
|
||||
# TODO(A3)：normalized 实例（Entity/ASSERTS，无 name_norm 索引，
|
||||
# 等值 6.81s）与空库实例的 Cypher 适配层，需 PM 决策后实现。
|
||||
raise Neo4jUnavailable(
|
||||
f"Neo4j 实例 {instance.name} 已连通，但其 schema（{instance.schema}）"
|
||||
"的 Cypher 适配尚未实现，暂不可用于图谱查询"
|
||||
)
|
||||
return instance
|
||||
|
||||
# ---------- 健康探测 ----------
|
||||
|
||||
async def probe(self) -> dict[str, Any]:
|
||||
"""连通 + 计数 + 索引检查（/api/health 调用），结果缓存 5s。
|
||||
|
||||
返回 {"ok": bool, "active": str|null, "detail": str|null}，
|
||||
主动探测时即使失败也返回 dict（health 聚合不因此抛错）。
|
||||
"""
|
||||
now = time.monotonic()
|
||||
if now - self._last_probe < PROBE_CACHE_TTL and self._last_probe_ok:
|
||||
return {"ok": True, "active": self._active.name, "detail": None}
|
||||
try:
|
||||
await self._probe()
|
||||
except Neo4jUnavailable as exc:
|
||||
return {"ok": False, "active": None, "detail": str(exc)}
|
||||
|
||||
instance = self._active
|
||||
detail: str | None = None
|
||||
if instance.schema == "canonical":
|
||||
try:
|
||||
counts = await self._counts(instance)
|
||||
indexes = await self._name_indexes(instance)
|
||||
# 成员判断而非全等：同一标签上存在唯一约束索引等其它索引是正常
|
||||
# 现象（DISCOVERY §2.2 实测有 srt2_canonical_entity_id），
|
||||
# 只有 name_norm 索引真正不在列表里才报缺失。
|
||||
if NAME_INDEX not in indexes:
|
||||
detail = f"name_norm 索引缺失：{indexes or '(无任何索引)'}"
|
||||
except Neo4jUnavailable as exc:
|
||||
# 连通了但计数/索引检查失败：不视为整体失败，附 detail
|
||||
detail = f"计数/索引检查失败：{exc}"
|
||||
return {"ok": True, "active": instance.name, "detail": detail}
|
||||
|
||||
async def _counts(self, instance: _Instance) -> dict[str, int]:
|
||||
rows = await self._query(
|
||||
"""
|
||||
CALL {
|
||||
MATCH (e:SRT2CanonicalEntity) RETURN count(e) AS entity_count
|
||||
}
|
||||
CALL {
|
||||
MATCH (a:SRT2CanonicalAssertion) RETURN count(a) AS assertion_count
|
||||
}
|
||||
RETURN entity_count, assertion_count
|
||||
""",
|
||||
timeout=QUERY_TIMEOUT,
|
||||
)
|
||||
row = rows[0] if rows else {}
|
||||
return {
|
||||
"entity_count": row.get("entity_count", 0),
|
||||
"assertion_count": row.get("assertion_count", 0),
|
||||
}
|
||||
|
||||
async def _name_indexes(self, instance: _Instance) -> list[str]:
|
||||
"""该实体标签上的所有索引名（含约束索引），滤掉空名。
|
||||
|
||||
SHOW INDEXES 的 labelsOrTypes 可以是字符串或字符串列表，用
|
||||
`ANY(x IN labelsOrTypes WHERE x = $label)` 兼容两者。
|
||||
"""
|
||||
rows = await self._query(
|
||||
"SHOW INDEXES YIELD name, labelsOrTypes "
|
||||
"WHERE ANY(x IN labelsOrTypes WHERE x = $label) "
|
||||
"RETURN name",
|
||||
{"label": instance.label},
|
||||
timeout=QUERY_TIMEOUT,
|
||||
)
|
||||
return [r["name"] for r in rows if r.get("name")]
|
||||
|
||||
# ---------- /api/kg/status ----------
|
||||
|
||||
async def kg_status(self) -> dict[str, Any]:
|
||||
"""labels / relationship_types / entity_count / assertion_count /
|
||||
name_index / fulltext_index（null，实例无全文索引）。"""
|
||||
instance = self._canonical(await self.select())
|
||||
counts = await self._counts(instance)
|
||||
indexes = await self._name_indexes(instance)
|
||||
name_index = NAME_INDEX if NAME_INDEX in indexes else None
|
||||
return {
|
||||
"ok": True,
|
||||
"active": instance.name,
|
||||
"schema": instance.schema,
|
||||
"labels": _CANONICAL_LABELS,
|
||||
"relationship_types": _CANONICAL_REL_TYPES,
|
||||
"entity_count": counts["entity_count"],
|
||||
"assertion_count": counts["assertion_count"],
|
||||
"name_index": name_index,
|
||||
"fulltext_index": FULLTEXT_INDEX,
|
||||
}
|
||||
|
||||
# ---------- /api/kg/overview ----------
|
||||
|
||||
async def kg_overview(self) -> dict[str, Any]:
|
||||
"""按 entity_type 与 relation_category 计数。
|
||||
|
||||
注意 1644 万边：只做 COUNT 聚合，绝不拉取全量边。
|
||||
"""
|
||||
self._canonical(await self.select())
|
||||
rows = await self._query(
|
||||
"""
|
||||
MATCH (e:SRT2CanonicalEntity)
|
||||
WITH e.entity_type AS t, count(*) AS c
|
||||
WHERE t IS NOT NULL
|
||||
RETURN t AS type, c AS count
|
||||
ORDER BY c DESC, t ASC
|
||||
LIMIT 200
|
||||
""",
|
||||
timeout=OVERVIEW_QUERY_TIMEOUT,
|
||||
)
|
||||
entities = [{"type": r["type"], "count": r["count"]} for r in rows]
|
||||
|
||||
# 关系类别取自边上的 relation_category（属于关系属性，不在边集合里拉边）
|
||||
rows = await self._query(
|
||||
"""
|
||||
MATCH ()-[r:SRT2_CANONICAL_ASSERTS]->()
|
||||
WITH r.relation_category AS cat, count(*) AS c
|
||||
WHERE cat IS NOT NULL
|
||||
RETURN cat AS type, c AS count
|
||||
ORDER BY c DESC, cat ASC
|
||||
LIMIT 200
|
||||
""",
|
||||
timeout=OVERVIEW_QUERY_TIMEOUT,
|
||||
)
|
||||
relations = [{"type": r["type"], "count": r["count"]} for r in rows]
|
||||
return {
|
||||
"entities_by_type": entities,
|
||||
"relations_by_category": relations,
|
||||
}
|
||||
|
||||
# ---------- /api/kg/entities ----------
|
||||
|
||||
async def search_entities(
|
||||
self, q: str, limit: int = 20
|
||||
) -> list[dict[str, Any]]:
|
||||
"""DISCOVERY.md §2.4 Cypher (a)：name_norm = 或 STARTS WITH（toLower 归一化）。
|
||||
|
||||
走 name_norm RANGE 索引（实测 0.09-0.10s）。结果按 name ASC 排序保证稳定。
|
||||
返回 [{id, name, type, support_count}]。
|
||||
"""
|
||||
self._canonical(await self.select())
|
||||
limit = min(max(int(limit), 1), 20)
|
||||
q_norm = q.lower()
|
||||
rows = await self._query(
|
||||
"""
|
||||
MATCH (e:SRT2CanonicalEntity)
|
||||
WHERE e.name_norm = $q_norm OR e.name_norm STARTS WITH $q_norm
|
||||
RETURN e.canonical_id AS canonical_id, e.name AS name,
|
||||
e.entity_type AS entity_type, e.support_count AS support_count
|
||||
ORDER BY e.name ASC
|
||||
LIMIT $limit
|
||||
""",
|
||||
{"q_norm": q_norm, "limit": limit},
|
||||
timeout=QUERY_TIMEOUT,
|
||||
)
|
||||
return [
|
||||
{
|
||||
"id": r.get("canonical_id"),
|
||||
"name": r.get("name"),
|
||||
"type": r.get("entity_type"),
|
||||
"support_count": r.get("support_count"),
|
||||
}
|
||||
for r in rows
|
||||
]
|
||||
|
||||
async def _resolve_id(self, id_or_name: str) -> str | None:
|
||||
"""canonical_id 或 name_norm 定位（DISCOVERY.md §2.4 Cypher b）。"""
|
||||
if id_or_name.startswith("cent_"):
|
||||
rows = await self._query(
|
||||
"MATCH (e:SRT2CanonicalEntity {canonical_id: $id}) "
|
||||
"RETURN e.canonical_id AS canonical_id",
|
||||
{"id": id_or_name},
|
||||
timeout=QUERY_TIMEOUT,
|
||||
)
|
||||
else:
|
||||
rows = await self._query(
|
||||
"MATCH (e:SRT2CanonicalEntity {name_norm: $name_norm}) "
|
||||
"RETURN e.canonical_id AS canonical_id LIMIT 5",
|
||||
{"name_norm": id_or_name.lower()},
|
||||
timeout=QUERY_TIMEOUT,
|
||||
)
|
||||
return rows[0].get("canonical_id") if rows else None
|
||||
|
||||
# ---------- /api/kg/entities/{id} ----------
|
||||
|
||||
async def entity(self, id_or_name: str) -> dict[str, Any] | None:
|
||||
"""按 canonical_id 或 name_norm 取实体全量属性；未命中返回 None。"""
|
||||
canonical_id = await self._resolve_id(id_or_name)
|
||||
if canonical_id is None:
|
||||
return None
|
||||
rows = await self._query(
|
||||
"MATCH (e:SRT2CanonicalEntity {canonical_id: $id}) "
|
||||
"RETURN e {.*} AS node",
|
||||
{"id": canonical_id},
|
||||
timeout=QUERY_TIMEOUT,
|
||||
)
|
||||
if not rows:
|
||||
return None
|
||||
props = rows[0].get("node") or {}
|
||||
return {
|
||||
"id": props.get("canonical_id"),
|
||||
"name": props.get("name"),
|
||||
"type": props.get("entity_type"),
|
||||
"properties": props,
|
||||
}
|
||||
|
||||
# ---------- /api/kg/entities/{id}/neighborhood ----------
|
||||
|
||||
async def neighborhood(
|
||||
self, id_or_name: str, depth: int = 1, limit: int = 200
|
||||
) -> dict[str, Any]:
|
||||
"""有界邻域扩展（DISCOVERY.md §2.4 Cypher (c) 思路）。
|
||||
|
||||
- depth ∈ {1,2}，默认 1；limit ≤ 200。
|
||||
- 每跳按 paper_count/support_count 降序截断（防 hub 吃满预算）。
|
||||
- OPTIONAL MATCH 保证孤立节点不丢。
|
||||
- 无向匹配（-()-），entity 按 canonical_id 去重。
|
||||
- depth=1 绝不执行第二跳（两跳实测约 13s，不默认开启）。
|
||||
- 返回 API.md §6 的 {center, center_name, depth, nodes, edges}。
|
||||
"""
|
||||
self._canonical(await self.select()) # 非 canonical 实例在此抛 Neo4jUnavailable
|
||||
depth = depth if depth in (1, 2) else 1
|
||||
limit = min(max(int(limit), 1), 200)
|
||||
canonical_id = await self._resolve_id(id_or_name)
|
||||
if canonical_id is None:
|
||||
return {"center": None, "center_name": None, "depth": depth, "nodes": [], "edges": []}
|
||||
|
||||
params = {"seed": canonical_id, "max_nodes": limit, "max_edges": limit}
|
||||
if depth == 1:
|
||||
cypher = """
|
||||
MATCH (s:SRT2CanonicalEntity {canonical_id: $seed})
|
||||
OPTIONAL MATCH (s)-[r1:SRT2_CANONICAL_ASSERTS]-(n1:SRT2CanonicalEntity)
|
||||
WITH s, r1, n1
|
||||
ORDER BY coalesce(r1.paper_count, 0) DESC, coalesce(r1.support_count, 0) DESC
|
||||
WITH s,
|
||||
collect(DISTINCT n1)[..$max_nodes] AS level1,
|
||||
[r IN collect(DISTINCT r1)[..$max_edges] | r {.*,
|
||||
source_canonical_id: CASE WHEN startNode(r).canonical_id = s.canonical_id
|
||||
THEN s.canonical_id ELSE endNode(r).canonical_id END,
|
||||
target_canonical_id: CASE WHEN startNode(r).canonical_id = s.canonical_id
|
||||
THEN endNode(r).canonical_id ELSE s.canonical_id END}] AS edges1
|
||||
RETURN s {.*} AS start,
|
||||
[n IN level1 | n {.*}] AS nodes,
|
||||
edges1 AS edges
|
||||
"""
|
||||
else:
|
||||
cypher = """
|
||||
MATCH (s:SRT2CanonicalEntity {canonical_id: $seed})
|
||||
OPTIONAL MATCH (s)-[r1:SRT2_CANONICAL_ASSERTS]-(n1:SRT2CanonicalEntity)
|
||||
WITH s, r1, n1
|
||||
ORDER BY coalesce(r1.paper_count, 0) DESC, coalesce(r1.support_count, 0) DESC
|
||||
WITH s,
|
||||
collect(DISTINCT n1)[..$max_nodes] AS level1,
|
||||
[r IN collect(DISTINCT r1)[..$max_edges] | r {.*,
|
||||
source_canonical_id: CASE WHEN startNode(r).canonical_id = s.canonical_id
|
||||
THEN s.canonical_id ELSE endNode(r).canonical_id END,
|
||||
target_canonical_id: CASE WHEN startNode(r).canonical_id = s.canonical_id
|
||||
THEN endNode(r).canonical_id ELSE s.canonical_id END}] AS edges1
|
||||
OPTIONAL MATCH (n1)-[r2:SRT2_CANONICAL_ASSERTS]-(n2:SRT2CanonicalEntity)
|
||||
WHERE n1 IN level1
|
||||
WITH s, level1, edges1, r2, n2
|
||||
ORDER BY coalesce(r2.paper_count, 0) DESC, coalesce(r2.support_count, 0) DESC
|
||||
WITH s, level1, edges1,
|
||||
collect(DISTINCT n2)[..$max_nodes] AS level2,
|
||||
[r IN collect(DISTINCT r2)[..$max_edges] | r {.*,
|
||||
source_canonical_id: startNode(r).canonical_id,
|
||||
target_canonical_id: endNode(r).canonical_id}] AS edges2
|
||||
RETURN s {.*} AS start,
|
||||
[n IN (level1 + level2) | n {.*}] AS nodes,
|
||||
(edges1 + edges2) AS edges
|
||||
"""
|
||||
rows = await self._query(
|
||||
cypher, params, timeout=NEIGHBORHOOD_QUERY_TIMEOUT
|
||||
)
|
||||
if not rows:
|
||||
return {"center": None, "center_name": None, "depth": depth, "nodes": [], "edges": []}
|
||||
row = rows[0]
|
||||
start = row.get("start") or {}
|
||||
|
||||
node_list: list[dict[str, Any]] = [
|
||||
{
|
||||
"id": start.get("canonical_id"),
|
||||
"label": start.get("name"),
|
||||
"type": start.get("entity_type"),
|
||||
}
|
||||
]
|
||||
node_index: dict[str, int] = {}
|
||||
if start.get("canonical_id") is not None:
|
||||
node_index[start["canonical_id"]] = 0
|
||||
for raw in row.get("nodes") or []:
|
||||
nid = raw.get("canonical_id")
|
||||
if nid is None or nid in node_index:
|
||||
continue
|
||||
node_index[nid] = len(node_list)
|
||||
node_list.append(
|
||||
{
|
||||
"id": nid,
|
||||
"label": raw.get("name"),
|
||||
"type": raw.get("entity_type"),
|
||||
}
|
||||
)
|
||||
|
||||
edges: list[dict[str, Any]] = []
|
||||
for raw in row.get("edges") or []:
|
||||
src = raw.get("source_canonical_id")
|
||||
tgt = raw.get("target_canonical_id")
|
||||
if src is None or tgt is None:
|
||||
continue
|
||||
if src not in node_index or tgt not in node_index:
|
||||
continue
|
||||
# 自环边（实体自身关系，如 gene_structure）对可视化无意义，过滤
|
||||
if src == tgt:
|
||||
continue
|
||||
edges.append(
|
||||
{
|
||||
"source": src,
|
||||
"target": tgt,
|
||||
"relation": raw.get("relation_category")
|
||||
if raw.get("relation_category") is not None
|
||||
else raw.get("relation_norm"),
|
||||
"label": raw.get("relation_category")
|
||||
if raw.get("relation_category") is not None
|
||||
else raw.get("relation_norm"),
|
||||
"assertion_id": raw.get("assertion_id"),
|
||||
"paper_count": raw.get("paper_count"),
|
||||
"support_count": raw.get("support_count"),
|
||||
"title": raw.get("title"),
|
||||
"doi": raw.get("doi"),
|
||||
"pmid": raw.get("pmid"),
|
||||
"chain_id": None,
|
||||
"hop": None,
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"center": start.get("canonical_id"),
|
||||
"center_name": start.get("name"),
|
||||
"depth": depth,
|
||||
"nodes": node_list,
|
||||
"edges": edges,
|
||||
}
|
||||
|
||||
# ---------- 调试/冒烟 ----------
|
||||
|
||||
async def raw_query(self, cypher: str, **params: Any) -> list[dict[str, Any]]:
|
||||
"""只读原始查询（冒烟测试用）。"""
|
||||
return await self._query(cypher, params, timeout=QUERY_TIMEOUT)
|
||||
|
||||
|
||||
# 模块级单例：health 路由与 kg 路由共用同一连接池（app 关闭时 close）。
|
||||
service = Neo4jService()
|
||||
|
||||
|
||||
async def close_neo4j() -> None:
|
||||
"""app shutdown 钩子。"""
|
||||
await service.close()
|
||||
|
||||
|
||||
async def probe_neo4j() -> dict[str, Any]:
|
||||
"""/api/health 用的探测（不抛错）。"""
|
||||
return await service.probe()
|
||||
@@ -0,0 +1,268 @@
|
||||
"""normalize.py —— 上游响应 → 契约形态的归一化器（API.md §7/§8/§9）。
|
||||
|
||||
真实上游键以 DISCOVERY.md §4.1 为准；别名规则按 API.md 逐条实现。
|
||||
本模块是纯函数集合（不发起 IO），供 qa 路由（A2）与测试共用。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def pick(obj: dict, *names: str, default=None):
|
||||
"""取第一个存在且非 None 的键（归一化核心工具）。保留 PM 脚手架语义。"""
|
||||
for n in names:
|
||||
if isinstance(obj, dict) and n in obj and obj[n] is not None:
|
||||
return obj[n]
|
||||
return default
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# §8 Source 归一化（缺失字段键必须存在，值为 null）
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def normalize_source(item: dict[str, Any] | None) -> dict[str, Any]:
|
||||
"""单个 source item → API.md §8 形态（所有键存在，缺失为 null）。
|
||||
|
||||
别名映射（按 API.md §8 精确实现）：
|
||||
- content|evidence|text|snippet → snippet
|
||||
- doc_id|paper_id|pubmed_id|id|chunk_id → paper_id（chunk_id 单独保留）
|
||||
- final_score|score|confidence|rerank_score → score（真实键 final_score）
|
||||
- title|paper_title|source → title（真实 item 的 source 是断言主体实体名；
|
||||
title 存在时优先取 title，两者并存时 source 落入 metadata.source_entity）
|
||||
- url|doi_url → url；只有 doi 时派生 url = https://doi.org/{doi}
|
||||
- evidence_id|id → id；relation_id|assertion_id → metadata.assertion_id
|
||||
- output_layer / claim_scope / cluster_id / support_count / section /
|
||||
species / vector_score / rrf_score / rerank_score / retrieval_methods /
|
||||
retrieval_queries / cluster_member_* 原样透传（或进 metadata）
|
||||
"""
|
||||
if not isinstance(item, dict):
|
||||
item = {}
|
||||
|
||||
source = pick(item, "source") # 断言主体实体名（上游语义）
|
||||
title = pick(item, "title", "paper_title")
|
||||
|
||||
doi = pick(item, "doi")
|
||||
url = pick(item, "url", "doi_url")
|
||||
if url is None and doi:
|
||||
url = f"https://doi.org/{doi}"
|
||||
|
||||
metadata: dict[str, Any] = {}
|
||||
if isinstance(pick(item, "metadata"), dict):
|
||||
metadata.update(item["metadata"])
|
||||
if source is not None and title is not None and source != title:
|
||||
metadata.setdefault("source_entity", source)
|
||||
for key in (
|
||||
"vector_score", "rrf_score", "rerank_score", "kb",
|
||||
"retrieval_methods", "retrieval_queries",
|
||||
"cluster_member_count", "cluster_member_assertion_ids",
|
||||
"cluster_member_paper_ids", "cluster_member_pmids",
|
||||
"output_layer", "claim_scope", "species",
|
||||
):
|
||||
if key in item and item[key] is not None:
|
||||
metadata.setdefault(key, item[key])
|
||||
|
||||
assertion_id = pick(item, "relation_id", "assertion_id")
|
||||
if assertion_id is not None:
|
||||
metadata.setdefault("assertion_id", assertion_id)
|
||||
|
||||
return {
|
||||
"id": pick(item, "evidence_id", "id"),
|
||||
"title": title,
|
||||
"paper_id": pick(item, "doc_id", "paper_id", "pubmed_id", "id", "chunk_id"),
|
||||
"chunk_id": pick(item, "chunk_id"),
|
||||
"pmid": pick(item, "pmid"),
|
||||
"pmcid": pick(item, "pmcid"),
|
||||
"doi": doi,
|
||||
"url": url,
|
||||
"snippet": pick(item, "content", "evidence", "text", "snippet"),
|
||||
"score": pick(item, "final_score", "score", "confidence", "rerank_score"),
|
||||
"locator": pick(item, "locator"),
|
||||
"section": pick(item, "section"),
|
||||
"support_count": pick(item, "support_count"),
|
||||
"type": pick(item, "type"),
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# §9 图谱：hop / chain / {nodes,edges,chains}
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def normalize_hop(hop: dict[str, Any]) -> dict[str, Any]:
|
||||
"""hop 归一化：
|
||||
from_name|source → source、to_name|target → target、
|
||||
from_id|source_canonical_id → source 节点 id、to_id|target_canonical_id → target 节点 id、
|
||||
relation 原样、source_type|target_type → source/target 节点 type、
|
||||
paper_id / pmid / pmcid / doi 透传。
|
||||
"""
|
||||
if not isinstance(hop, dict):
|
||||
hop = {}
|
||||
return {
|
||||
"source": pick(hop, "from_name", "source"),
|
||||
"target": pick(hop, "to_name", "target"),
|
||||
"source_id": pick(hop, "from_id", "source_canonical_id"),
|
||||
"target_id": pick(hop, "to_id", "target_canonical_id"),
|
||||
"source_type": pick(hop, "source_type"),
|
||||
"target_type": pick(hop, "target_type"),
|
||||
"relation": pick(hop, "relation"),
|
||||
"assertion_id": pick(hop, "assertion_id"),
|
||||
"paper_id": pick(hop, "paper_id"),
|
||||
"pmid": pick(hop, "pmid"),
|
||||
"pmcid": pick(hop, "pmcid"),
|
||||
"doi": pick(hop, "doi"),
|
||||
}
|
||||
|
||||
|
||||
def normalize_chain_item(item: dict[str, Any]) -> dict[str, Any]:
|
||||
"""单个 chain item（真实形态：chain_id / hops[] 等）→ {chain_id, hops}。"""
|
||||
if not isinstance(item, dict):
|
||||
item = {}
|
||||
hops = [normalize_hop(h) for h in item.get("hops", []) if isinstance(h, dict)]
|
||||
return {"chain_id": pick(item, "chain_id"), "hops": hops}
|
||||
|
||||
|
||||
def extract_chains(raw: dict[str, Any] | None) -> list[dict[str, Any]]:
|
||||
"""从可能的两类源提取 chain 列表：
|
||||
- `native_graph_chains`：列表（真实形态，305 条）
|
||||
- `graph_knowledge_chains`：dict 形态 `{mode, chains[], chain_count, …}`，取 `.chains`
|
||||
两种形态都要能提取出 hop 列表。
|
||||
"""
|
||||
if not isinstance(raw, dict):
|
||||
return []
|
||||
chains = pick(raw, "native_graph_chains", "graph_knowledge_chains")
|
||||
if isinstance(chains, dict):
|
||||
chains = chains.get("chains")
|
||||
if not isinstance(chains, list):
|
||||
return []
|
||||
return [normalize_chain_item(c) for c in chains]
|
||||
|
||||
|
||||
def normalize_graph(raw: dict[str, Any] | None) -> dict[str, Any]:
|
||||
"""从 `knowledge_graph` / `graph_edges` / `native_graph_chains`（列表形态）/
|
||||
`graph_knowledge_chains`（dict 形态）提取并归一化为 {nodes, edges, chains}。
|
||||
|
||||
- 节点去重（按 id；id 缺失的 hop 跳过）；label 取 from_name/to_name、type 取
|
||||
source_type/target_type。
|
||||
- 边携带 chain_id 与 hop 序号（hop_index 缺失时按列表序补）。
|
||||
- 链保持 chain_id + hops 列表。
|
||||
"""
|
||||
raw = raw if isinstance(raw, dict) else {}
|
||||
chains = extract_chains(raw)
|
||||
|
||||
nodes: dict[str, dict[str, Any]] = {}
|
||||
edges: list[dict[str, Any]] = []
|
||||
|
||||
def add_node(node_id: Any, name: Any, typ: Any) -> None:
|
||||
if not node_id:
|
||||
return
|
||||
key = str(node_id)
|
||||
if key not in nodes:
|
||||
nodes[key] = {"id": node_id, "label": name, "type": typ}
|
||||
|
||||
for chain in chains:
|
||||
for idx, hop in enumerate(chain["hops"]):
|
||||
hop_index = hop.get("hop_index")
|
||||
if hop_index is None:
|
||||
hop_index = idx
|
||||
# 节点 id 优先 canonical id；label 优先名字，退化到 id
|
||||
src_id = pick(hop, "source_id", "source") or hop.get("source")
|
||||
tgt_id = pick(hop, "target_id", "target") or hop.get("target")
|
||||
src_label = hop.get("source") or src_id
|
||||
tgt_label = hop.get("target") or tgt_id
|
||||
add_node(src_id, src_label, hop.get("source_type"))
|
||||
add_node(tgt_id, tgt_label, hop.get("target_type"))
|
||||
if src_id and tgt_id:
|
||||
edges.append({
|
||||
"source": src_id,
|
||||
"target": tgt_id,
|
||||
"relation": hop.get("relation"),
|
||||
"label": hop.get("relation"),
|
||||
"chain_id": chain.get("chain_id"),
|
||||
"hop": hop_index,
|
||||
"paper_id": hop.get("paper_id"),
|
||||
"pmid": hop.get("pmid"),
|
||||
"pmcid": hop.get("pmcid"),
|
||||
"doi": hop.get("doi"),
|
||||
})
|
||||
|
||||
# 兼容显式 {nodes, edges} 形态（knowledge_graph / graph_edges 直接给边列表时）
|
||||
explicit_nodes = raw.get("nodes")
|
||||
explicit_edges = raw.get("edges")
|
||||
if isinstance(explicit_nodes, list):
|
||||
for n in explicit_nodes:
|
||||
if isinstance(n, dict) and n.get("id"):
|
||||
key = str(n["id"])
|
||||
if key not in nodes:
|
||||
nodes[key] = {
|
||||
"id": n["id"],
|
||||
"label": pick(n, "label"),
|
||||
"type": pick(n, "type"),
|
||||
}
|
||||
if isinstance(explicit_edges, list):
|
||||
edges.extend(
|
||||
{
|
||||
"source": pick(e, "source"),
|
||||
"target": pick(e, "target"),
|
||||
"relation": pick(e, "relation"),
|
||||
"label": pick(e, "label", "relation"),
|
||||
"chain_id": pick(e, "chain_id"),
|
||||
"hop": pick(e, "hop"),
|
||||
"paper_id": pick(e, "paper_id"),
|
||||
"pmid": pick(e, "pmid"),
|
||||
"pmcid": pick(e, "pmcid"),
|
||||
"doi": pick(e, "doi"),
|
||||
}
|
||||
for e in explicit_edges
|
||||
if isinstance(e, dict)
|
||||
)
|
||||
|
||||
return {
|
||||
"nodes": list(nodes.values()),
|
||||
"edges": edges,
|
||||
"chains": chains,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# §7 归一化 QA 响应（/chat、/chat/stream complete、/search 共用）
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def normalize_qa_response(raw: dict[str, Any] | None) -> dict[str, Any]:
|
||||
"""上游 31 键响应 → API.md §7 形态。
|
||||
|
||||
别名（真实键优先）：
|
||||
- answer|response|message → answer
|
||||
- sources|citations → sources；used_sources|references → used_sources
|
||||
- query_plan|agentic_plan|retrieval_plan → retrieval_plan（真实键 query_plan）
|
||||
- retrieval_notes|notes → retrieval_notes
|
||||
- timings|timing → timing（真实键 timings）
|
||||
- citation_validation|citation_map → citation_map（真实键 citation_validation）
|
||||
- 图谱四源 → graph + chains（见 normalize_graph）
|
||||
- success 透传；raw 保存原始完整响应
|
||||
"""
|
||||
raw = raw if isinstance(raw, dict) else {}
|
||||
|
||||
sources = [normalize_source(s) for s in (pick(raw, "sources", "citations") or [])]
|
||||
used_sources = [
|
||||
normalize_source(s) for s in (pick(raw, "used_sources", "references") or [])
|
||||
]
|
||||
graph = normalize_graph(raw)
|
||||
retrieval_plan = pick(raw, "query_plan", "agentic_plan", "retrieval_plan")
|
||||
|
||||
return {
|
||||
"query_id": pick(raw, "query_id"),
|
||||
"version": pick(raw, "version"),
|
||||
"mode": pick(raw, "mode", default="chat"),
|
||||
"answer": pick(raw, "answer", "response", "message", default=""),
|
||||
"source_count": pick(raw, "source_count", default=len(sources)),
|
||||
"sources": sources,
|
||||
"used_sources": used_sources,
|
||||
"chains": graph["chains"],
|
||||
"graph": {"nodes": graph["nodes"], "edges": graph["edges"], "chains": graph["chains"]},
|
||||
"retrieval_plan": retrieval_plan,
|
||||
"retrieval_notes": pick(raw, "retrieval_notes", "notes"),
|
||||
"timing": pick(raw, "timings", "timing"),
|
||||
"citation_map": pick(raw, "citation_validation", "citation_map"),
|
||||
"success": pick(raw, "success", default=True),
|
||||
"raw": raw,
|
||||
}
|
||||
@@ -0,0 +1,590 @@
|
||||
"""services/sse_relay.py —— 上游 QA SSE → 前端契约帧中继（A2 归属）。
|
||||
|
||||
上游：127.0.0.1:7920 POST /api/v1/qa-v1-1-73/chat/stream
|
||||
（真实帧序列见 DISCOVERY.md §4.2：`event: <name>` + `data: <JSON>` 两行一组，
|
||||
data 可能多行，期间有 `: keepalive` 注释帧，终端为 `event: complete` + 完整 JSON。）
|
||||
|
||||
下游（前端契约帧，API.md §4）：
|
||||
event: progress data: {stage, phase, label, details, at} ← phase 由 BFF 计算
|
||||
event: complete data: {<同 /chat 的归一化响应>}
|
||||
event: raw data: {<无法分类的上游 JSON>} ← 绝不丢数据
|
||||
event: error data: {code, message, retryable}
|
||||
event: timeout data: {message}
|
||||
: ping （15s 保活注释帧）
|
||||
|
||||
规则：
|
||||
- 上游 HTTP 头阶段失败且 detail.error_code=="invalid_llm_citation" → 自动重试一次，
|
||||
仍失败 → `event: error`（INVALID_CITATION）。
|
||||
- 看门狗：15s 无输出发 `: ping`；240s 无任何上游事件 → `event: timeout` 并关闭。
|
||||
- 客户端断开（StreamingResponse 取消 → CancelledError 传播）→ `async with` 退出，
|
||||
上游流随之关闭。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from collections import deque
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, AsyncIterator, Callable, Optional
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 常量（TODO(A2→A1 对齐点)：A1 的 app.config.settings 就绪后，
|
||||
# 上游地址与「超时阶梯」改由 settings 提供，此处仅作兜底默认值）
|
||||
# ---------------------------------------------------------------------------
|
||||
QA_BASE_URL = "http://127.0.0.1:7920"
|
||||
QA_UPSTREAM_PREFIX = "/api/v1/qa-v1-1-73"
|
||||
CONNECT_TIMEOUT = 10.0
|
||||
STREAM_IDLE_TIMEOUT = 240.0
|
||||
STREAM_KEEPALIVE_INTERVAL = 15.0
|
||||
INVALID_CITATION_ERROR_CODE = "invalid_llm_citation"
|
||||
|
||||
# API.md §4：SSE 响应头
|
||||
sse_headers = {
|
||||
"Cache-Control": "no-cache",
|
||||
"X-Accel-Buffering": "no",
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 阶段 → 相位映射（API.md §4 表；与 i18n/labels.js 同源，PM 归属）。
|
||||
# 未知 stage → phase "retrieval"，label = stage 原文。
|
||||
# ---------------------------------------------------------------------------
|
||||
STAGE_PHASE: dict[str, str] = {
|
||||
"understanding_started": "planning",
|
||||
"entities_identified": "planning",
|
||||
"related_search_started": "planning",
|
||||
"concepts_found": "retrieval",
|
||||
"subquestion_retrieved": "retrieval",
|
||||
"agent_review_started": "retrieval",
|
||||
"subquestion_retry_started": "retrieval",
|
||||
"evidence_accepted": "evidence",
|
||||
"evidence_review_completed": "evidence",
|
||||
"answer_generation_started": "generation",
|
||||
"completed": "done",
|
||||
}
|
||||
|
||||
STAGE_LABEL: dict[str, str] = {
|
||||
"understanding_started": "开始理解问题",
|
||||
"entities_identified": "已识别实体",
|
||||
"related_search_started": "开始关联检索",
|
||||
"concepts_found": "发现相关概念",
|
||||
"subquestion_retrieved": "子问题召回完成",
|
||||
"agent_review_started": "开始智能体审核",
|
||||
"subquestion_retry_started": "子问题重试中",
|
||||
"evidence_accepted": "证据已采纳",
|
||||
"evidence_review_completed": "证据审核完成",
|
||||
"answer_generation_started": "开始生成答案",
|
||||
"completed": "检索完成",
|
||||
}
|
||||
|
||||
UNKNOWN_STAGE_PHASE = "retrieval"
|
||||
|
||||
|
||||
def phase_for(stage: str) -> tuple[str, str]:
|
||||
"""stage → (phase, label)；未知 stage 回落 retrieval / stage 原文。"""
|
||||
return (
|
||||
STAGE_PHASE.get(stage, UNKNOWN_STAGE_PHASE),
|
||||
STAGE_LABEL.get(stage, stage),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 错误文案与载荷（TODO(A2→A1 对齐点)：HTTP JSON 信封换用 app.upstream.envelope()）
|
||||
# ---------------------------------------------------------------------------
|
||||
ERROR_MESSAGES: dict[str, tuple[str, bool]] = {
|
||||
"UPSTREAM_DOWN": ("知识问答服务当前不可用，请稍后重试", True),
|
||||
"UPSTREAM_ERROR": ("知识问答服务返回错误", False),
|
||||
"TIMEOUT": ("检索超时", False),
|
||||
"INVALID_CITATION": ("引用校验失败，请重试", False),
|
||||
"VALIDATION": ("请求参数不合法", False),
|
||||
"INTERNAL": ("内部错误", False),
|
||||
}
|
||||
|
||||
RETRYABLE_ERRORS = {"TIMEOUT", "INVALID_CITATION", "UPSTREAM_ERROR"}
|
||||
|
||||
|
||||
def error_payload(code: str, message: Optional[str] = None) -> dict:
|
||||
"""SSE `event: error` 的 data 载荷（API.md §4：{code, message, retryable}）。"""
|
||||
default_message, _degraded = ERROR_MESSAGES.get(code, ("内部错误", False))
|
||||
return {
|
||||
"code": code,
|
||||
"message": message or default_message,
|
||||
"retryable": code in RETRYABLE_ERRORS,
|
||||
}
|
||||
|
||||
|
||||
def is_invalid_citation(body: Any) -> bool:
|
||||
"""上游错误体 `{"detail": {"error_code": "invalid_llm_citation"}}` 判定。"""
|
||||
if not isinstance(body, dict):
|
||||
return False
|
||||
detail = body.get("detail")
|
||||
if isinstance(detail, dict):
|
||||
return detail.get("error_code") == INVALID_CITATION_ERROR_CODE
|
||||
return body.get("error_code") == INVALID_CITATION_ERROR_CODE
|
||||
|
||||
|
||||
def status_to_error(status: int, detail: Any) -> tuple[str, Optional[str]]:
|
||||
"""HTTP≥400 头阶段失败 → (code, message)。"""
|
||||
if is_invalid_citation(detail):
|
||||
return "INVALID_CITATION", None
|
||||
if isinstance(detail, dict):
|
||||
d = detail.get("detail")
|
||||
if isinstance(d, dict):
|
||||
return "UPSTREAM_ERROR", d.get("message") or d.get("error")
|
||||
return "UPSTREAM_ERROR", detail.get("message") or detail.get("error")
|
||||
return "UPSTREAM_ERROR", None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 上游 client 与地址（TODO(A2→A1 对齐点)：A1 的 app.upstream client 单例 /
|
||||
# settings 上游地址就绪后，以下两函数直接改为其代理）
|
||||
# ---------------------------------------------------------------------------
|
||||
_qa_client: Optional[httpx.AsyncClient] = None
|
||||
|
||||
|
||||
def get_qa_client() -> httpx.AsyncClient:
|
||||
"""QA 上游 httpx client。
|
||||
|
||||
注意（2026-08-25 PM 集成修复）：**不复用 A1 upstream 单例**。
|
||||
A1 的 get_client() 是 async 函数且默认 read=300s，而 QA 流式路径
|
||||
需要 read=None + 应用层看门狗（长静默 SSE 会被 300s 读超时误杀）。
|
||||
两个 client 语义不同，各自独立持有。
|
||||
"""
|
||||
global _qa_client
|
||||
if _qa_client is None or _qa_client.is_closed:
|
||||
_qa_client = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(connect=CONNECT_TIMEOUT, read=None, write=30.0, pool=CONNECT_TIMEOUT),
|
||||
)
|
||||
return _qa_client
|
||||
|
||||
|
||||
def qa_base_url() -> str:
|
||||
"""QA 上游 base URL。A1 settings 就绪后从配置读取。"""
|
||||
try:
|
||||
from ..config import settings # A1 归属
|
||||
|
||||
for attr in ("qa_base_url", "qa_url", "upstream_qa_url", "qa_upstream_base"):
|
||||
value = getattr(settings, attr, None)
|
||||
if value:
|
||||
return str(value).rstrip("/")
|
||||
except Exception:
|
||||
pass
|
||||
return QA_BASE_URL
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 归一化（TODO(A2→A1 对齐点)：A1 的 app.services.normalize.normalize_qa_response
|
||||
# 就绪后，get_qa_normalizer 直接返回之，以下 fallback 全部删除）
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def get_qa_normalizer() -> Callable[[dict], dict]:
|
||||
"""QA 响应归一化函数：优先 A1 实现，未就绪时用等价本地逻辑。"""
|
||||
try:
|
||||
from .normalize import normalize_qa_response # A1 归属
|
||||
|
||||
return normalize_qa_response
|
||||
except (ImportError, AttributeError):
|
||||
return _fallback_normalize_qa
|
||||
|
||||
|
||||
def _fallback_normalize_qa(payload: dict) -> dict:
|
||||
"""等价本地归一化（API.md §7 顶层别名 / §8 source / §9 图谱链条双形态容错）。"""
|
||||
from .normalize import pick # A1 已提供
|
||||
|
||||
src = payload if isinstance(payload, dict) else {}
|
||||
sources = [_fallback_source(i) for i in (pick(src, "sources", "citations", default=[]) or [])]
|
||||
used = [_fallback_source(i) for i in (pick(src, "used_sources", "references", default=[]) or [])]
|
||||
chains, graph = _fallback_graph_chains(src)
|
||||
return {
|
||||
"query_id": pick(src, "query_id"),
|
||||
"version": pick(src, "version"),
|
||||
"mode": "chat", # 调用方（qa.py）按端点覆盖 chat|search
|
||||
"answer": pick(src, "answer", "response", "message", default=""),
|
||||
"source_count": pick(src, "source_count", default=len(sources)),
|
||||
"sources": sources,
|
||||
"used_sources": used,
|
||||
"chains": chains,
|
||||
"graph": graph,
|
||||
"retrieval_plan": pick(src, "retrieval_plan", "agentic_plan", "query_plan"),
|
||||
"retrieval_notes": pick(src, "retrieval_notes", "notes"),
|
||||
"timing": pick(src, "timing", "timings"),
|
||||
"citation_map": pick(src, "citation_map", "citation_validation"),
|
||||
"raw": src,
|
||||
}
|
||||
|
||||
|
||||
_SOURCE_PASSTHROUGH_KEYS = (
|
||||
"output_layer", "claim_scope", "cluster_id", "support_count", "section", "species",
|
||||
)
|
||||
_SOURCE_METADATA_KEYS = (
|
||||
"vector_score", "rrf_score", "rerank_score", "relation_id", "assertion_id",
|
||||
"cluster_member_count", "cluster_member_assertion_ids", "cluster_member_paper_ids",
|
||||
"cluster_member_pmids", "retrieval_methods", "retrieval_queries", "kb",
|
||||
)
|
||||
|
||||
|
||||
def _fallback_source(item: Any) -> dict:
|
||||
"""§8 source 别名归一化；缺失字段键必须存在（值为 null）。"""
|
||||
from .normalize import pick # A1 已提供
|
||||
|
||||
if not isinstance(item, dict):
|
||||
item = {}
|
||||
doi = pick(item, "doi")
|
||||
url = pick(item, "url", "doi_url")
|
||||
if url is None and doi:
|
||||
url = f"https://doi.org/{doi}"
|
||||
title = pick(item, "title", "paper_title")
|
||||
metadata: dict[str, Any] = {}
|
||||
for k in _SOURCE_PASSTHROUGH_KEYS:
|
||||
if k in item:
|
||||
metadata[k] = item[k]
|
||||
for k in _SOURCE_METADATA_KEYS:
|
||||
if k in item:
|
||||
metadata[k] = item[k]
|
||||
if "source" in item:
|
||||
# 真实 item 的 source 是断言主体实体名而非论文标题
|
||||
if title is None:
|
||||
title = item["source"]
|
||||
else:
|
||||
metadata.setdefault("source_entity", item["source"])
|
||||
return {
|
||||
"id": pick(item, "evidence_id", "id"),
|
||||
"title": title,
|
||||
"paper_id": pick(item, "paper_id", "doc_id", "pubmed_id", "id", "chunk_id"),
|
||||
"chunk_id": pick(item, "chunk_id"),
|
||||
"pmid": pick(item, "pmid"),
|
||||
"pmcid": pick(item, "pmcid"),
|
||||
"doi": doi,
|
||||
"url": url,
|
||||
"snippet": pick(item, "content", "evidence", "text", "snippet"),
|
||||
"score": pick(item, "final_score", "score", "confidence", "rerank_score"),
|
||||
"locator": pick(item, "locator"),
|
||||
"section": pick(item, "section"),
|
||||
"support_count": pick(item, "support_count"),
|
||||
"type": pick(item, "type"),
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
|
||||
def _fallback_graph_chains(src: dict) -> tuple[list[dict], dict]:
|
||||
"""§9 图谱与链条：native_graph_chains(list) / graph_knowledge_chains(dict) 双形态容错。
|
||||
|
||||
返回 (chains, graph)；graph = {nodes, edges, chains}。
|
||||
"""
|
||||
from .normalize import pick # A1 已提供
|
||||
|
||||
raw_chains: list = pick(src, "native_graph_chains", default=[]) or []
|
||||
if isinstance(raw_chains, dict): # 容错：意外为 dict
|
||||
raw_chains = raw_chains.get("chains") or []
|
||||
gkc = pick(src, "graph_knowledge_chains")
|
||||
if isinstance(gkc, dict):
|
||||
raw_chains = raw_chains + (gkc.get("chains") or [])
|
||||
elif isinstance(gkc, list):
|
||||
raw_chains = raw_chains + gkc
|
||||
|
||||
chains: list[dict] = []
|
||||
nodes: dict[str, dict] = {}
|
||||
edges: list[dict] = []
|
||||
seen_chain: set = set()
|
||||
|
||||
for idx, chain in enumerate(raw_chains):
|
||||
if not isinstance(chain, dict):
|
||||
continue
|
||||
chain_id = pick(chain, "chain_id", default=f"chain_{idx}")
|
||||
if chain_id in seen_chain: # native 与 graph_knowledge 两形态可能镜像同数据
|
||||
continue
|
||||
seen_chain.add(chain_id)
|
||||
hops_out: list[dict] = []
|
||||
for hop in (pick(chain, "hops", default=[]) or []):
|
||||
if not isinstance(hop, dict):
|
||||
continue
|
||||
source_name = pick(hop, "from_name", "source")
|
||||
target_name = pick(hop, "to_name", "target")
|
||||
src_id = pick(hop, "from_id", "source_canonical_id") or source_name
|
||||
tgt_id = pick(hop, "to_id", "target_canonical_id") or target_name
|
||||
hops_out.append({
|
||||
"source": source_name,
|
||||
"target": target_name,
|
||||
"relation": pick(hop, "relation"),
|
||||
"source_type": pick(hop, "source_type"),
|
||||
"target_type": pick(hop, "target_type"),
|
||||
"paper_id": pick(hop, "paper_id"),
|
||||
"pmid": pick(hop, "pmid"),
|
||||
"pmcid": pick(hop, "pmcid"),
|
||||
"doi": pick(hop, "doi"),
|
||||
})
|
||||
nodes.setdefault(src_id, {"id": src_id, "label": source_name, "type": pick(hop, "source_type")})
|
||||
nodes.setdefault(tgt_id, {"id": tgt_id, "label": target_name, "type": pick(hop, "target_type")})
|
||||
edges.append({
|
||||
"source": src_id,
|
||||
"target": tgt_id,
|
||||
"relation": pick(hop, "relation"),
|
||||
"label": pick(hop, "relation"),
|
||||
"chain_id": chain_id,
|
||||
"hop": pick(hop, "hop_index"),
|
||||
"paper_id": pick(hop, "paper_id"),
|
||||
"pmid": pick(hop, "pmid"),
|
||||
"pmcid": pick(hop, "pmcid"),
|
||||
"doi": pick(hop, "doi"),
|
||||
})
|
||||
chains.append({"chain_id": chain_id, "hops": hops_out})
|
||||
|
||||
# knowledge_graph 直出形态（上游若给出现成 nodes/edges，直接用）
|
||||
kg = pick(src, "knowledge_graph")
|
||||
if isinstance(kg, dict):
|
||||
for n in (kg.get("nodes") or []):
|
||||
if isinstance(n, dict) and n.get("id") not in nodes:
|
||||
nodes[n["id"]] = n
|
||||
edges = (kg.get("edges") or []) + edges
|
||||
elif not edges:
|
||||
edges = pick(src, "graph_edges", default=[]) or []
|
||||
|
||||
graph = {"nodes": list(nodes.values()), "edges": edges, "chains": chains}
|
||||
return chains, graph
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SSE 帧工具
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def utc_now_iso() -> str:
|
||||
return datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z")
|
||||
|
||||
|
||||
def sse_frame(event: str, data: Any = None) -> str:
|
||||
text = data if isinstance(data, str) else json.dumps(data, ensure_ascii=False)
|
||||
return f"event: {event}\ndata: {text}\n\n"
|
||||
|
||||
|
||||
def _progress_payload(payload: dict, at: str) -> dict:
|
||||
stage = payload.get("stage")
|
||||
phase, label = phase_for(stage)
|
||||
details = payload.get("details")
|
||||
if not isinstance(details, dict):
|
||||
details = {}
|
||||
return {"stage": stage, "phase": phase, "label": label, "details": details, "at": at}
|
||||
|
||||
|
||||
def _classify(event_name: Optional[str], data_text: str) -> Optional[tuple[str, Any]]:
|
||||
"""一条上游事件 → (kind, payload)；kind ∈ progress|complete|error|raw。
|
||||
|
||||
None 表示无内容可发（空帧）。无 event 名时按 payload 形态分类。
|
||||
"""
|
||||
if not data_text.strip():
|
||||
if not event_name:
|
||||
return None
|
||||
return ("raw", {"payload": data_text, "event_name": event_name})
|
||||
try:
|
||||
payload = json.loads(data_text)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
# 不是 JSON —— 透传不丢数据
|
||||
return ("raw", {"payload": data_text, "event_name": event_name or None})
|
||||
if not isinstance(payload, dict):
|
||||
return ("raw", payload)
|
||||
if "stage" in payload:
|
||||
return ("progress", payload)
|
||||
if event_name == "error":
|
||||
return ("error", _error_payload_from_upstream(payload))
|
||||
if event_name == "complete" or "answer" in payload or "success" in payload:
|
||||
return ("complete", payload)
|
||||
return ("raw", payload)
|
||||
|
||||
|
||||
def _error_payload_from_upstream(payload: dict) -> dict:
|
||||
if is_invalid_citation(payload):
|
||||
return error_payload("INVALID_CITATION")
|
||||
detail = payload.get("detail")
|
||||
if isinstance(detail, dict):
|
||||
return error_payload("UPSTREAM_ERROR", detail.get("message") or detail.get("error"))
|
||||
return error_payload("UPSTREAM_ERROR", payload.get("message") or payload.get("error"))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 中继核心
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def relay_stream(
|
||||
body: dict[str, Any],
|
||||
*,
|
||||
client: Optional[httpx.AsyncClient] = None,
|
||||
base_url: Optional[str] = None,
|
||||
normalizer: Optional[Callable[[dict], dict]] = None,
|
||||
idle_timeout: Optional[float] = None,
|
||||
keepalive_interval: Optional[float] = None,
|
||||
now_fn: Callable[[], str] = utc_now_iso,
|
||||
) -> AsyncIterator[str]:
|
||||
"""POST 上游 /chat/stream 并把帧中继为前端契约帧（API.md §4）。
|
||||
|
||||
客户端断开 → StreamingResponse 取消本生成器 → CancelledError 传播 →
|
||||
`async with client.stream(...)` 退出，上游流随之关闭。
|
||||
"""
|
||||
client = client or get_qa_client()
|
||||
normalizer = normalizer or get_qa_normalizer()
|
||||
idle_timeout = STREAM_IDLE_TIMEOUT if idle_timeout is None else idle_timeout
|
||||
keepalive_interval = STREAM_KEEPALIVE_INTERVAL if keepalive_interval is None else keepalive_interval
|
||||
url = f"{base_url or qa_base_url()}{QA_UPSTREAM_PREFIX}/chat/stream"
|
||||
timeout = httpx.Timeout(connect=CONNECT_TIMEOUT, read=None, write=30.0, pool=CONNECT_TIMEOUT)
|
||||
|
||||
for attempt in (1, 2):
|
||||
try:
|
||||
async with client.stream(
|
||||
"POST", url, json=body, timeout=timeout, headers={"Accept": "text/event-stream"}
|
||||
) as resp:
|
||||
if resp.status_code >= 400:
|
||||
detail = await _safe_aread(resp)
|
||||
if attempt == 1 and is_invalid_citation(detail):
|
||||
logger.warning(
|
||||
"上游 /chat/stream 返回 %d invalid_llm_citation，自动重试（第 1 次）",
|
||||
resp.status_code,
|
||||
)
|
||||
continue
|
||||
code, message = status_to_error(resp.status_code, detail)
|
||||
yield sse_frame("error", error_payload(code, message))
|
||||
return
|
||||
async for frame_text in _pump_events(
|
||||
resp, normalizer, idle_timeout, keepalive_interval, now_fn
|
||||
):
|
||||
yield frame_text
|
||||
return
|
||||
except httpx.ReadTimeout:
|
||||
yield sse_frame("error", error_payload("TIMEOUT"))
|
||||
return
|
||||
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout) as exc:
|
||||
logger.warning("上游 /chat/stream 连接失败: %s", exc)
|
||||
yield sse_frame("error", error_payload("UPSTREAM_DOWN"))
|
||||
return
|
||||
except httpx.HTTPError as exc:
|
||||
logger.warning("上游 /chat/stream 请求失败: %s", exc)
|
||||
yield sse_frame("error", error_payload("UPSTREAM_ERROR"))
|
||||
return
|
||||
# 理论上不可达（attempt 循环内必然 return）；防御性兜底
|
||||
yield sse_frame("error", error_payload("UPSTREAM_ERROR"))
|
||||
|
||||
|
||||
async def _safe_aread(resp: httpx.Response) -> Any:
|
||||
try:
|
||||
raw = await resp.aread()
|
||||
if not raw:
|
||||
return None
|
||||
return json.loads(raw.decode("utf-8", errors="replace"))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def _pump_events(
|
||||
resp: httpx.Response,
|
||||
normalizer: Callable[[dict], dict],
|
||||
idle_timeout: float,
|
||||
keepalive_interval: float,
|
||||
now_fn: Callable[[], str],
|
||||
) -> AsyncIterator[str]:
|
||||
"""消费上游 SSE 字节流并产出契约帧。
|
||||
|
||||
设计要点：上游读取放在独立的 producer task（一次 `anext` 可能挂数分钟——
|
||||
上游逐字节慢速输出是常态）。看门狗只在 `Event.wait()` 上施加 `fail_after`，
|
||||
避免取消传播进 httpx 的字节生成器（asyncio 后端下那会破坏后续所有 await，
|
||||
使 timeout/ping 帧发不出去）。
|
||||
|
||||
看门狗：
|
||||
- 上游字节静默 `keepalive_interval` → `: ping`；
|
||||
- 自上次完整上游事件起 `idle_timeout` 无事件 → `event: timeout` 并关闭。
|
||||
"""
|
||||
chunks = resp.aiter_raw()
|
||||
queue: deque[tuple[str, bool]] = deque() # (frame_text, is_terminal)
|
||||
wakeup = anyio.Event() # 持久事件：consumer 等待前先 clear，避免残留 set 造成忙循环
|
||||
|
||||
def enqueue(frame_text: str, terminal: bool) -> None:
|
||||
queue.append((frame_text, terminal))
|
||||
wakeup.set()
|
||||
|
||||
async def producer() -> None:
|
||||
buf = b""
|
||||
event_name: Optional[str] = None
|
||||
data_lines: list[str] = []
|
||||
|
||||
def dispatch_local() -> list[tuple[str, bool]]:
|
||||
nonlocal event_name, data_lines
|
||||
kind = _classify(event_name, "\n".join(data_lines))
|
||||
event_name, data_lines = None, []
|
||||
if kind is None:
|
||||
return []
|
||||
kind_name, payload = kind
|
||||
if kind_name == "progress":
|
||||
return [(sse_frame("progress", _progress_payload(payload, now_fn())), False)]
|
||||
if kind_name == "complete":
|
||||
try:
|
||||
normalized = normalizer(payload)
|
||||
except Exception:
|
||||
logger.exception("归一化上游 complete 帧失败")
|
||||
return [(sse_frame("error", error_payload("INTERNAL")), True)]
|
||||
return [(sse_frame("complete", normalized), True)]
|
||||
if kind_name == "error":
|
||||
return [(sse_frame("error", payload), True)]
|
||||
return [(sse_frame("raw", payload), False)]
|
||||
|
||||
try:
|
||||
async for chunk in chunks:
|
||||
buf += chunk
|
||||
while b"\n" in buf:
|
||||
raw_line, buf = buf.split(b"\n", 1)
|
||||
line = raw_line.rstrip(b"\r").decode("utf-8", errors="replace")
|
||||
if line == "":
|
||||
for frame_text, terminal in dispatch_local():
|
||||
enqueue(frame_text, terminal)
|
||||
if queue and any(t for _, t in queue):
|
||||
return # 终态已入队，停止读取，连接由 aclose 关闭
|
||||
elif line.startswith(":"):
|
||||
continue # 上游注释帧（keepalive）忽略；保活由本中继自己的 : ping 承担
|
||||
elif line.startswith("event:"):
|
||||
event_name = line[len("event:"):].strip() or None
|
||||
elif line.startswith("data:"):
|
||||
data_lines.append(line[len("data:"):].strip())
|
||||
# 其他 SSE 字段（id:/retry:）忽略
|
||||
# 上游流在终态帧之前关闭：残余帧照发（绝不丢数据），再报错
|
||||
for frame_text, terminal in dispatch_local():
|
||||
enqueue(frame_text, terminal)
|
||||
enqueue(sse_frame("error", error_payload("UPSTREAM_ERROR", "知识问答服务连接意外关闭")), True)
|
||||
except httpx.ReadTimeout:
|
||||
enqueue(sse_frame("error", error_payload("TIMEOUT")), True)
|
||||
except (httpx.ConnectError, httpx.ConnectTimeout):
|
||||
enqueue(sse_frame("error", error_payload("UPSTREAM_DOWN")), True)
|
||||
except httpx.HTTPError as exc:
|
||||
logger.warning("上游 SSE 流中断: %s", exc)
|
||||
enqueue(sse_frame("error", error_payload("UPSTREAM_ERROR")), True)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(producer)
|
||||
last_upstream_event = time.monotonic()
|
||||
while True:
|
||||
if queue:
|
||||
last_upstream_event = time.monotonic()
|
||||
frame_text, terminal = queue.popleft()
|
||||
yield frame_text
|
||||
if terminal:
|
||||
tg.cancel_scope.cancel()
|
||||
return
|
||||
continue
|
||||
now = time.monotonic()
|
||||
if now - last_upstream_event >= idle_timeout:
|
||||
yield sse_frame("timeout", {"message": "检索超时"})
|
||||
tg.cancel_scope.cancel()
|
||||
return
|
||||
wakeup = anyio.Event()
|
||||
try:
|
||||
with anyio.fail_after(min(keepalive_interval, idle_timeout - (now - last_upstream_event))):
|
||||
await wakeup.wait()
|
||||
except TimeoutError:
|
||||
now = time.monotonic()
|
||||
if now - last_upstream_event >= idle_timeout:
|
||||
yield sse_frame("timeout", {"message": "检索超时"})
|
||||
tg.cancel_scope.cancel()
|
||||
return
|
||||
yield ": ping\n\n"
|
||||
@@ -0,0 +1,382 @@
|
||||
"""services/test_extraction.py —— 抽取代理单元测试（respx mock HTTP）。
|
||||
|
||||
覆盖：progress_pct 派生（四状态 + failed=null）、8000 down → 7900 回落命中、
|
||||
alt 失败缓存（30s 内不再试 alt）、jobs 列表归一化、job 详情归一化、
|
||||
plan 请求 409 → null 不报错、上传流式转发 + message 附加。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
from app.services.extraction import (
|
||||
ALT_FAILURE_TTL,
|
||||
ExtractionService,
|
||||
ExtractionUnavailable,
|
||||
derive_progress_pct,
|
||||
)
|
||||
|
||||
PRIMARY = "http://primary.test"
|
||||
ALT = "http://alt.test"
|
||||
PREFIX = "/api/v1/methodology"
|
||||
|
||||
|
||||
def make_service(primary: str = PRIMARY, alt: str = ALT) -> ExtractionService:
|
||||
return ExtractionService(primary_base_url=primary, alt_base_url=alt)
|
||||
|
||||
|
||||
# ---------- progress_pct 派生（API.md §10） ----------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"stage,expected",
|
||||
[
|
||||
("queued", 5),
|
||||
("mineru", 35),
|
||||
("llm", 75),
|
||||
("completed", 100),
|
||||
("failed", None),
|
||||
(None, None),
|
||||
("unknown_stage", None),
|
||||
],
|
||||
)
|
||||
def test_derive_progress_pct(stage, expected):
|
||||
assert derive_progress_pct(stage) == expected
|
||||
|
||||
|
||||
# ---------- 回落逻辑 ----------
|
||||
|
||||
async def test_primary_down_falls_back_to_alt(respx_mock):
|
||||
# 8000 down（连接失败），7900 命中
|
||||
respx_mock.get(f"{PRIMARY}{PREFIX}/extract-jobs?limit=20").mock(
|
||||
side_effect=httpx.ConnectError("primary down")
|
||||
)
|
||||
respx_mock.get(f"{ALT}{PREFIX}/extract-jobs?limit=20").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"jobs": [
|
||||
{
|
||||
"job_id": "2943168e9c48",
|
||||
"status": "completed",
|
||||
"stage": "completed",
|
||||
"substage": "done",
|
||||
"message": "MinerU+LLM 方法学提取完成",
|
||||
"paper_id": "demo_paper",
|
||||
"family": "",
|
||||
"source_filename": "wang2019.pdf",
|
||||
"stage_count": 6,
|
||||
"species_count": 1,
|
||||
"figure_table_caption_count": 9,
|
||||
"figure_table_stage_count": 12,
|
||||
"updated_at": "2026-08-07T09:42:17",
|
||||
"config_url": "/api/v1/methodology/extract-jobs/2943168e9c48/bbav5-config",
|
||||
"methodology_plan_url": "/api/v1/methodology/extract-jobs/2943168e9c48/plan",
|
||||
}
|
||||
],
|
||||
"count": 1,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
service = make_service()
|
||||
result = await service.list_jobs(limit=20)
|
||||
assert result["count"] == 1
|
||||
job = result["jobs"][0]
|
||||
assert job["title"] == "wang2019.pdf" # source_filename 优先
|
||||
assert job["progress_pct"] == 100
|
||||
assert job["status"] == "completed"
|
||||
await service.close()
|
||||
|
||||
async def test_primary_5xx_falls_back_to_alt(respx_mock):
|
||||
respx_mock.get(f"{PRIMARY}{PREFIX}/extract-jobs?limit=20").mock(
|
||||
return_value=httpx.Response(500)
|
||||
)
|
||||
respx_mock.get(f"{ALT}{PREFIX}/extract-jobs?limit=20").mock(
|
||||
return_value=httpx.Response(200, json={"jobs": [], "count": 0})
|
||||
)
|
||||
service = make_service()
|
||||
result = await service.list_jobs(limit=20)
|
||||
assert result == {"jobs": [], "count": 0}
|
||||
await service.close()
|
||||
|
||||
async def test_both_down_raises_unavailable(respx_mock):
|
||||
respx_mock.get(f"{PRIMARY}{PREFIX}/extract-jobs?limit=20").mock(
|
||||
side_effect=httpx.ConnectError("down")
|
||||
)
|
||||
respx_mock.get(f"{ALT}{PREFIX}/extract-jobs?limit=20").mock(
|
||||
side_effect=httpx.ConnectError("down too")
|
||||
)
|
||||
service = make_service()
|
||||
with pytest.raises(ExtractionUnavailable):
|
||||
await service.list_jobs(limit=20)
|
||||
await service.close()
|
||||
|
||||
async def test_alt_failure_not_cached_over_30s(respx_mock):
|
||||
"""alt 失败后 30s 内不再试 alt；超过 30s 恢复尝试。"""
|
||||
respx_mock.get(f"{PRIMARY}{PREFIX}/extract-jobs?limit=20").mock(
|
||||
side_effect=httpx.ConnectError("down")
|
||||
)
|
||||
alt_route = respx_mock.get(f"{ALT}{PREFIX}/extract-jobs?limit=20").mock(
|
||||
side_effect=httpx.ConnectError("down too")
|
||||
)
|
||||
|
||||
service = make_service()
|
||||
with pytest.raises(ExtractionUnavailable):
|
||||
await service.list_jobs(limit=20)
|
||||
assert alt_route.called is True # 第一次尝试了 alt
|
||||
|
||||
# 立刻再请求：alt 在失败窗口内，不应再被尝试
|
||||
before = alt_route.call_count
|
||||
with pytest.raises(ExtractionUnavailable):
|
||||
await service.list_jobs(limit=20)
|
||||
assert alt_route.call_count == before # 没有新增 alt 请求
|
||||
|
||||
# 模拟时间流逝超过 ALT_FAILURE_TTL
|
||||
service._alt_dead_until -= ALT_FAILURE_TTL + 1
|
||||
alt_route.side_effect = None # 必须先清 side_effect，respx 优先走 side_effect
|
||||
alt_route.return_value = httpx.Response(200, json={"jobs": [], "count": 0})
|
||||
result = await service.list_jobs(limit=20)
|
||||
assert result == {"jobs": [], "count": 0}
|
||||
assert alt_route.call_count == before + 1 # 恢复后重新尝试 alt
|
||||
await service.close()
|
||||
|
||||
|
||||
# ---------- jobs 归一化 ----------
|
||||
|
||||
|
||||
def _raw_job(**overrides):
|
||||
base = {
|
||||
"job_id": "a1b2c3d4e5f6",
|
||||
"status": "running",
|
||||
"stage": "llm",
|
||||
"substage": None,
|
||||
"message": "LLM 关系抽取中…",
|
||||
"paper_id": "pecan_spl",
|
||||
"family": "SPL",
|
||||
"source_filename": "Pecan_SPL_2021.pdf",
|
||||
"stage_count": 4,
|
||||
"species_count": 1,
|
||||
"figure_table_caption_count": 3,
|
||||
"figure_table_stage_count": 4,
|
||||
"updated_at": "2026-08-25T09:52:00",
|
||||
"config_url": None,
|
||||
"methodology_plan_url": None,
|
||||
}
|
||||
base.update(overrides)
|
||||
return base
|
||||
|
||||
async def test_jobs_normalization_stages(respx_mock):
|
||||
respx_mock.get(f"{PRIMARY}{PREFIX}/extract-jobs?limit=20").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"jobs": [
|
||||
_raw_job(stage="queued", status="queued", substage=None,
|
||||
message="排队等待处理"),
|
||||
_raw_job(stage="mineru", status="running"),
|
||||
_raw_job(stage="failed", status="failed",
|
||||
source_filename=None),
|
||||
],
|
||||
"count": 3,
|
||||
},
|
||||
)
|
||||
)
|
||||
service = make_service()
|
||||
result = await service.list_jobs(limit=20)
|
||||
jobs = result["jobs"]
|
||||
assert [j["progress_pct"] for j in jobs] == [5, 35, None]
|
||||
# failed 且无 source_filename：title 回退 paper_id
|
||||
assert jobs[2]["title"] == "pecan_spl"
|
||||
assert jobs[0]["stage"] == "queued"
|
||||
await service.close()
|
||||
|
||||
async def test_jobs_title_fallback_chain(respx_mock):
|
||||
respx_mock.get(f"{PRIMARY}{PREFIX}/extract-jobs?limit=20").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"jobs": [
|
||||
_raw_job(job_id="onlyjobid", source_filename=None, paper_id=None),
|
||||
],
|
||||
"count": 1,
|
||||
},
|
||||
)
|
||||
)
|
||||
service = make_service()
|
||||
result = await service.list_jobs(limit=20)
|
||||
assert result["jobs"][0]["title"] == "onlyjobid" # 最后回退 job_id
|
||||
await service.close()
|
||||
|
||||
|
||||
# ---------- job 详情与 plan ----------
|
||||
|
||||
|
||||
def _detail_job(**overrides):
|
||||
base = {
|
||||
"updated_at": "2026-08-07T09:42:17",
|
||||
"job_id": "2943168e9c48",
|
||||
"status": "completed",
|
||||
"stage": "completed",
|
||||
"substage": "done",
|
||||
"message": "MinerU+LLM 方法学提取完成",
|
||||
"source_path": "/home/xzq/5090/PlantAI-rag/data/evaluation/upload/wang2019.pdf",
|
||||
"content_list": "/home/xzq/5090/PlantAI-rag/.../wang2019_content_list.json",
|
||||
"output_dir": "/home/xzq/5090/PlantAI-rag/methodology_batch_demo_paper_2943168e9c48",
|
||||
"paper_id": "demo_paper",
|
||||
"family": "",
|
||||
"stage_count": 6,
|
||||
"species_count": 1,
|
||||
"figure_table_caption_count": 9,
|
||||
"figure_table_stage_count": 12,
|
||||
"figure_table_warning": "",
|
||||
"source_filename": "wang2019.pdf",
|
||||
"config_url": "/api/v1/methodology/extract-jobs/2943168e9c48/bbav5-config",
|
||||
"methodology_plan_url": "/api/v1/methodology/extract-jobs/2943168e9c48/plan",
|
||||
}
|
||||
base.update(overrides)
|
||||
return base
|
||||
|
||||
async def test_get_job_completed_fetches_plan(respx_mock):
|
||||
respx_mock.get(f"{PRIMARY}{PREFIX}/extract-jobs/2943168e9c48").mock(
|
||||
return_value=httpx.Response(200, json=_detail_job())
|
||||
)
|
||||
respx_mock.get(f"{PRIMARY}{PREFIX}/extract-jobs/2943168e9c48/plan").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"plan_id": "plan_2943168e9c48",
|
||||
"paper_id": "demo_paper",
|
||||
"gene_family": "",
|
||||
"stages": [],
|
||||
},
|
||||
)
|
||||
)
|
||||
service = make_service()
|
||||
detail = await service.get_job("2943168e9c48")
|
||||
assert detail["job_id"] == "2943168e9c48"
|
||||
assert detail["title"] == "wang2019.pdf"
|
||||
assert detail["progress_pct"] == 100
|
||||
assert detail["stats"] == {
|
||||
"stage_count": 6,
|
||||
"species_count": 1,
|
||||
"figure_table_caption_count": 9,
|
||||
"figure_table_stage_count": 12,
|
||||
}
|
||||
assert detail["detail"]["status"] == "completed" # 上游原始全量
|
||||
assert detail["plan"] == {
|
||||
"plan_id": "plan_2943168e9c48",
|
||||
"paper_id": "demo_paper",
|
||||
"gene_family": "",
|
||||
"stages": [],
|
||||
}
|
||||
await service.close()
|
||||
|
||||
async def test_get_job_incomplete_plan_409_is_null(respx_mock):
|
||||
"""未完成时上游 plan 返回 409：BFF 置 null 不报错（API.md §10）。"""
|
||||
respx_mock.get(f"{PRIMARY}{PREFIX}/extract-jobs/a1b2c3d4e5f6").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json=_detail_job(
|
||||
job_id="a1b2c3d4e5f6",
|
||||
status="running",
|
||||
stage="llm",
|
||||
substage=None,
|
||||
message="LLM 关系抽取中…",
|
||||
),
|
||||
)
|
||||
)
|
||||
respx_mock.get(f"{PRIMARY}{PREFIX}/extract-jobs/a1b2c3d4e5f6/plan").mock(
|
||||
return_value=httpx.Response(
|
||||
409, json={"detail": "Extraction job is not completed"}
|
||||
)
|
||||
)
|
||||
service = make_service()
|
||||
detail = await service.get_job("a1b2c3d4e5f6")
|
||||
assert detail["status"] == "running"
|
||||
assert detail["progress_pct"] == 75
|
||||
assert detail["plan"] is None # 409 被吞掉，不报错
|
||||
await service.close()
|
||||
|
||||
async def test_get_job_plan_upstream_down_yields_null_plan(respx_mock):
|
||||
respx_mock.get(f"{PRIMARY}{PREFIX}/extract-jobs/2943168e9c48").mock(
|
||||
return_value=httpx.Response(200, json=_detail_job())
|
||||
)
|
||||
respx_mock.get(f"{PRIMARY}{PREFIX}/extract-jobs/2943168e9c48/plan").mock(
|
||||
side_effect=httpx.ConnectError("plan service down")
|
||||
)
|
||||
respx_mock.get(f"{ALT}{PREFIX}/extract-jobs/2943168e9c48/plan").mock(
|
||||
side_effect=httpx.ConnectError("alt down")
|
||||
)
|
||||
service = make_service()
|
||||
detail = await service.get_job("2943168e9c48")
|
||||
assert detail["plan"] is None # plan 拉不到不影响 job 详情
|
||||
await service.close()
|
||||
|
||||
|
||||
# ---------- 透传 ----------
|
||||
|
||||
async def test_passthrough_run_log(respx_mock):
|
||||
respx_mock.get(f"{PRIMARY}{PREFIX}/run-log").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"text": "2026-08-25 10:00:01 INFO MinerU 解析完成\n"
|
||||
"2026-08-25 10:12:09 INFO 关系抽取 751 relations\n"
|
||||
},
|
||||
)
|
||||
)
|
||||
service = make_service()
|
||||
result = await service._passthrough("/run-log")
|
||||
assert result["text"].startswith("2026-08-25 10:00:01")
|
||||
await service.close()
|
||||
|
||||
|
||||
# ---------- 上传 ----------
|
||||
|
||||
async def _file_stream(chunks: list[bytes]):
|
||||
for c in chunks:
|
||||
yield c
|
||||
|
||||
async def test_upload_streams_and_appends_message(respx_mock):
|
||||
captured: dict = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured["url"] = str(request.url)
|
||||
captured["body"] = request.content # 多部分请求体已完整接收
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"success": True,
|
||||
"job_id": "d4e5f6a7b8c9",
|
||||
"status_url": "/api/v1/methodology/extract-jobs/d4e5f6a7b8c9",
|
||||
},
|
||||
)
|
||||
|
||||
respx_mock.post(f"{PRIMARY}{PREFIX}/extract-upload").mock(side_effect=handler)
|
||||
|
||||
service = make_service()
|
||||
result = await service.upload(
|
||||
filename="wang2019.pdf",
|
||||
stream=_file_stream([b"%PDF-1.4 mock bytes"]),
|
||||
content_type="application/pdf",
|
||||
paper_id="demo_paper",
|
||||
family="SPL",
|
||||
output_name="wang2019",
|
||||
)
|
||||
assert result == {
|
||||
"success": True,
|
||||
"job_id": "d4e5f6a7b8c9",
|
||||
"status_url": "/api/v1/methodology/extract-jobs/d4e5f6a7b8c9",
|
||||
"message": "已提交抽取任务",
|
||||
}
|
||||
# 请求被发到 primary（8000），multipart 含文件名、form 字段与文件字节
|
||||
assert captured["url"] == f"{PRIMARY}{PREFIX}/extract-upload"
|
||||
body = captured["body"]
|
||||
assert b'name="file"' in body
|
||||
assert b"wang2019.pdf" in body
|
||||
assert b'name="paper_id"' in body and b"demo_paper" in body
|
||||
assert b'name="family"' in body and b"SPL" in body
|
||||
assert b"%PDF-1.4 mock bytes" in body
|
||||
await service.close()
|
||||
@@ -0,0 +1,588 @@
|
||||
"""services/test_neo4j.py —— Neo4jService 单元测试（FakeDriver，不连真实实例）。
|
||||
|
||||
覆盖：实例选择与回退、5s 探测缓存、kg_status / kg_overview /
|
||||
search_entities / entity / neighborhood 的 Cypher 与结果组装、
|
||||
非 canonical 实例的 TODO 行为、超时与驱动错误 → Neo4jUnavailable。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from app.services.neo4j import (
|
||||
CONNECT_TIMEOUT,
|
||||
DIRECT_ACTIVE_NAME,
|
||||
FULLTEXT_INDEX,
|
||||
NAME_INDEX,
|
||||
NORMALIZED_ACTIVE_NAME,
|
||||
Neo4jService,
|
||||
Neo4jUnavailable,
|
||||
_Instance,
|
||||
)
|
||||
|
||||
|
||||
def inst(
|
||||
name: str = DIRECT_ACTIVE_NAME,
|
||||
uri: str = "bolt://fake-direct:18788",
|
||||
schema: str = "canonical",
|
||||
label: str | None = "SRT2CanonicalEntity",
|
||||
rel: str | None = "SRT2_CANONICAL_ASSERTS",
|
||||
) -> _Instance:
|
||||
return _Instance(name, uri, "neo4j", "pw", "neo4j", schema, label, rel)
|
||||
|
||||
|
||||
# ---------- FakeDriver ----------
|
||||
|
||||
|
||||
class FakeResult:
|
||||
"""模拟 neo4j AsyncResult：可迭代 record（dict 样式）。"""
|
||||
|
||||
def __init__(self, rows: list[dict]):
|
||||
self.rows = rows
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self.rows)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.rows)
|
||||
|
||||
|
||||
class FakeDriver:
|
||||
"""内存脚本表：{(cypher 关键片段): rows}。匹配取最长注册片段（避免短键误配）。"""
|
||||
|
||||
_NO_SCRIPT = object()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scripts: dict[str, list[dict]] | None = None,
|
||||
delays: dict[str, float] | None = None,
|
||||
errors: dict[str, Exception] | None = None,
|
||||
connect_error: Exception | None = None,
|
||||
):
|
||||
self.scripts = scripts or {}
|
||||
self.delays = delays or {}
|
||||
self.errors = errors or {}
|
||||
self.connect_error = connect_error
|
||||
self.calls: list[tuple[str, dict]] = []
|
||||
self.closed = False
|
||||
|
||||
def _match(self, cypher: str):
|
||||
# 大小写不敏感 + 最长片段优先，避免 "canonical_id: $id" 误配邻域查询
|
||||
lowered = cypher.lower()
|
||||
matches = [k for k in self.scripts if k.lower() in lowered]
|
||||
return max(matches, key=len) if matches else self._NO_SCRIPT
|
||||
|
||||
async def execute_query(
|
||||
self, cypher: str, parameters_: dict | None = None, **kwargs
|
||||
) -> tuple[FakeResult, object, object]:
|
||||
self.calls.append((cypher, dict(parameters_ or {})))
|
||||
if self.connect_error is not None:
|
||||
raise self.connect_error
|
||||
key = self._match(cypher)
|
||||
if key is self._NO_SCRIPT:
|
||||
raise AssertionError(
|
||||
f"未注册的 Cypher：{cypher!r}（已注册 {list(self.scripts)}）"
|
||||
)
|
||||
if key in self.errors:
|
||||
raise self.errors[key]
|
||||
if key in self.delays:
|
||||
await asyncio.sleep(self.delays[key])
|
||||
return FakeResult(self.scripts[key]), None, None
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
def _canonical_instance(
|
||||
name: str = DIRECT_ACTIVE_NAME,
|
||||
uri: str = "bolt://fake-direct:18788",
|
||||
) -> _Instance:
|
||||
return inst(name, uri, "canonical", "SRT2CanonicalEntity", "SRT2_CANONICAL_ASSERTS")
|
||||
|
||||
|
||||
def make_service(
|
||||
driver: FakeDriver,
|
||||
instances: list[_Instance] | None = None,
|
||||
) -> Neo4jService:
|
||||
"""构建已选中实例（跳过探测）的服务：默认只有 canonical 实例，
|
||||
连接指向测试 driver。"""
|
||||
if instances is None:
|
||||
instances = [_canonical_instance()]
|
||||
service = Neo4jService(instances=instances)
|
||||
service._drivers = {i.name: driver for i in instances}
|
||||
service._selected = True
|
||||
service._active = instances[0]
|
||||
return service
|
||||
|
||||
|
||||
COUNTS_SCRIPT = "RETURN entity_count, assertion_count"
|
||||
INDEXES_SCRIPT = "SHOW INDEXES"
|
||||
|
||||
|
||||
ENTITY_ROW = {
|
||||
"canonical_id": "cent_ospl14",
|
||||
"name": "OsSPL14",
|
||||
"entity_type": "gene",
|
||||
"support_count": 18,
|
||||
}
|
||||
|
||||
|
||||
def node_row(canonical_id: str, name: str, entity_type: str):
|
||||
return {
|
||||
"canonical_id": canonical_id,
|
||||
"name": name,
|
||||
"name_norm": name.lower(),
|
||||
"entity_type": entity_type,
|
||||
"support_count": 1,
|
||||
}
|
||||
|
||||
|
||||
def edge_row(src: str, tgt: str, category: str, *, title: str = "Paper"):
|
||||
return {
|
||||
"assertion_id": f"cass_{src}_{tgt}",
|
||||
"source_canonical_id": src,
|
||||
"target_canonical_id": tgt,
|
||||
"relation_norm": category,
|
||||
"relation_category": category,
|
||||
"paper_count": 3,
|
||||
"support_count": 2,
|
||||
"title": title,
|
||||
"doi": "10.1/x",
|
||||
"pmid": "1",
|
||||
}
|
||||
|
||||
|
||||
# ---------- 实例选择与探测 ----------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tcp_ok(monkeypatch):
|
||||
"""绕过 TCP 预检（fake URI 无法真实连接），让 FakeDriver 接管。"""
|
||||
|
||||
async def reachable(uri: str) -> bool:
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(Neo4jService, "_tcp_reachable", staticmethod(reachable))
|
||||
|
||||
|
||||
async def test_select_defaults_to_direct():
|
||||
service = make_service(FakeDriver({"RETURN 1": [{"x": 1}]}))
|
||||
instance = await service.select()
|
||||
assert instance.name == DIRECT_ACTIVE_NAME
|
||||
assert instance.uri == "bolt://fake-direct:18788"
|
||||
|
||||
|
||||
async def test_probe_falls_back_to_next_instance_on_connect_error(tcp_ok):
|
||||
from neo4j import exceptions as neo4j_exc
|
||||
|
||||
direct = _canonical_instance()
|
||||
failing = FakeDriver(
|
||||
{"RETURN 1": [{"x": 1}], COUNTS_SCRIPT: [], INDEXES_SCRIPT: []},
|
||||
connect_error=neo4j_exc.ServiceUnavailable("down"),
|
||||
)
|
||||
healthy = FakeDriver(
|
||||
{"RETURN 1": [{"x": 1}], COUNTS_SCRIPT: [], INDEXES_SCRIPT: []}
|
||||
)
|
||||
normal = inst(NORMALIZED_ACTIVE_NAME, "bolt://fake-normalized:19788",
|
||||
schema="normalized", label=None, rel=None)
|
||||
service = Neo4jService([direct, normal])
|
||||
service._drivers = {direct.name: failing, normal.name: healthy}
|
||||
|
||||
result = await service.probe()
|
||||
assert result["ok"] is True
|
||||
assert result["active"] == NORMALIZED_ACTIVE_NAME
|
||||
|
||||
async def test_probe_all_down_returns_ok_false(tcp_ok):
|
||||
from neo4j import exceptions as neo4j_exc
|
||||
|
||||
direct = _canonical_instance()
|
||||
failing = FakeDriver(
|
||||
{"RETURN 1": [{"x": 1}]},
|
||||
connect_error=neo4j_exc.ServiceUnavailable("down"),
|
||||
)
|
||||
service = Neo4jService([direct])
|
||||
service._drivers = {direct.name: failing}
|
||||
|
||||
result = await service.probe()
|
||||
assert result["ok"] is False
|
||||
assert result["active"] is None
|
||||
assert result["detail"]
|
||||
|
||||
# 失败结果缓存 5s：select() 也应抛
|
||||
with pytest.raises(Neo4jUnavailable):
|
||||
await service.select()
|
||||
|
||||
async def test_probe_result_cached_5s(tcp_ok):
|
||||
direct = _canonical_instance()
|
||||
driver = FakeDriver({"RETURN 1": [{"x": 1}], COUNTS_SCRIPT: [
|
||||
{"entity_count": 1, "assertion_count": 2}
|
||||
], INDEXES_SCRIPT: [{"name": NAME_INDEX}]})
|
||||
service = Neo4jService([direct])
|
||||
service._drivers = {direct.name: driver}
|
||||
|
||||
result1 = await service.probe()
|
||||
assert result1["ok"] is True
|
||||
assert result1["active"] == DIRECT_ACTIVE_NAME
|
||||
assert result1["detail"] is None
|
||||
n_calls = len(driver.calls)
|
||||
result2 = await service.probe() # 缓存命中：不再打实例
|
||||
assert result2 == result1
|
||||
assert len(driver.calls) == n_calls
|
||||
|
||||
|
||||
async def test_probe_extra_indexes_not_reported_as_missing(tcp_ok):
|
||||
"""BUG-R1-08 回归：同一标签上存在唯一约束索引等其它索引是正常现象
|
||||
（DISCOVERY §2.2 实测有 srt2_canonical_entity_id），probe 不得把
|
||||
「索引列表不止 name_norm 一个」误判为缺失或异常。"""
|
||||
direct = _canonical_instance()
|
||||
driver = FakeDriver({
|
||||
"RETURN 1": [{"x": 1}],
|
||||
COUNTS_SCRIPT: [{"entity_count": 1, "assertion_count": 2}],
|
||||
# 真实 SHOW INDEXES：name_norm RANGE 索引 + canonical_id 唯一约束索引
|
||||
INDEXES_SCRIPT: [
|
||||
{"name": NAME_INDEX},
|
||||
{"name": "srt2_canonical_entity_id"},
|
||||
],
|
||||
})
|
||||
service = Neo4jService([direct])
|
||||
service._drivers = {direct.name: driver}
|
||||
|
||||
result = await service.probe()
|
||||
assert result["ok"] is True
|
||||
assert result["detail"] is None # 修复前此处误报「缺失或异常」
|
||||
|
||||
|
||||
async def test_probe_missing_name_index_reports_detail(tcp_ok):
|
||||
direct = _canonical_instance()
|
||||
driver = FakeDriver({
|
||||
"RETURN 1": [{"x": 1}],
|
||||
COUNTS_SCRIPT: [{"entity_count": 1, "assertion_count": 2}],
|
||||
# 只有约束索引，没有 name_norm 索引
|
||||
INDEXES_SCRIPT: [{"name": "srt2_canonical_entity_id"}],
|
||||
})
|
||||
service = Neo4jService([direct])
|
||||
service._drivers = {direct.name: driver}
|
||||
|
||||
result = await service.probe()
|
||||
assert result["ok"] is True # 索引缺失不导致整体失败
|
||||
assert result["detail"] is not None
|
||||
assert "name_norm 索引缺失" in result["detail"]
|
||||
|
||||
|
||||
# ---------- kg_status ----------
|
||||
|
||||
async def test_kg_status_shape():
|
||||
driver = FakeDriver(
|
||||
{
|
||||
COUNTS_SCRIPT: [
|
||||
{"entity_count": 3820996, "assertion_count": 8246875}
|
||||
],
|
||||
INDEXES_SCRIPT: [{"name": NAME_INDEX}],
|
||||
}
|
||||
)
|
||||
service = make_service(driver)
|
||||
status = await service.kg_status()
|
||||
assert status == {
|
||||
"ok": True,
|
||||
"active": DIRECT_ACTIVE_NAME,
|
||||
"schema": "canonical",
|
||||
"labels": ["SRT2CanonicalEntity", "SRT2CanonicalAssertion", "SRT2Paper"],
|
||||
"relationship_types": ["SRT2_CANONICAL_ASSERTS", "SRT2_SUPPORTED_BY"],
|
||||
"entity_count": 3820996,
|
||||
"assertion_count": 8246875,
|
||||
"name_index": NAME_INDEX,
|
||||
"fulltext_index": FULLTEXT_INDEX,
|
||||
}
|
||||
|
||||
async def test_kg_status_missing_index_reports_null():
|
||||
driver = FakeDriver(
|
||||
{
|
||||
COUNTS_SCRIPT: [{"entity_count": 1, "assertion_count": 2}],
|
||||
INDEXES_SCRIPT: [],
|
||||
}
|
||||
)
|
||||
service = make_service(driver)
|
||||
status = await service.kg_status()
|
||||
assert status["name_index"] is None
|
||||
assert status["fulltext_index"] is None
|
||||
|
||||
async def test_non_canonical_instance_raises_unavailable():
|
||||
normal = inst(NORMALIZED_ACTIVE_NAME, "bolt://fake-normalized:19788",
|
||||
schema="normalized", label=None, rel=None)
|
||||
driver = FakeDriver({"RETURN 1": [{"x": 1}]})
|
||||
service = make_service(driver, [normal])
|
||||
with pytest.raises(Neo4jUnavailable):
|
||||
await service.kg_status()
|
||||
with pytest.raises(Neo4jUnavailable):
|
||||
await service.search_entities("OsSPL14")
|
||||
|
||||
|
||||
# ---------- kg_overview ----------
|
||||
|
||||
async def test_kg_overview_counts():
|
||||
driver = FakeDriver(
|
||||
{
|
||||
"WITH e.entity_type AS t, count(*) AS c": [
|
||||
{"type": "gene", "count": 43210},
|
||||
{"type": "gene_family", "count": 1288},
|
||||
],
|
||||
"WITH r.relation_category AS cat, count(*) AS c": [
|
||||
{"type": "characterized_in", "count": 12033},
|
||||
],
|
||||
}
|
||||
)
|
||||
service = make_service(driver)
|
||||
overview = await service.kg_overview()
|
||||
assert overview["entities_by_type"] == [
|
||||
{"type": "gene", "count": 43210},
|
||||
{"type": "gene_family", "count": 1288},
|
||||
]
|
||||
assert overview["relations_by_category"] == [
|
||||
{"type": "characterized_in", "count": 12033}
|
||||
]
|
||||
# 不能有拉全量边/实体的危险查询
|
||||
for cypher, _ in driver.calls:
|
||||
assert "RETURN e" not in cypher and "RETURN r" not in cypher
|
||||
|
||||
|
||||
# ---------- search_entities ----------
|
||||
|
||||
async def test_search_entities_maps_fields_and_normalizes_query():
|
||||
driver = FakeDriver(
|
||||
{
|
||||
"e.name_norm = $q_norm OR e.name_norm STARTS WITH $q_norm": [
|
||||
ENTITY_ROW,
|
||||
{**ENTITY_ROW, "canonical_id": "cent_ospl7", "name": "OsSPL7", "support_count": 9},
|
||||
]
|
||||
}
|
||||
)
|
||||
service = make_service(driver)
|
||||
results = await service.search_entities("OsSPL14", limit=5)
|
||||
assert results[0] == {
|
||||
"id": "cent_ospl14",
|
||||
"name": "OsSPL14",
|
||||
"type": "gene",
|
||||
"support_count": 18,
|
||||
}
|
||||
assert len(results) == 2
|
||||
# toLower 归一化已传到 Cypher 参数
|
||||
call_params = driver.calls[0][1]
|
||||
assert call_params["q_norm"] == "osspl14"
|
||||
assert call_params["limit"] == 5
|
||||
|
||||
async def test_search_entities_caps_limit_at_20():
|
||||
driver = FakeDriver(
|
||||
{"e.name_norm = $q_norm OR e.name_norm STARTS WITH $q_norm": [ENTITY_ROW]}
|
||||
)
|
||||
service = make_service(driver)
|
||||
await service.search_entities("x", limit=999)
|
||||
params = driver.calls[0][1]
|
||||
assert params["limit"] == 20
|
||||
|
||||
|
||||
# ---------- entity ----------
|
||||
|
||||
async def test_entity_by_canonical_id():
|
||||
props = {**node_row("cent_ospl14", "OsSPL14", "gene"), "rule": "exact"}
|
||||
driver = FakeDriver(
|
||||
{
|
||||
"canonical_id: $id": [{"canonical_id": "cent_ospl14"}],
|
||||
"RETURN e {.*} AS node": [{"node": props}],
|
||||
}
|
||||
)
|
||||
service = make_service(driver)
|
||||
entity = await service.entity("cent_ospl14")
|
||||
assert entity == {
|
||||
"id": "cent_ospl14",
|
||||
"name": "OsSPL14",
|
||||
"type": "gene",
|
||||
"properties": props,
|
||||
}
|
||||
|
||||
async def test_entity_by_name_uses_name_norm():
|
||||
driver = FakeDriver(
|
||||
{
|
||||
"name_norm: $name_norm": [{"canonical_id": "cent_ospl14"}],
|
||||
"RETURN e {.*} AS node": [{"node": node_row("cent_ospl14", "OsSPL14", "gene")}],
|
||||
}
|
||||
)
|
||||
service = make_service(driver)
|
||||
entity = await service.entity("OsSPL14")
|
||||
assert entity is not None and entity["id"] == "cent_ospl14"
|
||||
assert driver.calls[0][1]["name_norm"] == "osspl14"
|
||||
|
||||
async def test_entity_miss_returns_none():
|
||||
driver = FakeDriver({"name_norm: $name_norm": []})
|
||||
service = make_service(driver)
|
||||
assert await service.entity("NoSuchGene") is None
|
||||
|
||||
|
||||
# ---------- neighborhood ----------
|
||||
|
||||
async def test_neighborhood_depth1_shape():
|
||||
driver = FakeDriver(
|
||||
{
|
||||
COUNTS_SCRIPT: [],
|
||||
INDEXES_SCRIPT: [],
|
||||
"name_norm: $name_norm": [{"canonical_id": "cent_ospl14"}],
|
||||
"OPTIONAL MATCH (s)-[r1:SRT2_CANONICAL_ASSERTS]": [
|
||||
{
|
||||
"start": node_row("cent_ospl14", "OsSPL14", "gene"),
|
||||
"nodes": [
|
||||
node_row("cent_panicle", "panicle", "experimental_material"),
|
||||
node_row("cent_ospl14", "OsSPL14", "gene"), # 重复，应去重
|
||||
],
|
||||
"edges": [
|
||||
edge_row("cent_ospl14", "cent_panicle", "expressed_in"),
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
service = make_service(driver)
|
||||
result = await service.neighborhood("OsSPL14", depth=1, limit=50)
|
||||
assert result["center"] == "cent_ospl14"
|
||||
assert result["center_name"] == "OsSPL14"
|
||||
assert result["depth"] == 1
|
||||
assert result["nodes"] == [
|
||||
{"id": "cent_ospl14", "label": "OsSPL14", "type": "gene"},
|
||||
{"id": "cent_panicle", "label": "panicle", "type": "experimental_material"},
|
||||
]
|
||||
assert len(result["edges"]) == 1
|
||||
edge = result["edges"][0]
|
||||
assert edge["source"] == "cent_ospl14"
|
||||
assert edge["target"] == "cent_panicle"
|
||||
assert edge["relation"] == "expressed_in"
|
||||
assert edge["label"] == "expressed_in"
|
||||
assert edge["assertion_id"] == "cass_cent_ospl14_cent_panicle"
|
||||
assert edge["paper_count"] == 3
|
||||
assert edge["support_count"] == 2
|
||||
assert edge["title"] == "Paper"
|
||||
assert edge["doi"] == "10.1/x"
|
||||
assert edge["pmid"] == "1"
|
||||
assert edge["chain_id"] is None and edge["hop"] is None
|
||||
# limit 已下压进 Cypher 参数
|
||||
params = driver.calls[-1][1]
|
||||
assert params["max_nodes"] == 50 and params["max_edges"] == 50
|
||||
# depth=1 绝不能执行第二跳（两跳实测约 13s，不默认开启）
|
||||
for cypher, _ in driver.calls:
|
||||
assert "r2:SRT2_CANONICAL_ASSERTS" not in cypher
|
||||
|
||||
|
||||
async def test_neighborhood_depth2_uses_second_hop():
|
||||
driver = FakeDriver(
|
||||
{
|
||||
COUNTS_SCRIPT: [],
|
||||
INDEXES_SCRIPT: [],
|
||||
"canonical_id: $id": [{"canonical_id": "cent_ospl14"}],
|
||||
"OPTIONAL MATCH (s)-[r1:SRT2_CANONICAL_ASSERTS]": [
|
||||
{
|
||||
"start": node_row("cent_ospl14", "OsSPL14", "gene"),
|
||||
"nodes": [
|
||||
node_row("cent_panicle", "panicle", "experimental_material"),
|
||||
node_row("cent_branching", "panicle branching", "process"),
|
||||
],
|
||||
"edges": [edge_row("cent_ospl14", "cent_panicle", "expressed_in")],
|
||||
}
|
||||
],
|
||||
"OPTIONAL MATCH (n1)-[r2:SRT2_CANONICAL_ASSERTS]": [
|
||||
{
|
||||
"start": node_row("cent_ospl14", "OsSPL14", "gene"),
|
||||
"nodes": [
|
||||
node_row("cent_panicle", "panicle", "experimental_material"),
|
||||
node_row("cent_mir156", "miR156", "sequence_feature"),
|
||||
],
|
||||
"edges": [
|
||||
edge_row("cent_ospl14", "cent_panicle", "expressed_in"),
|
||||
edge_row("cent_panicle", "cent_mir156", "regulated_by"),
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
service = make_service(driver)
|
||||
result = await service.neighborhood("cent_ospl14", depth=2, limit=50)
|
||||
assert result["depth"] == 2
|
||||
# 一跳 + 二跳节点合并去重
|
||||
ids = [n["id"] for n in result["nodes"]]
|
||||
assert "cent_mir156" in ids
|
||||
assert len(ids) == len(set(ids))
|
||||
# 二跳边也出现（最长匹配选中的是 r2 脚本）
|
||||
relations = {e["relation"] for e in result["edges"]}
|
||||
assert "regulated_by" in relations
|
||||
# 确认深度 2 时执行了带 r2 的 Cypher
|
||||
assert any("r2:SRT2_CANONICAL_ASSERTS" in c for c, _ in driver.calls)
|
||||
|
||||
async def test_neighborhood_isolated_node_keeps_center():
|
||||
driver = FakeDriver(
|
||||
{
|
||||
"canonical_id: $id": [{"canonical_id": "cent_lonely"}],
|
||||
"OPTIONAL MATCH (s)-[r1:SRT2_CANONICAL_ASSERTS]": [
|
||||
{"start": node_row("cent_lonely", "lonely gene", "gene"),
|
||||
"nodes": [], "edges": []}
|
||||
],
|
||||
}
|
||||
)
|
||||
service = make_service(driver)
|
||||
result = await service.neighborhood("cent_lonely", depth=1, limit=50)
|
||||
assert result["center"] == "cent_lonely"
|
||||
assert result["nodes"] == [
|
||||
{"id": "cent_lonely", "label": "lonely gene", "type": "gene"}
|
||||
]
|
||||
assert result["edges"] == []
|
||||
|
||||
async def test_neighborhood_miss_returns_empty_center():
|
||||
driver = FakeDriver({"name_norm: $name_norm": []})
|
||||
service = make_service(driver)
|
||||
result = await service.neighborhood("NoSuch", depth=1, limit=50)
|
||||
assert result["center"] is None
|
||||
assert result["nodes"] == [] and result["edges"] == []
|
||||
|
||||
async def test_neighborhood_clamps_depth_and_limit():
|
||||
driver = FakeDriver(
|
||||
{
|
||||
"canonical_id: $id": [{"canonical_id": "cent_x"}],
|
||||
"OPTIONAL MATCH (s)-[r1:SRT2_CANONICAL_ASSERTS]": [
|
||||
{"start": node_row("cent_x", "x", "gene"), "nodes": [], "edges": []}
|
||||
],
|
||||
}
|
||||
)
|
||||
service = make_service(driver)
|
||||
result = await service.neighborhood("cent_x", depth=5, limit=9999)
|
||||
assert result["depth"] == 1
|
||||
params = driver.calls[-1][1]
|
||||
assert params["max_nodes"] == 200
|
||||
|
||||
|
||||
# ---------- 超时与驱动错误 ----------
|
||||
|
||||
async def test_query_timeout_raises_unavailable(monkeypatch):
|
||||
import app.services.neo4j as neo4j_mod
|
||||
|
||||
# 调短查询超时以在单测内触发 asyncio.wait_for
|
||||
monkeypatch.setattr(neo4j_mod, "QUERY_TIMEOUT", 0.05)
|
||||
slow = FakeDriver(
|
||||
{"name_norm: $name_norm": []},
|
||||
delays={"name_norm: $name_norm": 0.5},
|
||||
)
|
||||
service = make_service(slow)
|
||||
with pytest.raises(Neo4jUnavailable):
|
||||
await service.entity("OsSPL14")
|
||||
|
||||
async def test_driver_error_raises_unavailable():
|
||||
from neo4j import exceptions as neo4j_exc
|
||||
|
||||
failing = FakeDriver(
|
||||
{"name_norm: $name_norm": []},
|
||||
errors={"name_norm: $name_norm": neo4j_exc.ServiceUnavailable("boom")},
|
||||
)
|
||||
service = make_service(failing)
|
||||
with pytest.raises(Neo4jUnavailable):
|
||||
await service.entity("OsSPL14")
|
||||
|
||||
async def test_close_closes_all_drivers():
|
||||
driver = FakeDriver({})
|
||||
service = make_service(driver)
|
||||
await service.close()
|
||||
assert driver.closed is True
|
||||
@@ -0,0 +1,205 @@
|
||||
"""services/test_normalize.py —— 归一化器单测。
|
||||
|
||||
输入用 backend/mock/fixtures/upstream_real_shape_sample.json（真实上游形态样本，
|
||||
DISCOVERY.md §4.1），断言 API.md §7/§8/§9 的归一化输出关键字段。
|
||||
fixture 是单一事实源（ARCHITECTURE.md §6），本测试同时守护 mock/BFF 契约一致。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from app.services.normalize import (
|
||||
extract_chains,
|
||||
normalize_graph,
|
||||
normalize_hop,
|
||||
normalize_qa_response,
|
||||
normalize_source,
|
||||
pick,
|
||||
)
|
||||
|
||||
FIXTURE = (
|
||||
Path(__file__).resolve().parents[2]
|
||||
/ "mock" / "fixtures" / "upstream_real_shape_sample.json"
|
||||
)
|
||||
|
||||
|
||||
def load_raw() -> dict:
|
||||
return json.loads(FIXTURE.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# pick
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_pick_first_existing_key():
|
||||
obj = {"b": 2, "c": 3}
|
||||
assert pick(obj, "a", "b", "c") == 2
|
||||
assert pick(obj, "a") is None
|
||||
assert pick(obj, "a", default="x") == "x"
|
||||
assert pick(obj, "n", "b", default="y") == 2
|
||||
|
||||
|
||||
def test_pick_skips_none_values():
|
||||
assert pick({"a": None, "b": 1}, "a", "b") == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# §8 SourceItem
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_normalize_source_all_keys_present():
|
||||
src = normalize_source(load_raw()["sources"][0])
|
||||
expected_keys = {
|
||||
"id", "title", "paper_id", "chunk_id", "pmid", "pmcid", "doi", "url",
|
||||
"snippet", "score", "locator", "section", "support_count", "type", "metadata",
|
||||
}
|
||||
assert expected_keys <= set(src.keys()) # 所有字段键必须存在（缺失为 null）
|
||||
assert src["id"] == "S1" # evidence_id → id
|
||||
assert src["title"] == "Genome-wide analysis of the SPL gene family in pecan"
|
||||
assert src["paper_id"] == "PMC5488377_12864_2017_Article_3882"
|
||||
assert src["pmid"] == "28655295"
|
||||
assert src["doi"] == "10.1186/s12864-017-3882-4"
|
||||
assert src["score"] == 0.982 # final_score → score
|
||||
assert src["snippet"].startswith("## Phylogenetic")
|
||||
assert src["url"] == "https://doi.org/10.1186/s12864-017-3882-4"
|
||||
# 上游 item 的 source 是断言主体实体名（非论文标题），落 metadata.source_entity
|
||||
assert src["metadata"]["source_entity"] == "SPL gene family"
|
||||
assert src["metadata"]["assertion_id"] == "cass_d12f62f9abffbd903b97"
|
||||
assert src["metadata"]["vector_score"] == 0.714
|
||||
assert src["metadata"]["rerank_score"] == 0.989
|
||||
assert src["metadata"]["output_layer"] == "main_kg"
|
||||
|
||||
|
||||
def test_normalize_source_title_preferred_over_source():
|
||||
src = normalize_source({"title": "Real Title", "source": "Entity Name"})
|
||||
assert src["title"] == "Real Title"
|
||||
assert src["metadata"]["source_entity"] == "Entity Name"
|
||||
|
||||
|
||||
def test_normalize_source_only_doi_derives_url():
|
||||
src = normalize_source({"doi": "10.1000/xyz"})
|
||||
assert src["url"] == "https://doi.org/10.1000/xyz"
|
||||
assert src["title"] is None
|
||||
assert src["metadata"] == {}
|
||||
|
||||
|
||||
def test_normalize_source_alias_chains():
|
||||
src = normalize_source({
|
||||
"content": "snippet text",
|
||||
"doc_id": "doc-1",
|
||||
"confidence": 0.5,
|
||||
"assertion_id": "cass_x",
|
||||
})
|
||||
assert src["snippet"] == "snippet text"
|
||||
assert src["paper_id"] == "doc-1"
|
||||
assert src["score"] == 0.5
|
||||
assert src["metadata"]["assertion_id"] == "cass_x"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# §9 图谱
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_normalize_hop_aliases():
|
||||
hop = normalize_hop({
|
||||
"from_name": "A", "from_id": "cent_a", "source_type": "gene",
|
||||
"to_name": "B", "to_id": "cent_b", "target_type": "species",
|
||||
"relation": "characterized_in", "paper_id": "P1", "pmid": "1",
|
||||
"pmcid": None, "doi": "10.1/x",
|
||||
})
|
||||
assert hop["source"] == "A"
|
||||
assert hop["target"] == "B"
|
||||
assert hop["source_id"] == "cent_a"
|
||||
assert hop["target_id"] == "cent_b"
|
||||
assert hop["source_type"] == "gene"
|
||||
assert hop["target_type"] == "species"
|
||||
assert hop["relation"] == "characterized_in"
|
||||
assert hop["paper_id"] == "P1"
|
||||
|
||||
|
||||
def test_extract_chains_both_shapes():
|
||||
raw = load_raw()
|
||||
chains = extract_chains(raw) # 优先 native_graph_chains（列表形态）
|
||||
assert len(chains) == 1
|
||||
assert chains[0]["chain_id"] == "G1"
|
||||
assert len(chains[0]["hops"]) == 1
|
||||
|
||||
# dict 形态（graph_knowledge_chains）→ 取 .chains
|
||||
chains_dict = extract_chains({"graph_knowledge_chains": raw["graph_knowledge_chains"]})
|
||||
assert len(chains_dict) == 1
|
||||
assert chains_dict[0]["chain_id"] == "G1"
|
||||
assert len(chains_dict[0]["hops"]) == 2
|
||||
|
||||
|
||||
def test_normalize_graph_nodes_edges_chains():
|
||||
g = normalize_graph(load_raw())
|
||||
assert isinstance(g, dict) and set(g.keys()) == {"nodes", "edges", "chains"}
|
||||
assert len(g["chains"]) == 1
|
||||
assert g["chains"][0]["chain_id"] == "G1"
|
||||
assert len(g["chains"][0]["hops"]) == 1 # native_graph_chains 里只有 1 跳
|
||||
|
||||
# 节点：按 id 去重，label 取名字、type 取 source_type/target_type
|
||||
ids = [n["id"] for n in g["nodes"]]
|
||||
assert "cent_spl" in ids and "cent_oryza" in ids
|
||||
assert len(ids) == len(set(ids))
|
||||
by_id = {n["id"]: n for n in g["nodes"]}
|
||||
assert by_id["cent_spl"]["label"] == "SPL gene"
|
||||
assert by_id["cent_spl"]["type"] == "gene_family"
|
||||
assert by_id["cent_oryza"]["type"] == "species"
|
||||
|
||||
# 边：带 chain_id / hop / paper_id
|
||||
assert len(g["edges"]) == 1
|
||||
edge = g["edges"][0]
|
||||
assert edge["source"] == "cent_spl"
|
||||
assert edge["target"] == "cent_oryza"
|
||||
assert edge["relation"] == "characterized_in"
|
||||
assert edge["label"] == "characterized_in"
|
||||
assert edge["chain_id"] == "G1"
|
||||
assert edge["hop"] == 0
|
||||
assert edge["paper_id"] == "PMC5488377_12864_2017_Article_3882"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# §7 归一化 QA 响应
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_normalize_qa_response():
|
||||
out = normalize_qa_response(load_raw())
|
||||
assert out["query_id"] == "20260825_000715_49f801a9"
|
||||
assert out["version"] == "1.1.73"
|
||||
assert out["answer"] == "" # search 模式空
|
||||
assert out["success"] is True
|
||||
assert out["source_count"] == 1
|
||||
assert isinstance(out["sources"], list) and len(out["sources"]) == 1
|
||||
assert out["sources"][0]["id"] == "S1"
|
||||
assert out["used_sources"] == []
|
||||
# timing：真实键 timings → timing
|
||||
assert out["timing"]["total"] == 91.9
|
||||
# citation_validation → citation_map
|
||||
assert out["citation_map"]["valid"] is True
|
||||
# query_plan → retrieval_plan
|
||||
assert "agentic_plan" in out["retrieval_plan"]
|
||||
# graph 与 chains 同源
|
||||
assert len(out["graph"]["chains"]) == 1
|
||||
assert out["chains"][0]["chain_id"] == "G1"
|
||||
# raw 保存原始完整响应
|
||||
assert out["raw"]["question"] == "What is the SPL gene family?"
|
||||
|
||||
|
||||
def test_normalize_qa_response_alt_aliases():
|
||||
out = normalize_qa_response({
|
||||
"response": "answer text",
|
||||
"citations": [{"evidence_id": "E1", "paper_title": "T", "final_score": 0.9}],
|
||||
"used_sources": [],
|
||||
"agentic_plan": {"anchor": "x"},
|
||||
"timing": {"total_s": 1.0},
|
||||
"citation_map": {"valid": False},
|
||||
"notes": "n",
|
||||
})
|
||||
assert out["answer"] == "answer text"
|
||||
assert out["sources"][0]["title"] == "T"
|
||||
assert out["retrieval_plan"] == {"anchor": "x"}
|
||||
assert out["timing"] == {"total_s": 1.0}
|
||||
assert out["citation_map"] == {"valid": False}
|
||||
assert out["retrieval_notes"] == "n"
|
||||
@@ -0,0 +1,498 @@
|
||||
"""test_sse_relay.py —— sse_relay 单元测试（A2 归属）。
|
||||
|
||||
全部测试对注入的假数据/假 client 运行，不依赖真实上游。看门狗测试用缩短的
|
||||
超时参数注入（idle_timeout / keepalive_interval），避免真实 240s。
|
||||
环境无 respx/pytest-asyncio：异步测试用 anyio 插件（pytest.mark.anyio）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from app.services import sse_relay
|
||||
from app.services.sse_relay import (
|
||||
STAGE_LABEL,
|
||||
STAGE_PHASE,
|
||||
UNKNOWN_STAGE_PHASE,
|
||||
_classify,
|
||||
phase_for,
|
||||
relay_stream,
|
||||
)
|
||||
|
||||
|
||||
def _frame_texts(frames: list[str]) -> list[str]:
|
||||
return [f for f in frames if f.strip()]
|
||||
|
||||
|
||||
def _event_data_map(frames: list[str]) -> dict[str, list[dict]]:
|
||||
"""把 SSE 帧序列解析为 {event_name: [data_dict, ...]}。"""
|
||||
out: dict[str, list[dict]] = {}
|
||||
event = None
|
||||
for frame in frames:
|
||||
for line in frame.split("\n"):
|
||||
if line.startswith("event:"):
|
||||
event = line[len("event:"):].strip()
|
||||
out.setdefault(event, [])
|
||||
elif line.startswith("data:") and event:
|
||||
out[event].append(json.loads(line[len("data:"):].strip()))
|
||||
return out
|
||||
|
||||
|
||||
def _sse_bytes(events: list[tuple[bytes, bytes]]) -> list[bytes]:
|
||||
return [b"event: " + e + b"\ndata: " + d + b"\n\n" for e, d in events]
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# STAGE_PHASE / STAGE_LABEL 全覆盖
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
def test_stage_phase_maps_all_11_known_stages():
|
||||
expected = {
|
||||
"understanding_started": ("planning", "开始理解问题"),
|
||||
"entities_identified": ("planning", "已识别实体"),
|
||||
"related_search_started": ("planning", "开始关联检索"),
|
||||
"concepts_found": ("retrieval", "发现相关概念"),
|
||||
"subquestion_retrieved": ("retrieval", "子问题召回完成"),
|
||||
"agent_review_started": ("retrieval", "开始智能体审核"),
|
||||
"subquestion_retry_started": ("retrieval", "子问题重试中"),
|
||||
"evidence_accepted": ("evidence", "证据已采纳"),
|
||||
"evidence_review_completed": ("evidence", "证据审核完成"),
|
||||
"answer_generation_started": ("generation", "开始生成答案"),
|
||||
"completed": ("done", "检索完成"),
|
||||
}
|
||||
assert len(STAGE_PHASE) == 11
|
||||
assert len(STAGE_LABEL) == 11
|
||||
for stage, (phase, label) in expected.items():
|
||||
assert STAGE_PHASE[stage] == phase
|
||||
assert STAGE_LABEL[stage] == label
|
||||
assert phase_for("not_a_real_stage") == (UNKNOWN_STAGE_PHASE, "not_a_real_stage")
|
||||
|
||||
|
||||
def test_unknown_stage_falls_back_to_retrieval():
|
||||
assert STAGE_PHASE.get("future_stage_xyz") is None
|
||||
phase, label = phase_for("future_stage_xyz")
|
||||
assert phase == "retrieval"
|
||||
assert label == "future_stage_xyz"
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 帧分类：无 event 名按形态、多行 data、: 注释
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
def test_classify_stage_payload():
|
||||
kind, payload = _classify("stage", '{"stage":"entities_identified","details":{"entities":["OsSPL14"]}}')
|
||||
assert kind == "progress"
|
||||
assert payload["stage"] == "entities_identified"
|
||||
|
||||
|
||||
def test_classify_terminal_complete():
|
||||
kind, payload = _classify("complete", '{"answer":"xxx","success":true,"sources":[]}')
|
||||
assert kind == "complete"
|
||||
|
||||
|
||||
def test_classify_no_event_name_by_shape():
|
||||
# 无 event 名：按 payload 形态分类
|
||||
assert _classify(None, '{"stage":"concepts_found","details":{}}')[0] == "progress"
|
||||
assert _classify(None, '{"answer":"...","sources":[]}')[0] == "complete"
|
||||
assert _classify(None, '{"success":false,"detail":{"error_code":"x"}}')[0] == "complete" # 含 success 判终态
|
||||
assert _classify(None, '{"foo":1}')[0] == "raw"
|
||||
|
||||
|
||||
def test_classify_error_event():
|
||||
kind, payload = _classify("error", '{"detail":{"error_code":"invalid_llm_citation"}}')
|
||||
assert kind == "error"
|
||||
assert payload["code"] == "INVALID_CITATION"
|
||||
|
||||
|
||||
def test_classify_raw_never_drops():
|
||||
# 非 JSON data
|
||||
kind, payload = _classify(None, "not-json-at-all")
|
||||
assert kind == "raw"
|
||||
assert payload["payload"] == "not-json-at-all"
|
||||
# 无法分类的 JSON
|
||||
kind, payload = _classify("mystery", '{"unexpected":true}')
|
||||
assert kind == "raw"
|
||||
assert payload == {"unexpected": True}
|
||||
# JSON 数组
|
||||
kind, payload = _classify(None, "[1,2,3]")
|
||||
assert kind == "raw"
|
||||
assert payload == [1, 2, 3]
|
||||
# 空 data 但有 event 名
|
||||
kind, payload = _classify("some_event", "")
|
||||
assert kind == "raw"
|
||||
assert payload["event_name"] == "some_event"
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 中继整体：mock client 注入
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class FakeUpstreamStream:
|
||||
"""httpx stream 上下文的假实现：按给定 chunks 序列吐字节。
|
||||
|
||||
silent=True 时 aiter_raw 永久挂起（模拟上游连接保持但完全静默）。
|
||||
"""
|
||||
|
||||
def __init__(self, chunks: list[bytes], status: int = 200, headers: dict | None = None,
|
||||
json_body: dict | None = None, silent: bool = False):
|
||||
self._chunks = list(chunks)
|
||||
self.status_code = status
|
||||
self.headers = headers or {}
|
||||
self._json = json_body
|
||||
self._silent = silent
|
||||
|
||||
async def aiter_raw(self):
|
||||
if self._silent:
|
||||
import anyio
|
||||
|
||||
await anyio.sleep_forever()
|
||||
return # pragma: no cover
|
||||
for c in self._chunks:
|
||||
yield c
|
||||
|
||||
async def aread(self):
|
||||
return json.dumps(self._json).encode() if self._json is not None else b""
|
||||
|
||||
|
||||
class _StreamCtx:
|
||||
def __init__(self, response: FakeUpstreamStream):
|
||||
self._resp = response
|
||||
|
||||
async def __aenter__(self):
|
||||
return self._resp
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
|
||||
class FakeClient:
|
||||
"""按顺序弹出响应的假 httpx.AsyncClient；记录请求供断言。"""
|
||||
|
||||
def __init__(self, responses: list[FakeUpstreamStream], exc: Exception | None = None):
|
||||
self._responses = list(responses)
|
||||
self._exc = exc
|
||||
self.calls: list[dict] = []
|
||||
|
||||
def stream(self, method, url, **kwargs):
|
||||
self.calls.append({"method": method, "url": url, **kwargs})
|
||||
if self._exc is not None:
|
||||
exc, self._exc = self._exc, None
|
||||
raise exc
|
||||
resp = self._responses.pop(0)
|
||||
return _StreamCtx(resp)
|
||||
|
||||
|
||||
def _collect(agen) -> list[str]:
|
||||
import asyncio
|
||||
|
||||
async def run():
|
||||
return [f async for f in agen]
|
||||
|
||||
return asyncio.run(run())
|
||||
|
||||
|
||||
def test_relay_emits_progress_complete_sequence():
|
||||
import anyio
|
||||
|
||||
frames = [
|
||||
("stage", '{"stage":"understanding_started","details":{}}'),
|
||||
("stage", '{"stage":"entities_identified","details":{"entities":["OsSPL14"]}}'),
|
||||
("stage", '{"stage":"completed","details":{"elapsed_seconds":38.3}}'),
|
||||
("complete", '{"answer":"final answer","success":true,"sources":[{"title":"T","score":0.9}]}'),
|
||||
]
|
||||
client = FakeClient([FakeUpstreamStream(_sse_bytes([(e.encode(), d.encode()) for e, d in frames]))])
|
||||
|
||||
async def run():
|
||||
out = [f async for f in relay_stream({"question": "q"}, client=client, base_url="http://up")]
|
||||
return out
|
||||
|
||||
out = anyio.run(run)
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
assert len(events["progress"]) == 3
|
||||
assert events["progress"][0]["stage"] == "understanding_started"
|
||||
assert events["progress"][0]["phase"] == "planning"
|
||||
assert events["progress"][1]["phase"] == "planning"
|
||||
assert events["progress"][2]["phase"] == "done"
|
||||
assert events["complete"][0]["mode"] == "chat"
|
||||
assert "at" in events["progress"][0]
|
||||
assert client.calls[0]["url"] == "http://up/api/v1/qa-v1-1-73/chat/stream"
|
||||
assert client.calls[0]["json"] == {"question": "q"}
|
||||
|
||||
|
||||
def test_relay_multiline_data_and_comment_frames():
|
||||
import anyio
|
||||
|
||||
comment = b": keepalive\n\n"
|
||||
chunks = [
|
||||
b"event: stage\ndata: {\"stage\":\"evidence_accepted\",\n",
|
||||
b"data: \"details\":{\"accepted_count\":12}}\n\n",
|
||||
comment,
|
||||
b"event: complete\ndata: {\"answer\":\"A\"}\n\n",
|
||||
]
|
||||
client = FakeClient([FakeUpstreamStream(chunks)])
|
||||
|
||||
async def run():
|
||||
return [f async for f in relay_stream({"question": "q"}, client=client, base_url="http://up")]
|
||||
|
||||
out = anyio.run(run)
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
assert "progress" in events
|
||||
assert events["progress"][0]["details"] == {"accepted_count": 12}
|
||||
assert "complete" in events
|
||||
assert all(not f.startswith(": keepalive") for f in _frame_texts(out))
|
||||
|
||||
|
||||
def test_relay_no_event_name_classified_by_shape():
|
||||
import anyio
|
||||
|
||||
# 上游不带 event: 行，纯 data: JSON（模拟变体）
|
||||
chunks = [
|
||||
b'data: {"stage":"concepts_found","details":{}}\n\n',
|
||||
b'data: {"answer":"A","success":true}\n\n',
|
||||
]
|
||||
client = FakeClient([FakeUpstreamStream(chunks)])
|
||||
|
||||
async def run():
|
||||
return [f async for f in relay_stream({"q": 1}, client=client, base_url="http://up")]
|
||||
|
||||
out = anyio.run(run)
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
assert events["progress"][0]["phase"] == "retrieval"
|
||||
assert events["complete"][0]["answer"] == "A"
|
||||
|
||||
|
||||
def test_relay_unknown_json_goes_raw():
|
||||
import anyio
|
||||
|
||||
chunks = [
|
||||
b"data: {\"mystery\":\"payload\"}\n\n",
|
||||
b"event: complete\ndata: {\"answer\":\"A\"}\n\n",
|
||||
]
|
||||
client = FakeClient([FakeUpstreamStream(chunks)])
|
||||
|
||||
async def run():
|
||||
return [f async for f in relay_stream({}, client=client, base_url="http://up")]
|
||||
|
||||
out = anyio.run(run)
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
assert events["raw"] == [{"mystery": "payload"}]
|
||||
assert "complete" in events
|
||||
|
||||
|
||||
def test_relay_unknown_stage_maps_to_retrieval_with_raw_label():
|
||||
import anyio
|
||||
|
||||
chunks = [b'data: {"stage":"brand_new_stage","details":{"x":1}}\n\n']
|
||||
client = FakeClient([FakeUpstreamStream(chunks)])
|
||||
|
||||
async def run():
|
||||
return [f async for f in relay_stream({}, client=client, base_url="http://up")]
|
||||
|
||||
out = anyio.run(run)
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
assert events["progress"][0]["phase"] == "retrieval"
|
||||
assert events["progress"][0]["label"] == "brand_new_stage"
|
||||
|
||||
|
||||
def test_relay_error_event_closes():
|
||||
import anyio
|
||||
|
||||
chunks = [
|
||||
b'event: error\ndata: {"detail":{"error_code":"invalid_llm_citation"}}\n\n',
|
||||
b"data: {\"stage\":\"completed\"}\n\n", # 关闭后不应被处理
|
||||
]
|
||||
client = FakeClient([FakeUpstreamStream(chunks)])
|
||||
|
||||
async def run():
|
||||
return [f async for f in relay_stream({}, client=client, base_url="http://up")]
|
||||
|
||||
out = anyio.run(run)
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
assert list(events.keys()) == ["error"]
|
||||
assert events["error"][0]["code"] == "INVALID_CITATION"
|
||||
|
||||
|
||||
def test_relay_upstream_close_before_terminal_reports_error():
|
||||
import anyio
|
||||
|
||||
chunks = [b'data: {"stage":"concepts_found","details":{}}\n\n'] # 无 complete 即 EOF
|
||||
client = FakeClient([FakeUpstreamStream(chunks)])
|
||||
|
||||
async def run():
|
||||
return [f async for f in relay_stream({}, client=client, base_url="http://up")]
|
||||
|
||||
out = anyio.run(run)
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
# 残余帧先转发（绝不丢数据），再报错误
|
||||
assert "progress" in events
|
||||
assert events["error"][0]["code"] == "UPSTREAM_ERROR"
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 看门狗（缩短超时注入）
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
def test_watchdog_timeout_after_idle():
|
||||
import anyio
|
||||
|
||||
# 上游连接保持但完全静默；idle_timeout=0.3s 注入
|
||||
client = FakeClient([FakeUpstreamStream([], silent=True)])
|
||||
|
||||
async def run():
|
||||
return [
|
||||
f
|
||||
async for f in relay_stream(
|
||||
{}, client=client, base_url="http://up", idle_timeout=0.3, keepalive_interval=0.15
|
||||
)
|
||||
]
|
||||
|
||||
out = anyio.run(run)
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
assert list(events.keys()) == ["timeout"]
|
||||
assert events["timeout"][0]["message"] == "检索超时"
|
||||
|
||||
|
||||
def test_watchdog_ping_before_timeout():
|
||||
import anyio
|
||||
|
||||
client = FakeClient([FakeUpstreamStream([], silent=True)])
|
||||
|
||||
async def run():
|
||||
return [
|
||||
f
|
||||
async for f in relay_stream(
|
||||
{}, client=client, base_url="http://up", idle_timeout=0.6, keepalive_interval=0.2
|
||||
)
|
||||
]
|
||||
|
||||
out = anyio.run(run)
|
||||
pings = [f for f in _frame_texts(out) if f.startswith(": ping")]
|
||||
assert len(pings) >= 1
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
assert "timeout" in events
|
||||
|
||||
|
||||
def test_watchdog_does_not_fire_when_events_keep_coming():
|
||||
import anyio
|
||||
|
||||
# 上游在 idle_timeout 窗口内持续给事件（逐块喂，模拟慢速流）
|
||||
chunks = [
|
||||
b'data: {"stage":"understanding_started"}\n\n',
|
||||
b'data: {"stage":"entities_identified"}\n\n',
|
||||
b'event: complete\ndata: {"answer":"A"}\n\n',
|
||||
]
|
||||
client = FakeClient([FakeUpstreamStream(chunks)])
|
||||
|
||||
async def run():
|
||||
return [
|
||||
f
|
||||
async for f in relay_stream(
|
||||
{}, client=client, base_url="http://up", idle_timeout=10.0, keepalive_interval=0.05
|
||||
)
|
||||
]
|
||||
|
||||
out = anyio.run(run)
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
assert "complete" in events
|
||||
assert "timeout" not in events
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 上游 HTTP 头阶段失败：invalid_llm_citation 重试一次 → 仍失败 → error
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
def _http_error_response(body: dict, status: int) -> FakeUpstreamStream:
|
||||
return FakeUpstreamStream([], status=status, json_body=body)
|
||||
|
||||
|
||||
def test_http_error_402_invalid_citation_retries_then_errors():
|
||||
import anyio
|
||||
|
||||
client = FakeClient([
|
||||
_http_error_response({"detail": {"error_code": "invalid_llm_citation"}}, 402),
|
||||
_http_error_response({"detail": {"error_code": "invalid_llm_citation"}}, 402),
|
||||
])
|
||||
|
||||
async def run():
|
||||
return [f async for f in relay_stream({}, client=client, base_url="http://up")]
|
||||
|
||||
out = anyio.run(run)
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
assert list(events.keys()) == ["error"]
|
||||
assert events["error"][0]["code"] == "INVALID_CITATION"
|
||||
assert events["error"][0]["retryable"] is True
|
||||
assert len(client.calls) == 2 # 重试一次 = 总共 2 次请求
|
||||
|
||||
|
||||
def test_http_error_invalid_citation_retries_then_succeeds():
|
||||
import anyio
|
||||
|
||||
good = FakeUpstreamStream(
|
||||
[b'event: complete\ndata: {"answer":"ok"}\n\n'],
|
||||
status=200,
|
||||
)
|
||||
client = FakeClient([
|
||||
_http_error_response({"detail": {"error_code": "invalid_llm_citation"}}, 402),
|
||||
good,
|
||||
])
|
||||
|
||||
async def run():
|
||||
return [f async for f in relay_stream({}, client=client, base_url="http://up")]
|
||||
|
||||
out = anyio.run(run)
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
assert "complete" in events
|
||||
assert "error" not in events
|
||||
assert len(client.calls) == 2
|
||||
|
||||
|
||||
def test_http_error_other_status_no_retry():
|
||||
import anyio
|
||||
|
||||
client = FakeClient([
|
||||
_http_error_response({"detail": "boom"}, 500),
|
||||
])
|
||||
|
||||
async def run():
|
||||
return [f async for f in relay_stream({}, client=client, base_url="http://up")]
|
||||
|
||||
out = anyio.run(run)
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
assert events["error"][0]["code"] == "UPSTREAM_ERROR"
|
||||
assert len(client.calls) == 1
|
||||
|
||||
|
||||
def test_connect_error_emits_upstream_down():
|
||||
import anyio
|
||||
import httpx
|
||||
|
||||
client = FakeClient([], exc=httpx.ConnectError("conn refused", request=None))
|
||||
|
||||
async def run():
|
||||
return [f async for f in relay_stream({}, client=client, base_url="http://up")]
|
||||
|
||||
out = anyio.run(run)
|
||||
events = _event_data_map(_frame_texts(out))
|
||||
assert events["error"][0]["code"] == "UPSTREAM_DOWN"
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 工具函数
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
def test_sse_headers_contract():
|
||||
assert sse_relay.sse_headers["Cache-Control"] == "no-cache"
|
||||
assert sse_relay.sse_headers["X-Accel-Buffering"] == "no"
|
||||
|
||||
|
||||
def test_error_payload_shapes():
|
||||
assert sse_relay.error_payload("TIMEOUT") == {"code": "TIMEOUT", "message": "检索超时", "retryable": True}
|
||||
assert sse_relay.error_payload("UPSTREAM_DOWN")["retryable"] is False
|
||||
assert sse_relay.error_payload("INVALID_CITATION")["retryable"] is True
|
||||
@@ -0,0 +1,75 @@
|
||||
"""app/test_db.py —— SQLite 层单测（aiosqlite + WAL，幂等迁移，JSON 编解码）。
|
||||
|
||||
覆盖：connect 建表/索引、重复 connect 幂等、外键级联删除、JSON 编解码容错。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from app.db import Database, dumps_json, loads_json
|
||||
|
||||
|
||||
async def test_connect_creates_schema_and_wal(tmp_path):
|
||||
db = Database(tmp_path / "t1" / "plantai.db")
|
||||
await db.connect()
|
||||
try:
|
||||
# WAL 模式
|
||||
row = await db.fetchone("PRAGMA journal_mode")
|
||||
assert row[0].lower() == "wal"
|
||||
# 表存在
|
||||
tables = await db.fetchall("SELECT name FROM sqlite_master WHERE type='table'")
|
||||
names = {t["name"] for t in tables}
|
||||
assert {"sessions", "messages"} <= names
|
||||
# 索引存在
|
||||
indexes = await db.fetchall(
|
||||
"SELECT name FROM sqlite_master WHERE type='index' AND name LIKE 'idx_%'"
|
||||
)
|
||||
inames = {i["name"] for i in indexes}
|
||||
assert "idx_messages_session" in inames
|
||||
assert "idx_sessions_updated" in inames
|
||||
finally:
|
||||
await db.close()
|
||||
|
||||
|
||||
async def test_connect_idempotent(tmp_path):
|
||||
db = Database(tmp_path / "t2" / "plantai.db")
|
||||
await db.connect()
|
||||
await db.close()
|
||||
await db.connect() # 二次迁移不报错
|
||||
await db.close()
|
||||
|
||||
|
||||
async def test_cascade_delete_messages(tmp_path):
|
||||
db = Database(tmp_path / "t3" / "plantai.db")
|
||||
await db.connect()
|
||||
try:
|
||||
await db.execute(
|
||||
"INSERT INTO sessions (id, title, status, meta, created_at, updated_at) "
|
||||
"VALUES ('s1', 't', 'idle', '{}', '2026-01-01T00:00:00Z', '2026-01-01T00:00:00Z')"
|
||||
)
|
||||
await db.execute(
|
||||
"INSERT INTO messages (id, session_id, role, kind, status, blocks, raw, text, "
|
||||
"created_at, updated_at) VALUES "
|
||||
"('m1', 's1', 'user', 'qa', 'done', '[]', NULL, 'hi', "
|
||||
"'2026-01-01T00:00:00Z', '2026-01-01T00:00:00Z')"
|
||||
)
|
||||
row = await db.fetchone("SELECT COUNT(*) AS n FROM messages WHERE session_id='s1'")
|
||||
assert row["n"] == 1
|
||||
await db.execute("DELETE FROM sessions WHERE id='s1'")
|
||||
row = await db.fetchone("SELECT COUNT(*) AS n FROM messages WHERE session_id='s1'")
|
||||
assert row["n"] == 0 # ON DELETE CASCADE
|
||||
finally:
|
||||
await db.close()
|
||||
|
||||
|
||||
def test_dumps_loads_json_roundtrip():
|
||||
payload = {"a": 1, "b": [1, 2], "c": {"d": "中文"}}
|
||||
assert loads_json(dumps_json(payload)) == payload
|
||||
assert loads_json(None, default=[]) == []
|
||||
assert loads_json("not-json", default={}) == {}
|
||||
|
||||
|
||||
def test_dumps_json_fallback_on_unserializable():
|
||||
import threading
|
||||
|
||||
assert dumps_json({"lock": threading.Lock()}, default=None) == "null"
|
||||
@@ -0,0 +1,243 @@
|
||||
"""app/test_sessions_api.py —— 会话/消息路由集成单测（API.md §5 / §1 信封）。
|
||||
|
||||
TestClient + 临时 SQLite（dependency_overrides 指向临时 Database），
|
||||
不触发 app lifespan（不跑 init_db/上游探测）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.db import Database, dumps_json, get_db
|
||||
from app.main import app
|
||||
from app.upstream import envelope
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
async def db(tmp_path):
|
||||
database = Database(tmp_path / "api" / "plantai.db")
|
||||
await database.connect()
|
||||
yield database
|
||||
await database.close()
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def client(db):
|
||||
app.dependency_overrides[get_db] = lambda: db
|
||||
try:
|
||||
yield TestClient(app)
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 错误信封构造器（§1）
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"code,expected_status",
|
||||
[
|
||||
("UPSTREAM_DOWN", 503),
|
||||
("UPSTREAM_ERROR", 502),
|
||||
("TIMEOUT", 504),
|
||||
("INVALID_CITATION", 502),
|
||||
("NOT_FOUND", 404),
|
||||
("VALIDATION", 422),
|
||||
("INTERNAL", 500),
|
||||
],
|
||||
)
|
||||
def test_envelope_defaults(code, expected_status):
|
||||
body, status = envelope(code)
|
||||
assert status == expected_status
|
||||
assert body["ok"] is False
|
||||
assert body["error"]["code"] == code
|
||||
assert isinstance(body["error"]["message"], str) and body["error"]["message"]
|
||||
assert "degraded" in body["error"]
|
||||
|
||||
|
||||
def test_envelope_custom():
|
||||
body, status = envelope("UPSTREAM_DOWN", message="自定义", degraded=True, status=503)
|
||||
assert body["error"]["message"] == "自定义"
|
||||
assert body["error"]["degraded"] is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 会话 CRUD 全链路
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_session_crud_full_flow(client):
|
||||
# POST（默认标题 + 自定义标题）
|
||||
r = client.post("/api/sessions", json={})
|
||||
assert r.status_code == 201, r.text
|
||||
sess = r.json()
|
||||
assert sess["title"] == "新会话"
|
||||
assert sess["status"] == "idle"
|
||||
assert sess["message_count"] == 0
|
||||
assert sess["meta"] == {"title_manual": False}
|
||||
|
||||
r2 = client.post("/api/sessions", json={"title": " 水稻SPL "})
|
||||
sess2 = r2.json()
|
||||
assert sess2["title"] == "水稻SPL"
|
||||
|
||||
# GET 列表（updated_at 倒序，最新在前）
|
||||
r = client.get("/api/sessions")
|
||||
assert r.status_code == 200
|
||||
items = r.json()
|
||||
assert [s["id"] for s in items][:2] == [sess2["id"], sess["id"]]
|
||||
|
||||
# POST user 消息
|
||||
r = client.post(f"/api/sessions/{sess['id']}/messages", json={"role": "user", "text": "问题原文"})
|
||||
assert r.status_code == 201, r.text
|
||||
msg = r.json()
|
||||
assert msg["role"] == "user" and msg["status"] == "done"
|
||||
assert msg["text"] == "问题原文"
|
||||
assert msg["blocks"] == [] and msg["raw"] is None
|
||||
|
||||
# POST assistant 流式骨架
|
||||
r = client.post(
|
||||
f"/api/sessions/{sess['id']}/messages",
|
||||
json={
|
||||
"role": "assistant", "kind": "qa", "status": "streaming",
|
||||
"blocks": [{"type": "pipeline", "status": "running", "steps": []}],
|
||||
},
|
||||
)
|
||||
assert r.status_code == 201, r.text
|
||||
skeleton = r.json()
|
||||
assert skeleton["status"] == "streaming"
|
||||
assert skeleton["blocks"][0]["type"] == "pipeline"
|
||||
assert "status" in skeleton["blocks"][0]
|
||||
|
||||
# 消息活动后会话状态 running、message_count 更新
|
||||
r = client.get(f"/api/sessions/{sess['id']}")
|
||||
detail = r.json()
|
||||
assert detail["status"] == "running"
|
||||
assert detail["message_count"] == 2
|
||||
assert [m["role"] for m in detail["messages"]] == ["user", "assistant"]
|
||||
|
||||
# PATCH 消息（流式检查点：status/blocks/raw）
|
||||
r = client.patch(
|
||||
f"/api/sessions/{sess['id']}/messages/{skeleton['id']}",
|
||||
json={
|
||||
"status": "done",
|
||||
"blocks": [{"type": "text", "html": "<p>最终答案</p>"}],
|
||||
"raw": {"answer": "最终答案"},
|
||||
},
|
||||
)
|
||||
assert r.status_code == 200, r.text
|
||||
updated = r.json()
|
||||
assert updated["status"] == "done"
|
||||
assert updated["blocks"] == [{"type": "text", "html": "<p>最终答案</p>"}]
|
||||
assert updated["raw"] == {"answer": "最终答案"}
|
||||
|
||||
# PATCH 会话标题 → title_manual=true
|
||||
r = client.patch(f"/api/sessions/{sess['id']}", json={"title": "手工标题"})
|
||||
assert r.status_code == 200
|
||||
assert r.json()["meta"] == {"title_manual": True}
|
||||
|
||||
# PATCH 会话状态
|
||||
r = client.patch(f"/api/sessions/{sess['id']}", json={"status": "done"})
|
||||
assert r.status_code == 200
|
||||
assert r.json()["status"] == "done"
|
||||
|
||||
# DELETE 204 + 级联
|
||||
r = client.delete(f"/api/sessions/{sess['id']}")
|
||||
assert r.status_code == 204
|
||||
r = client.get(f"/api/sessions/{sess['id']}")
|
||||
assert r.status_code == 404
|
||||
body = r.json()
|
||||
assert body["ok"] is False
|
||||
assert body["error"]["code"] == "NOT_FOUND"
|
||||
|
||||
|
||||
def test_not_found_envelopes(client):
|
||||
cases = [
|
||||
("get", "/api/sessions/nope", None),
|
||||
("patch", "/api/sessions/nope", {"title": "x"}),
|
||||
("delete", "/api/sessions/nope", None),
|
||||
("post", "/api/sessions/nope/messages", {"role": "user", "text": "hi"}),
|
||||
("get", "/api/sessions/nope/sources", None),
|
||||
]
|
||||
for method, url, body in cases:
|
||||
kwargs = {"json": body} if body is not None else {}
|
||||
r = getattr(client, method)(url, **kwargs)
|
||||
assert r.status_code == 404, (method, url)
|
||||
payload = r.json()
|
||||
assert payload["ok"] is False
|
||||
assert payload["error"]["code"] == "NOT_FOUND"
|
||||
|
||||
r = client.patch("/api/sessions/nope/messages/nope-msg", json={"status": "done"})
|
||||
assert r.status_code == 404
|
||||
assert r.json()["error"]["code"] == "NOT_FOUND"
|
||||
|
||||
|
||||
def test_message_404_when_session_exists_but_message_missing(client):
|
||||
r = client.post("/api/sessions", json={"title": "t"})
|
||||
sid = r.json()["id"]
|
||||
r = client.patch(f"/api/sessions/{sid}/messages/ghost", json={"status": "done"})
|
||||
assert r.status_code == 404
|
||||
assert r.json()["error"]["code"] == "NOT_FOUND"
|
||||
|
||||
|
||||
def test_empty_title_falls_back_to_default(client):
|
||||
r = client.post("/api/sessions", json={"title": " "})
|
||||
assert r.json()["title"] == "新会话"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 来源聚合（GET /api/sessions/{id}/sources）
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _source_item(i: str, paper: str | None, title: str) -> dict:
|
||||
return {
|
||||
"id": i, "title": title, "paper_id": paper, "chunk_id": None,
|
||||
"pmid": None, "pmcid": None, "doi": None, "url": None,
|
||||
"snippet": "…", "score": 0.9, "locator": None, "section": None,
|
||||
"support_count": None, "type": "assertion", "metadata": {},
|
||||
}
|
||||
|
||||
|
||||
async def test_sources_aggregation(client, db):
|
||||
r = client.post("/api/sessions", json={"title": "聚合"})
|
||||
sid = r.json()["id"]
|
||||
|
||||
# 两个助手消息各带一个 sources 块
|
||||
blocks1 = [
|
||||
{
|
||||
"type": "sources",
|
||||
"items": [
|
||||
_source_item("a1", "PAPER_A", "Paper A"),
|
||||
_source_item("b1", "PAPER_B", "Paper B"),
|
||||
_source_item("x1", None, "No paper id"),
|
||||
],
|
||||
}
|
||||
]
|
||||
blocks2 = [
|
||||
{
|
||||
"type": "sources",
|
||||
"items": [
|
||||
_source_item("a2", "PAPER_A", "Paper A"),
|
||||
_source_item("x2", None, "No paper id"),
|
||||
],
|
||||
}
|
||||
]
|
||||
await db.execute(
|
||||
"INSERT INTO messages (id, session_id, role, kind, status, blocks, raw, text, "
|
||||
"created_at, updated_at) VALUES "
|
||||
"('m1', ?, 'assistant', 'qa', 'done', ?, NULL, NULL, '2026-01-01T00:00:00Z', '2026-01-01T00:00:00Z'), "
|
||||
"('m2', ?, 'assistant', 'qa', 'done', ?, NULL, NULL, '2026-01-01T00:00:00Z', '2026-01-01T00:00:00Z')",
|
||||
(sid, dumps_json(blocks1), sid, dumps_json(blocks2)),
|
||||
)
|
||||
|
||||
r = client.get(f"/api/sessions/{sid}/sources")
|
||||
assert r.status_code == 200, r.text
|
||||
body = r.json()
|
||||
assert body["session_id"] == sid
|
||||
# PAPER_A 出现 2 次（2 个消息各 1 条），PAPER_B 1 次；无 paper_id 的按 id 各 1 次
|
||||
assert body["total"] == 5
|
||||
by_key = {(s["paper_id"] or s["id"]): s for s in body["sources"]}
|
||||
assert by_key["PAPER_A"]["count"] == 2
|
||||
assert by_key["PAPER_A"]["title"] == "Paper A"
|
||||
assert by_key["PAPER_B"]["count"] == 1
|
||||
assert by_key["x1"]["count"] == 1
|
||||
assert by_key["x2"]["count"] == 1
|
||||
assert len(body["sources"]) == 4
|
||||
@@ -0,0 +1,195 @@
|
||||
"""upstream.py —— 上游健康探测（httpx 单例 + 5s TTL 缓存）与统一错误信封。
|
||||
|
||||
- `probe_all()`：并发探测 qa（/health 或 /info）、rag（/health）、ragAlt（/health）、
|
||||
Neo4j（延迟导入 `app.services.neo4j.probe_neo4j()`，A3 提供；若模块缺失则置
|
||||
`{ok:false}`）。探测超时 1s（settings.timeout_health），结果缓存 5s。
|
||||
- `envelope(code, message, degraded, status)`：API.md §1 错误信封构造器，供所有 router 用。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Awaitable, Callable
|
||||
|
||||
import httpx
|
||||
|
||||
from .config import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_client: httpx.AsyncClient | None = None
|
||||
_httpx_lock = asyncio.Lock()
|
||||
|
||||
|
||||
async def get_client() -> httpx.AsyncClient:
|
||||
"""httpx.AsyncClient 单例（供 qa/kg/extraction 各代理层共用）。"""
|
||||
global _client
|
||||
if _client is None or _client.is_closed:
|
||||
async with _httpx_lock:
|
||||
if _client is None or _client.is_closed:
|
||||
settings = get_settings()
|
||||
_client = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(
|
||||
connect=settings.timeout_health,
|
||||
read=settings.timeout_qa_chat,
|
||||
write=10.0,
|
||||
pool=5.0,
|
||||
),
|
||||
follow_redirects=True,
|
||||
)
|
||||
return _client
|
||||
|
||||
|
||||
async def close_client() -> None:
|
||||
"""app shutdown 钩子。"""
|
||||
global _client
|
||||
if _client is not None and not _client.is_closed:
|
||||
await _client.aclose()
|
||||
_client = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 错误信封（API.md §1）——供所有 router 用
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_ERROR_MESSAGES: dict[str, str] = {
|
||||
"UPSTREAM_DOWN": "知识问答服务当前不可用，请稍后重试",
|
||||
"UPSTREAM_ERROR": "上游服务返回错误",
|
||||
"TIMEOUT": "检索超时",
|
||||
"INVALID_CITATION": "引用校验失败",
|
||||
"NOT_FOUND": "资源不存在",
|
||||
"VALIDATION": "请求参数不合法",
|
||||
"INTERNAL": "服务内部错误",
|
||||
}
|
||||
|
||||
|
||||
def envelope(
|
||||
code: str,
|
||||
message: str | None = None,
|
||||
degraded: bool = False,
|
||||
status: int | None = None,
|
||||
) -> tuple[dict[str, Any], int]:
|
||||
"""构造 `{"ok": false, "error": {code, message, degraded}}` 信封 + HTTP 状态。
|
||||
|
||||
status 缺省时按 code 推导（API.md §1 表）：
|
||||
UPSTREAM_DOWN=503, UPSTREAM_ERROR=502, TIMEOUT=504, INVALID_CITATION=502,
|
||||
NOT_FOUND=404, VALIDATION=422, INTERNAL=500。
|
||||
"""
|
||||
if status is None:
|
||||
status = {
|
||||
"UPSTREAM_DOWN": 503,
|
||||
"UPSTREAM_ERROR": 502,
|
||||
"TIMEOUT": 504,
|
||||
"INVALID_CITATION": 502,
|
||||
"NOT_FOUND": 404,
|
||||
"VALIDATION": 422,
|
||||
"INTERNAL": 500,
|
||||
}.get(code, 500)
|
||||
if message is None:
|
||||
message = _ERROR_MESSAGES.get(code, "未知错误")
|
||||
return {
|
||||
"ok": False,
|
||||
"error": {"code": code, "message": message, "degraded": degraded},
|
||||
}, status
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 健康探测
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_probe_lock = asyncio.Lock()
|
||||
_probe_at: float = 0.0
|
||||
_probe_result: dict[str, Any] | None = None
|
||||
|
||||
|
||||
async def _get_json(url: str) -> tuple[bool, Any]:
|
||||
"""1s 超时 GET + JSON 解析；任何失败返回 (False, None)，不抛错。"""
|
||||
try:
|
||||
client = await get_client()
|
||||
resp = await client.get(url, timeout=httpx.Timeout(get_settings().timeout_health))
|
||||
if resp.status_code >= 400:
|
||||
return False, None
|
||||
return True, resp.json()
|
||||
except Exception:
|
||||
return False, None
|
||||
|
||||
|
||||
async def _probe_qa() -> dict[str, Any]:
|
||||
settings = get_settings()
|
||||
base = settings.qa_base_url.rstrip("/")
|
||||
# 按 DISCOVERY.md：health 为可靠探测端点；其不可用时回落 info
|
||||
ok, data = await _get_json(f"{base}{settings.qa_api_path}/health")
|
||||
if not ok:
|
||||
ok, data = await _get_json(f"{base}{settings.qa_api_path}/info")
|
||||
if not ok or not isinstance(data, dict):
|
||||
return {"ok": False}
|
||||
version = data.get("version") or data.get("app_version")
|
||||
return {"ok": True, "version": version}
|
||||
|
||||
|
||||
async def _probe_rag(base_url: str) -> dict[str, Any]:
|
||||
ok, data = await _get_json(f"{base_url.rstrip('/')}/health")
|
||||
if not ok:
|
||||
return {"ok": False}
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
async def _probe_neo4j() -> dict[str, Any]:
|
||||
"""Neo4j 状态由 A3 提供（app.services.neo4j.probe_neo4j）。
|
||||
延迟导入：A3 未实现时 catch ImportError，置 {ok:false}。
|
||||
"""
|
||||
try:
|
||||
from .services.neo4j import probe_neo4j
|
||||
except ImportError:
|
||||
return {"ok": False}
|
||||
try:
|
||||
result = await probe_neo4j()
|
||||
except Exception:
|
||||
return {"ok": False}
|
||||
if not isinstance(result, dict):
|
||||
return {"ok": False}
|
||||
out: dict[str, Any] = {"ok": bool(result.get("ok"))}
|
||||
if result.get("active"):
|
||||
out["active"] = result["active"]
|
||||
if result.get("detail"):
|
||||
out["detail"] = result["detail"]
|
||||
return out
|
||||
|
||||
|
||||
async def _probe_all_uncached() -> dict[str, Any]:
|
||||
settings = get_settings()
|
||||
qa_t, rag_t, alt_t, neo_t = await asyncio.gather(
|
||||
_probe_qa(),
|
||||
_probe_rag(settings.rag_base_url),
|
||||
_probe_rag(settings.rag_alt_base_url),
|
||||
_probe_neo4j(),
|
||||
)
|
||||
return {"qa": qa_t, "rag": rag_t, "ragAlt": alt_t, "neo4j": neo_t}
|
||||
|
||||
|
||||
async def probe_all(force: bool = False) -> dict[str, Any]:
|
||||
"""并发探测全部上游；结果缓存 settings.health_cache_ttl 秒（默认 5）。"""
|
||||
global _probe_at, _probe_result
|
||||
now = time.monotonic()
|
||||
if not force and _probe_result is not None:
|
||||
if now - _probe_at < get_settings().health_cache_ttl:
|
||||
return _probe_result
|
||||
async with _probe_lock:
|
||||
now = time.monotonic()
|
||||
if not force and _probe_result is not None:
|
||||
if now - _probe_at < get_settings().health_cache_ttl:
|
||||
return _probe_result
|
||||
result = await _probe_all_uncached()
|
||||
_probe_at = time.monotonic()
|
||||
_probe_result = result
|
||||
return result
|
||||
|
||||
|
||||
async def probe_all_refresh() -> dict[str, Any]:
|
||||
"""强制刷新（无缓存）。"""
|
||||
return await probe_all(force=True)
|
||||
|
||||
|
||||
# 供 A2/A3 复用的类型别名
|
||||
UpstreamProbe = Callable[..., Awaitable[dict[str, Any]]]
|
||||
Reference in New Issue
Block a user