- config.py 默认值改终态(浸泡期 .env 覆盖 7921/8001/7911 继续生效) - extraction _DEFAULT_ALT=8093;三个 start.sh 默认端口改终态 - 运行中服务零触碰:浸泡端口不变,BFF 未重启,旧服务不动 - 验证:sdk pytest 125 绿,BFF 五上游仍绿(env 覆盖优先) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
326 lines
11 KiB
Python
326 lines
11 KiB
Python
"""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:8093" # 终态(浸泡期由 RAG_ALT_BASE_URL env 覆盖 7911)
|
||
|
||
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()
|