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,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
|
||||
Reference in New Issue
Block a user