Files
PlantAI-rag-new/sdk/app/services/extraction.py
jerryandClaude Fable 5 a990a00eb5 端口终态修订:BFF 8090 / QA 8091 / Methodology 8092 / legacy 8093
- 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>
2026-08-31 16:41:37 +08:00

326 lines
11 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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()