Compare commits
101
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3c00e1675e | ||
|
|
8ab4279b04 | ||
|
|
c943e8d363 | ||
|
|
9ac36e99ba | ||
|
|
9cb0bf2e35 | ||
|
|
6b4b5e571f | ||
|
|
9797b1b414 | ||
|
|
b9bc90ffb0 | ||
|
|
821f3c1c4e | ||
|
|
da8f531bbc | ||
|
|
21c79582bf | ||
|
|
1ecc4cb29e | ||
|
|
75eb489fc3 | ||
|
|
e4ed3589d5 | ||
|
|
f27336e961 | ||
|
|
963a87855c | ||
|
|
60375b04d7 | ||
|
|
1bd4aea5fd | ||
|
|
51532a30bc | ||
|
|
d37e0d9c3a | ||
|
|
c923419d2d | ||
|
|
43a7c6c132 | ||
|
|
8b228923bd | ||
|
|
2e5372cf90 | ||
|
|
91f6a1c242 | ||
|
|
0282cfeeeb | ||
|
|
0f6e16ce79 | ||
|
|
964346dc77 | ||
|
|
f10b86c088 | ||
|
|
d0e6c76ec1 | ||
|
|
16bde32020 | ||
|
|
ba43ba90bf | ||
|
|
36d3d8d62a | ||
|
|
f2c101ec65 | ||
|
|
750e218541 | ||
|
|
becf76ecd9 | ||
|
|
dcb810a758 | ||
|
|
3b512fffd0 | ||
|
|
04fb306624 | ||
|
|
2ee9af4bd0 | ||
|
|
957f3cb062 | ||
|
|
e3870d19de | ||
|
|
a008fbba2a | ||
|
|
3989dd3c3a | ||
|
|
8a58703d61 | ||
|
|
4266a89872 | ||
|
|
1bce4bf587 | ||
|
|
9bbe2c83f6 | ||
|
|
82b954b957 | ||
|
|
d12071a4d7 | ||
|
|
76529f78ab | ||
|
|
835d55cb5b | ||
|
|
653d2211cf | ||
|
|
4128fbf3bf | ||
|
|
68b45c60a0 | ||
|
|
c0b3026209 | ||
|
|
bb560d7bea | ||
|
|
2a6315e615 | ||
|
|
812decf310 | ||
|
|
3b16397baf | ||
|
|
4fa758d602 | ||
|
|
d95c2b5761 | ||
|
|
4866b0408a | ||
|
|
03e6ff508a | ||
|
|
79b4c7c52c | ||
|
|
f0a1711ebc | ||
|
|
688edba169 | ||
|
|
730692a1f6 | ||
|
|
a9a0ece4b4 | ||
|
|
00c77aa323 | ||
|
|
6bcb88cf76 | ||
|
|
037790c609 | ||
|
|
afa7cce938 | ||
|
|
7444edf525 | ||
|
|
8d115f8f1c | ||
|
|
18e69046f1 | ||
|
|
4b5b501463 | ||
|
|
7037a1d5d8 | ||
|
|
36fc2a7250 | ||
|
|
9d7ed1bdd3 | ||
|
|
f74b27daa1 | ||
|
|
00fe6a379e | ||
|
|
faa8686461 | ||
|
|
8e4fce062a | ||
|
|
04bf099faf | ||
|
|
d54419f040 | ||
|
|
0252d52425 | ||
|
|
2895cd8497 | ||
|
|
aec01ca0df | ||
|
|
1fe80b293e | ||
|
|
8d04eb6845 | ||
|
|
c620461794 | ||
|
|
025e2d516c | ||
|
|
f18286f022 | ||
|
|
28bee31d00 | ||
|
|
00eb0bc1d0 | ||
|
|
b5dc12e8a5 | ||
|
|
60f7360cec | ||
|
|
3889a2a29b | ||
|
|
75877ecab7 | ||
|
|
7840527f82 |
+26
@@ -21,3 +21,29 @@ bioflow_service/bioflow.log
|
||||
|
||||
# QA 数据库裸进程离线二进制载体(构建机从镜像/官网提取,不入 git)
|
||||
deploy/qa-db/bin/
|
||||
|
||||
# 水产库大文件目录(source 26G zip 解压产物、build 清单、registry 建库生成、data 运行数据)
|
||||
kb-libraries/aqua/source/
|
||||
kb-libraries/aqua/build/
|
||||
kb-libraries/aqua/registry/*
|
||||
!kb-libraries/aqua/registry/README.md
|
||||
kb-libraries/aqua/data/
|
||||
|
||||
# 18 阶段拓展分析测试用例数据(2.6G,不入库;测试时从本机同步)
|
||||
tests/extension_cases_18/
|
||||
|
||||
# 开发过程备份文件与本地运行库
|
||||
*.bak-*
|
||||
bioflow_service/bioflow.db
|
||||
*.before-*
|
||||
|
||||
# paper_writer_service 运行时产物(不入库)
|
||||
paper_writer_service/paper.db
|
||||
paper_writer_service/paper.db-shm
|
||||
paper_writer_service/paper.db-wal
|
||||
paper_writer_service/paper.log
|
||||
|
||||
# 部署方案文档(含主机凭据,不入库)
|
||||
docs/DEPLOY_PAPER_176.md
|
||||
# Python 覆盖率运行时产物(coverage.py --cov 生成)
|
||||
.coveragevstore/vstore.log
|
||||
|
||||
@@ -29,7 +29,7 @@ PAPERS_CSV = PROJECT_ROOT / "data" / "evaluation" / "methodology_papers.csv"
|
||||
GOLD_CSV = PROJECT_ROOT / "data" / "evaluation" / "methodology_gold.csv"
|
||||
REVIEW_RESULTS_JSON = BATCH_DIR / "review_results.json"
|
||||
REVIEW_RESULTS_CSV = BATCH_DIR / "review_results.csv"
|
||||
# 2026-08-31 新工程适配:上传目录/产出根支持 env 覆盖
|
||||
# 2026-08-31 新工程适配:上传目录/产出根支持 env 覆盖(数据根迁移时无需改代码)
|
||||
UPLOAD_JOBS_DIR = Path(os.environ.get("METHODOLOGY_UPLOAD_JOBS_DIR", PROJECT_ROOT / "data" / "evaluation" / "methodology_upload_jobs"))
|
||||
UPLOAD_OUTPUT_ROOT = Path(os.environ.get("METHODOLOGY_UPLOAD_OUTPUT_ROOT", PROJECT_ROOT / "data" / "evaluation" / "0505"))
|
||||
DEFAULT_MINERU_CMD = str(conda_env_path("minerupy312") / "bin" / "mineru")
|
||||
@@ -351,6 +351,22 @@ def mineru_gpu_preflight_error(
|
||||
)
|
||||
|
||||
|
||||
# 缺陷 #38(2026-09-11):成功任务的 job.log 必然有 doc_analyze_streaming 行(176 实测)。
|
||||
# returncode≠0 且日志中无此行 = 启动期失败(pipeline 起来前就死了,真因被块缓冲吞掉)。
|
||||
# 判定条件(methodology 与 legacy-rag 两处 build_mineru_runtime/run_uploaded_methodology_job 同改)
|
||||
MINERU_PIPELINE_LOG_MARKER = "doc_analyze_streaming"
|
||||
MINERU_STARTUP_RETRIES = 1
|
||||
|
||||
|
||||
def job_log_has_pipeline_line(log_path: Path) -> bool:
|
||||
"""job.log 是否已出现 pipeline 日志行(def _job_log_has_pipeline_line)。"""
|
||||
try:
|
||||
text = log_path.read_text(encoding="utf-8", errors="replace")
|
||||
except OSError:
|
||||
return False
|
||||
return MINERU_PIPELINE_LOG_MARKER in text
|
||||
|
||||
|
||||
def summarize_mineru_failure(log_path: Path, returncode: int) -> str:
|
||||
message = f"MinerU failed with exit code {returncode}; see {log_path}"
|
||||
if not log_path.exists():
|
||||
@@ -368,6 +384,28 @@ def summarize_mineru_failure(log_path: Path, returncode: int) -> str:
|
||||
return message
|
||||
|
||||
|
||||
def summarize_mineru_startup_failure(log_path: Path, returncode: int, retried: bool) -> str:
|
||||
"""启动期失败(无 pipeline 行)的中文可读错误(缺陷 #38)。
|
||||
|
||||
附上 job.log 中可见的原始错误尾部,让用户能看到真因(修复前真因
|
||||
被 stdout 块缓冲 + CLI 立即杀 API 进程永久丢失)。
|
||||
"""
|
||||
tail = ""
|
||||
try:
|
||||
text = log_path.read_text(encoding="utf-8", errors="replace")
|
||||
lines = [ln for ln in text.splitlines() if ln.strip()]
|
||||
if lines:
|
||||
tail = " | ".join(ln.strip() for ln in lines[-3:])[:500]
|
||||
except OSError:
|
||||
tail = ""
|
||||
detail = f"原始日志尾部:{tail}" if tail else f"原始错误见日志:{log_path}"
|
||||
action = "重试后" if retried else ""
|
||||
return (
|
||||
f"提取启动阶段失败{action},可能受主机 GPU/显示重配置干扰,请重试"
|
||||
f"(MinerU 退出码 {returncode};{detail})。"
|
||||
)
|
||||
|
||||
|
||||
def build_mineru_runtime() -> tuple[list[str], dict[str, str], dict[str, str]]:
|
||||
env = os.environ.copy()
|
||||
env.setdefault("MINERU_MODEL_SOURCE", "modelscope")
|
||||
@@ -384,6 +422,9 @@ def build_mineru_runtime() -> tuple[list[str], dict[str, str], dict[str, str]]:
|
||||
else:
|
||||
env.setdefault("CUDA_VISIBLE_DEVICES", "0")
|
||||
env.setdefault("TORCH_CUDNN_V8_API_DISABLED", "1")
|
||||
# 缺陷 #38(2026-09-11):子进程 API 的异常日志即写即落(不再被 stdout
|
||||
# 块缓冲吞掉)——CLI 失败后立即杀 API 进程前真因已落到 job.log。
|
||||
env["PYTHONUNBUFFERED"] = "1"
|
||||
# mineru 3.1.13 CLI 陷阱(部署 agent 实测):带 MINERU_API_SHUTDOWN_ON_STDIN_EOF=1
|
||||
# 时,stdin EOF 会让 API 进程静默 shutdown——subprocess 调用必须 stdin=DEVNULL
|
||||
env.setdefault("MINERU_API_SHUTDOWN_ON_STDIN_EOF", "0")
|
||||
@@ -566,7 +607,34 @@ def run_uploaded_methodology_job(
|
||||
stdout=handle,
|
||||
stderr=subprocess.STDOUT,
|
||||
)
|
||||
# 缺陷 #38:启动期失败自动重试 1 次——returncode≠0 且 job.log 无
|
||||
# doc_analyze_streaming 行 = pipeline 起来前就失败(真因曾因块缓冲
|
||||
# 丢失)。重试的 stdout 继续 append 到同一 job.log,日志/状态留
|
||||
# retry 标记(G6-2:已有 pipeline 行的中途失败不得触发重试)。
|
||||
startup_retried = False
|
||||
if result.returncode != 0 and not job_log_has_pipeline_line(log_path):
|
||||
startup_retried = True
|
||||
with log_path.open("a", encoding="utf-8") as handle:
|
||||
handle.write(
|
||||
f"[retry] MinerU 启动期失败(exit {result.returncode},"
|
||||
"无 pipeline 阶段日志),自动重试 1/1\n"
|
||||
)
|
||||
handle.write("$ " + " ".join(cmd) + "\n")
|
||||
handle.flush()
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
cwd=PROJECT_ROOT,
|
||||
env=mineru_env,
|
||||
text=True,
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=handle,
|
||||
stderr=subprocess.STDOUT,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
if not job_log_has_pipeline_line(log_path):
|
||||
raise RuntimeError(
|
||||
summarize_mineru_startup_failure(log_path, result.returncode, retried=startup_retried)
|
||||
)
|
||||
raise RuntimeError(summarize_mineru_failure(log_path, result.returncode))
|
||||
matches = sorted(mineru_out.rglob("*_content_list.json"))
|
||||
if not matches:
|
||||
|
||||
@@ -0,0 +1,328 @@
|
||||
"""test_mineru_retry.py —— MinerU 启动期失败重试与可读错误(缺陷 #38,验收 §4 + 回归 G6)。
|
||||
|
||||
运行(worktree 根目录,无需部署、不触服务):
|
||||
/home/xzq/5090/teamWorkspace/zh/plantai-rag-new-venv/bin/python -m pytest agent/methodology/api/test_mineru_retry.py -v
|
||||
/home/xzq/5090/teamWorkspace/zh/plantai-rag-new-venv/bin/python -m pytest agent/legacy-rag/api/test_mineru_retry.py -v
|
||||
说明:agent/methodology 与 agent/legacy-rag 共用同一份 routes(文件名同、内容同),
|
||||
同一测试文件复制两处、两处同跑(各用各树的 src/)。
|
||||
|
||||
断言:
|
||||
1. build_mineru_runtime 返回 env 含 PYTHONUNBUFFERED=1;
|
||||
2. subprocess.run 首次失败 + job.log 无 pipeline 行 → 自动重试 1 次
|
||||
(两次启动行、job.log 有 [retry] 标记,第二次 stdout 追加同一 job.log);
|
||||
3. 重试仍失败 → status failed + 中文可读「提取启动阶段失败」+ 原始错误尾部;
|
||||
4. 已有 pipeline 行的中途失败 → 不重试,走原错误路径(含 CUDA OOM 中文错误)。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
import api.methodology_routes as mr
|
||||
|
||||
from api.methodology_routes import (
|
||||
MINERU_PIPELINE_LOG_MARKER,
|
||||
build_mineru_runtime,
|
||||
job_log_has_pipeline_line,
|
||||
run_uploaded_methodology_job,
|
||||
summarize_mineru_startup_failure,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 固定夹具
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class FakeResult:
|
||||
returncode: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class RunnerRecorder:
|
||||
"""按调用顺序返回预设 returncode;每次调用把 stdout 写进所给 handle(模拟子进程日志)。"""
|
||||
|
||||
sequence: list
|
||||
calls: list = field(default_factory=list)
|
||||
|
||||
def run(self, cmd, **kwargs):
|
||||
n = len(self.calls)
|
||||
self.calls.append({"cmd": cmd, "kwargs": kwargs})
|
||||
rc = int(self.sequence[min(n, len(self.sequence) - 1)])
|
||||
handle = kwargs.get("stdout")
|
||||
if handle is not None and kwargs.get("text"):
|
||||
handle.write(f"[fake stdout call#{n}] rc={rc}\n")
|
||||
handle.flush()
|
||||
return FakeResult(returncode=rc)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_success(monkeypatch):
|
||||
"""LLM 批次与 bbav5 配置阶段全部假成功(不触真实脚本/目录/真实函数)。
|
||||
|
||||
- gf_python_path / _preflight_executable 桩掉(不依赖真实 python 路径)
|
||||
- build_bbav5_methodology_config 桩掉(不要求真实批次产物目录)
|
||||
- figure_table 阶段走 routes 内部 try/except,不阻断
|
||||
"""
|
||||
captured = {}
|
||||
|
||||
def fake_batch(cmd, **kwargs):
|
||||
captured["batch_cmd"] = cmd
|
||||
return FakeResult(returncode=0)
|
||||
|
||||
monkeypatch.setattr(mr, "gf_python_path", lambda: Path("/usr/bin/python3"))
|
||||
monkeypatch.setattr(mr, "_preflight_executable", lambda cmd, hint, name: None)
|
||||
monkeypatch.setattr(mr.subprocess, "run", fake_batch)
|
||||
monkeypatch.setattr(
|
||||
mr,
|
||||
"build_bbav5_methodology_config",
|
||||
lambda output_dir: captured.setdefault("config", {}),
|
||||
)
|
||||
return captured
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def job_root(monkeypatch, tmp_path):
|
||||
"""把 UPLOAD_JOBS_DIR 重定向到 tmp(绝不写真实数据目录)。"""
|
||||
root = tmp_path / "upload_jobs"
|
||||
monkeypatch.setattr(mr, "UPLOAD_JOBS_DIR", root)
|
||||
return root
|
||||
|
||||
|
||||
def read_status(job_dir: Path) -> dict:
|
||||
return json.loads((job_dir / "status.json").read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def run_job(monkeypatch, tmp_path, runner: RunnerRecorder, source_path: Path, fake_success, job_root):
|
||||
"""以 runner 重绑 subprocess.run 跑完整任务流(fake_success 保证后续阶段假成功)。
|
||||
|
||||
注意:fake_success 已把 mr.subprocess.run 桩为假批次;这里再次 setattr
|
||||
为 runner(同一模块对象),此后所有子进程调用(mineru 两次 + LLM 批次)
|
||||
都由 runner 按 sequence 依次应答——测试 sequence 需覆盖全部调用。
|
||||
"""
|
||||
job_dir = job_root / "testjob"
|
||||
# 模拟 mineru 成功产出:在 job/mineru 下伪造 *_content_list.json
|
||||
# (测试不真正执行 mineru)
|
||||
(job_dir / "mineru").mkdir(parents=True, exist_ok=True)
|
||||
(job_dir / "mineru" / "fake_content_list.json").write_text("{}", encoding="utf-8")
|
||||
# GPU 预检禁用(否则 nvidia-smi 会消耗 runner 的 sequence 槽位)
|
||||
monkeypatch.setenv("METHODOLOGY_MINERU_GPU_PREFLIGHT", "0")
|
||||
monkeypatch.setattr(subprocess, "run", runner.run)
|
||||
run_uploaded_methodology_job(
|
||||
job_id="testjob",
|
||||
source_path=source_path,
|
||||
paper_id="paper1",
|
||||
family="",
|
||||
output_dir=tmp_path / "out",
|
||||
llm_model="deepseek-v4-flash",
|
||||
llm_request_timeout=75,
|
||||
sandbox=False,
|
||||
sandbox_run_id="",
|
||||
)
|
||||
return job_dir
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pdf_source(tmp_path):
|
||||
src = tmp_path / "sample.pdf"
|
||||
src.write_bytes(b"%PDF-1.4\nfake pdf bytes for startup-failure test\n")
|
||||
return src
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def json_source(tmp_path):
|
||||
src = tmp_path / "sample.json"
|
||||
src.write_text("{}", encoding="utf-8")
|
||||
return src
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 修复点 1:env 注入 PYTHONUNBUFFERED
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_build_mineru_runtime_env_unbuffered(monkeypatch):
|
||||
monkeypatch.delenv("MINERU_CMD", raising=False)
|
||||
cmd, env, runtime = build_mineru_runtime()
|
||||
assert env.get("PYTHONUNBUFFERED") == "1"
|
||||
assert cmd # 默认 MINERU_CMD 路径(不要求存在,build 阶段不校验)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 修复点 2:启动期失败自动重试 1 次
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_startup_failure_retries_once(monkeypatch, tmp_path, pdf_source, fake_success, job_root):
|
||||
runner = RunnerRecorder(sequence=[1, 0, 0]) # mineru 首调失败 → 重试成功 → LLM 批次成功
|
||||
job_dir = run_job(monkeypatch, tmp_path, runner, pdf_source, fake_success, job_root)
|
||||
|
||||
log_text = (job_dir / "job.log").read_text(encoding="utf-8")
|
||||
startup_lines = [ln for ln in log_text.splitlines() if ln.startswith("$ ")]
|
||||
assert len(startup_lines) >= 2, f"应出现两次 mineru 启动(1 次自动重试),实际:{startup_lines}"
|
||||
assert "[retry]" in log_text, "job.log 应有可识别重试标记"
|
||||
# 第二次 subprocess.run 的 stdout 追加到同一 job.log
|
||||
assert "[fake stdout call#0]" in log_text
|
||||
assert "[fake stdout call#1]" in log_text
|
||||
|
||||
status = read_status(job_dir)
|
||||
assert status.get("status") == "completed"
|
||||
|
||||
|
||||
def test_startup_failure_retry_then_fail_readable_error(monkeypatch, tmp_path, pdf_source, fake_success, job_root):
|
||||
runner = RunnerRecorder(sequence=[1, 1]) # 两次都失败,且无 pipeline 行
|
||||
job_dir = run_job(monkeypatch, tmp_path, runner, pdf_source, fake_success, job_root)
|
||||
|
||||
log_text = (job_dir / "job.log").read_text(encoding="utf-8")
|
||||
assert len([ln for ln in log_text.splitlines() if ln.startswith("$ /")]) == 2 # 两次 mineru 启动
|
||||
assert "[retry]" in log_text
|
||||
# pipeline 判定:仅检查子进程输出行(启动命令行本身含 marker 关键字)
|
||||
stdout_lines = [ln for ln in log_text.splitlines() if "[fake stdout" in ln or ln.startswith("[retry]")]
|
||||
assert MINERU_PIPELINE_LOG_MARKER not in "\n".join(stdout_lines)
|
||||
|
||||
status = read_status(job_dir)
|
||||
assert status.get("status") == "failed"
|
||||
message = str(status.get("message") or "")
|
||||
assert "提取启动阶段失败" in message, f"中文可读错误缺失:{message}"
|
||||
assert "退出码" in message and "1" in message
|
||||
assert "原始日志尾部" in message or "原始错误见日志" in message
|
||||
assert message # 非空
|
||||
|
||||
|
||||
def test_startup_failure_with_marker_retries_once_then_no_second_retry(monkeypatch, tmp_path, pdf_source, job_root):
|
||||
"""G6-2 边界:第一次失败无 marker → 触发重试;重试后才有 marker 但 rc≠0 →
|
||||
不二次重试,走 summarize_mineru_failure(有 marker 的中途失败路径)。"""
|
||||
sequence = [1, 1]
|
||||
calls = []
|
||||
|
||||
def run_with_marker(cmd, **kwargs):
|
||||
n = len(calls)
|
||||
rc = int(sequence[min(n, len(sequence) - 1)])
|
||||
calls.append({"cmd": cmd, "kwargs": kwargs})
|
||||
handle = kwargs.get("stdout")
|
||||
if handle is not None and kwargs.get("text"):
|
||||
if n == 1: # 第二次调用(重试)时写入 pipeline 行
|
||||
handle.write(f"[fake] {MINERU_PIPELINE_LOG_MARKER} started\n")
|
||||
handle.write(f"[fake stdout call#{n}] rc={rc}\n")
|
||||
handle.flush()
|
||||
return FakeResult(returncode=rc)
|
||||
|
||||
monkeypatch.setenv("METHODOLOGY_MINERU_GPU_PREFLIGHT", "0")
|
||||
monkeypatch.setattr(subprocess, "run", run_with_marker)
|
||||
monkeypatch.setattr(mr, "_preflight_executable", lambda cmd, hint, name: None)
|
||||
monkeypatch.setattr(
|
||||
mr,
|
||||
"summarize_mineru_failure",
|
||||
lambda log_path, rc: f"CUDA out of memory fake path rc={rc}",
|
||||
)
|
||||
job_dir = job_root / "testjob"
|
||||
run_uploaded_methodology_job(
|
||||
job_id="testjob",
|
||||
source_path=pdf_source,
|
||||
paper_id="paper1",
|
||||
family="",
|
||||
output_dir=tmp_path / "out",
|
||||
llm_model="deepseek-v4-flash",
|
||||
llm_request_timeout=75,
|
||||
sandbox=False,
|
||||
sandbox_run_id="",
|
||||
)
|
||||
log_text = (job_dir / "job.log").read_text(encoding="utf-8")
|
||||
assert MINERU_PIPELINE_LOG_MARKER in log_text
|
||||
assert log_text.count("[retry]") == 1, "有 pipeline 行后不得再次重试"
|
||||
# 中途失败(有 marker)→ 原错误路径(CUDA 中文错误)
|
||||
status = read_status(job_dir)
|
||||
assert status.get("status") == "failed"
|
||||
assert "CUDA" in str(status.get("message") or "")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 修复点 4:有 pipeline 行的中途失败不重试
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_mid_failure_with_pipeline_line_no_retry(monkeypatch, tmp_path, pdf_source, job_root):
|
||||
sequence = [1]
|
||||
calls = []
|
||||
|
||||
def run_with_marker(cmd, **kwargs):
|
||||
n = len(calls)
|
||||
rc = int(sequence[min(n, len(sequence) - 1)])
|
||||
calls.append({"cmd": cmd, "kwargs": kwargs})
|
||||
handle = kwargs.get("stdout")
|
||||
if handle is not None and kwargs.get("text"):
|
||||
if n == 0: # 第一次调用即写 pipeline 行
|
||||
handle.write(f"[fake] {MINERU_PIPELINE_LOG_MARKER} started\n")
|
||||
handle.write(f"[fake stdout call#{n}] rc={rc}\n")
|
||||
handle.flush()
|
||||
return FakeResult(returncode=rc)
|
||||
|
||||
monkeypatch.setenv("METHODOLOGY_MINERU_GPU_PREFLIGHT", "0")
|
||||
monkeypatch.setattr(subprocess, "run", run_with_marker)
|
||||
monkeypatch.setattr(mr, "_preflight_executable", lambda cmd, hint, name: None)
|
||||
monkeypatch.setattr(
|
||||
mr,
|
||||
"summarize_mineru_failure",
|
||||
lambda log_path, rc: f"CUDA out of memory fake path rc={rc}",
|
||||
)
|
||||
job_dir = job_root / "testjob"
|
||||
run_uploaded_methodology_job(
|
||||
job_id="testjob",
|
||||
source_path=pdf_source,
|
||||
paper_id="paper1",
|
||||
family="",
|
||||
output_dir=tmp_path / "out",
|
||||
llm_model="deepseek-v4-flash",
|
||||
llm_request_timeout=75,
|
||||
sandbox=False,
|
||||
sandbox_run_id="",
|
||||
)
|
||||
log_text = (job_dir / "job.log").read_text(encoding="utf-8")
|
||||
assert MINERU_PIPELINE_LOG_MARKER in log_text
|
||||
assert "[retry]" not in log_text, "已有 pipeline 行的中途失败不得触发重试"
|
||||
# 原始启动块 = 1 条 `$ env ...` + 1 条 `$ <cmd>`;无重试追加
|
||||
assert len([ln for ln in log_text.splitlines() if ln.startswith("$ /")]) == 1
|
||||
status = read_status(job_dir)
|
||||
assert status.get("status") == "failed"
|
||||
assert "CUDA" in str(status.get("message") or "")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# JSON 直接投递(跳过 mineru 段)不重试、不破坏状态机
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_json_source_skips_mineru_no_retry(monkeypatch, tmp_path, json_source, fake_success, job_root):
|
||||
runner = RunnerRecorder(sequence=[0]) # 仅 LLM 批次一次
|
||||
job_dir = run_job(monkeypatch, tmp_path, runner, json_source, fake_success, job_root)
|
||||
log_text = (job_dir / "job.log").read_text(encoding="utf-8")
|
||||
assert "[retry]" not in log_text
|
||||
assert MINERU_PIPELINE_LOG_MARKER not in log_text
|
||||
status = read_status(job_dir)
|
||||
assert status.get("status") == "completed"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 工具函数单元
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_job_log_has_pipeline_line(tmp_path):
|
||||
log_path = tmp_path / "job.log"
|
||||
assert job_log_has_pipeline_line(log_path) is False
|
||||
log_path.write_text("noise\n", encoding="utf-8")
|
||||
assert job_log_has_pipeline_line(log_path) is False
|
||||
log_path.write_text(f"2026-09-11 [INFO] {MINERU_PIPELINE_LOG_MARKER} begin\n", encoding="utf-8")
|
||||
assert job_log_has_pipeline_line(log_path) is True
|
||||
|
||||
|
||||
def test_summarize_startup_failure_readable(tmp_path):
|
||||
log_path = tmp_path / "job.log"
|
||||
log_path.write_text("line one\nline two\nTraceback: boom\n", encoding="utf-8")
|
||||
msg = summarize_mineru_startup_failure(log_path, 1, retried=True)
|
||||
assert "提取启动阶段失败" in msg
|
||||
assert "重试后" in msg
|
||||
assert "退出码 1" in msg
|
||||
assert "Traceback: boom" in msg
|
||||
missing = summarize_mineru_startup_failure(tmp_path / "nope.log", 1, retried=False)
|
||||
assert "提取启动阶段失败" in missing
|
||||
assert "重试后" not in missing
|
||||
@@ -0,0 +1,226 @@
|
||||
#!/usr/bin/env python3
|
||||
"""V16.aqua.1.0 batch extraction with resume support.
|
||||
|
||||
Input manifest: JSON array of {paper_id, path, kind: "md"|"pdf"}.
|
||||
Each paper runs in its own subprocess (run_v16_aqua_one_paper.py); stdout/stderr
|
||||
are captured to <output_root>/<paper_id>/run.log. Progress is persisted to
|
||||
<output_root>/batch_progress.json after every completed future; rows whose
|
||||
status is "success" are skipped on startup, so re-running the same command
|
||||
resumes an interrupted batch. On completion, BATCH_RUN_SUMMARY.json is written.
|
||||
Failed rows keep their status and are retried on the next run.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1] # agent/legacy-rag(one_paper 的 sys.path 根)
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3] # 仓库根(PlantAI-rag-new):manifest 相对路径以此解析
|
||||
|
||||
|
||||
def _row_id(row: dict) -> str:
|
||||
return str(row.get("paper_id") or Path(str(row.get("path", ""))).stem)
|
||||
|
||||
|
||||
def _safe_dir(paper_id: str) -> str:
|
||||
"""Filesystem-safe directory name for a paper_id (DOIs contain '/')."""
|
||||
return re.sub(r"[^A-Za-z0-9._-]+", "_", str(paper_id))
|
||||
|
||||
|
||||
def build_run_command(
|
||||
row: dict,
|
||||
root: Path,
|
||||
output_dir: Path,
|
||||
python: Path,
|
||||
*,
|
||||
use_pc_data: bool = False,
|
||||
disable_llm: bool = False,
|
||||
) -> list[str]:
|
||||
kind = str(row.get("kind") or "md").lower()
|
||||
source = Path(str(row["path"]))
|
||||
if not source.is_absolute():
|
||||
source = REPO_ROOT / source # manifest 路径为仓库根相对路径
|
||||
if kind == "pdf":
|
||||
flag, value = "--pdf", source
|
||||
else:
|
||||
flag, value = "--md", source
|
||||
command = [
|
||||
str(python),
|
||||
str(root / "scripts/run_v16_aqua_one_paper.py"),
|
||||
flag,
|
||||
str(value),
|
||||
"--output-dir",
|
||||
str(output_dir),
|
||||
]
|
||||
pmid = str(row.get("pmid") or "").strip()
|
||||
doi = str(row.get("doi") or "").strip()
|
||||
paper_id = str(row.get("paper_id") or "").strip()
|
||||
if pmid:
|
||||
command.extend(["--pmid", pmid])
|
||||
if doi:
|
||||
command.extend(["--doi", doi])
|
||||
if paper_id:
|
||||
command.extend(["--paper-id", paper_id])
|
||||
if use_pc_data:
|
||||
command.append("--use-pc-data")
|
||||
if disable_llm:
|
||||
command.append("--disable-llm")
|
||||
return command
|
||||
|
||||
|
||||
def run_one(
|
||||
row: dict,
|
||||
root: Path,
|
||||
output_root: Path,
|
||||
python: Path,
|
||||
*,
|
||||
use_pc_data: bool = False,
|
||||
disable_llm: bool = False,
|
||||
) -> dict:
|
||||
paper_id = _row_id(row)
|
||||
out = output_root / _safe_dir(paper_id)
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
command = build_run_command(
|
||||
row, root, out, python, use_pc_data=use_pc_data, disable_llm=disable_llm
|
||||
)
|
||||
log_path = out / "run.log"
|
||||
started = time.time()
|
||||
with log_path.open("w", encoding="utf-8") as log:
|
||||
process = subprocess.run(
|
||||
command,
|
||||
cwd=root,
|
||||
env=os.environ.copy(),
|
||||
stdout=log,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
)
|
||||
ended = time.time()
|
||||
result_path = out / "result.json"
|
||||
result = {
|
||||
**row,
|
||||
"paper_id": paper_id,
|
||||
"output_dir": str(out),
|
||||
"result_json": str(result_path),
|
||||
"return_code": process.returncode,
|
||||
"wall_seconds": round(ended - started, 3),
|
||||
"success": process.returncode == 0 and result_path.exists(),
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--manifest", required=True, type=Path)
|
||||
parser.add_argument("--output-root", required=True, type=Path)
|
||||
parser.add_argument("--workers", type=int, default=4)
|
||||
parser.add_argument("--use-pc-data", action="store_true")
|
||||
parser.add_argument("--disable-llm", action="store_true")
|
||||
parser.add_argument(
|
||||
"--project-root",
|
||||
type=Path,
|
||||
default=Path(os.environ.get("PLANTAI_PROJECT_ROOT", PROJECT_ROOT)),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--python",
|
||||
type=Path,
|
||||
default=Path(os.environ.get("PLANTAI_PYTHON", sys.executable)),
|
||||
)
|
||||
args = parser.parse_args()
|
||||
if args.workers < 1:
|
||||
raise SystemExit("--workers must be >= 1")
|
||||
rows = json.loads(args.manifest.read_text(encoding="utf-8"))
|
||||
# Resolve to absolute at the caller's cwd, then to the caller's cwd so the
|
||||
# subprocess (cwd=project_root) writes into the same tree the progress file
|
||||
# lives in. Relative output roots would otherwise land under project_root.
|
||||
args.output_root = args.output_root.resolve()
|
||||
args.manifest = args.manifest.resolve()
|
||||
args.output_root.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Resume: load previous progress, skip success rows.
|
||||
progress_path = args.output_root / "batch_progress.json"
|
||||
prior = {}
|
||||
if progress_path.exists():
|
||||
try:
|
||||
prior = {_row_id(item): item for item in json.loads(progress_path.read_text(encoding="utf-8"))}
|
||||
except Exception:
|
||||
prior = {}
|
||||
todo = [row for row in rows if prior.get(_row_id(row), {}).get("status") != "success"]
|
||||
done_results = [prior[_row_id(row)] for row in rows if prior.get(_row_id(row), {}).get("status") == "success"]
|
||||
print(
|
||||
f"manifest={len(rows)} already_success={len(done_results)} to_run={len(todo)}",
|
||||
flush=True,
|
||||
)
|
||||
if not todo:
|
||||
print("nothing to run; writing summary and exiting", flush=True)
|
||||
else:
|
||||
results: list[dict] = []
|
||||
started = time.time()
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as pool:
|
||||
future_map = {
|
||||
pool.submit(
|
||||
run_one,
|
||||
row,
|
||||
args.project_root,
|
||||
args.output_root,
|
||||
args.python,
|
||||
use_pc_data=args.use_pc_data,
|
||||
disable_llm=args.disable_llm,
|
||||
): row
|
||||
for row in todo
|
||||
}
|
||||
for future in concurrent.futures.as_completed(future_map):
|
||||
row = future_map[future]
|
||||
try:
|
||||
result = future.result()
|
||||
result["status"] = "success" if result["success"] else "failed"
|
||||
except Exception as exc:
|
||||
result = {
|
||||
**row,
|
||||
"paper_id": _row_id(row),
|
||||
"return_code": -1,
|
||||
"wall_seconds": 0,
|
||||
"success": False,
|
||||
"status": "failed",
|
||||
"error": repr(exc),
|
||||
}
|
||||
results.append(result)
|
||||
all_results = done_results + results
|
||||
(progress_path).write_text(
|
||||
json.dumps(all_results, ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
print(
|
||||
f"[{len(all_results)}/{len(rows)}] {result['paper_id']} "
|
||||
f"status={result['status']} wall={result['wall_seconds']}s",
|
||||
flush=True,
|
||||
)
|
||||
done_results.extend(results)
|
||||
|
||||
summary = {
|
||||
"started_epoch": time.time(),
|
||||
"workers": args.workers,
|
||||
"total": len(done_results),
|
||||
"success": sum(bool(item.get("status") == "success") for item in done_results),
|
||||
"failed": sum(bool(item.get("status") != "success") for item in done_results),
|
||||
"results": done_results,
|
||||
}
|
||||
(args.output_root / "BATCH_RUN_SUMMARY.json").write_text(
|
||||
json.dumps(summary, ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
print(
|
||||
f"BATCH DONE success={summary['success']} failed={summary['failed']}",
|
||||
flush=True,
|
||||
)
|
||||
return 0 if summary["failed"] == 0 else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,217 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run the V16.aqua.1.0 aquaculture-science SRT extraction pipeline for one paper.
|
||||
|
||||
Entry points:
|
||||
- --md <path> : read the Markdown directly, skip MinerU. Text goes through the
|
||||
block builder's markdown path (_from_markdown).
|
||||
- --pdf <path> : full MinerU path (same shape as run_v16_1_2_one_paper.py).
|
||||
|
||||
Cost instrumentation: every chat.completions.create response in the LLM lane is
|
||||
appended to <output-dir>/cost_log.jsonl as
|
||||
{paper_id, model, prompt_tokens, completion_tokens, timestamp}.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
if str(PROJECT_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(PROJECT_ROOT))
|
||||
|
||||
from src.rag.extractors.v16_aqua_1_0.v16_extractor import V16AquaExtractor
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Run the formal V16.aqua.1.0 aquaculture-science SRT extraction pipeline."
|
||||
)
|
||||
group = parser.add_mutually_exclusive_group(required=True)
|
||||
group.add_argument("--md", type=Path, help="Input Markdown file (MinerU skipped).")
|
||||
group.add_argument("--pdf", type=Path, help="Input PDF file (MinerU path).")
|
||||
parser.add_argument("--output-dir", required=True, type=Path)
|
||||
parser.add_argument("--pmid", default="")
|
||||
parser.add_argument("--doi", default="")
|
||||
# Stable batch paper_id (e.g. doi:{doi_norm}): used for chunk/block IDs and
|
||||
# the cost log. Falls back to the source file stem when omitted.
|
||||
parser.add_argument("--paper-id", default="")
|
||||
parser.add_argument("--disable-llm", action="store_true")
|
||||
parser.add_argument("--use-pc-data", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
output_dir = args.output_dir.resolve()
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if args.md:
|
||||
md_path = args.md.resolve()
|
||||
if not md_path.exists():
|
||||
raise SystemExit(f"Markdown does not exist: {md_path}")
|
||||
return _run_md(
|
||||
md_path=md_path,
|
||||
output_dir=output_dir,
|
||||
pmid=args.pmid or None,
|
||||
doi=args.doi or None,
|
||||
paper_id=args.paper_id or None,
|
||||
disable_llm=args.disable_llm,
|
||||
use_pc_data=args.use_pc_data,
|
||||
)
|
||||
pdf = args.pdf.resolve()
|
||||
if not pdf.exists():
|
||||
raise SystemExit(f"PDF does not exist: {pdf}")
|
||||
return _run_pdf(
|
||||
pdf=pdf,
|
||||
output_dir=output_dir,
|
||||
pmid=args.pmid or None,
|
||||
doi=args.doi or None,
|
||||
paper_id=args.paper_id or None,
|
||||
disable_llm=args.disable_llm,
|
||||
use_pc_data=args.use_pc_data,
|
||||
)
|
||||
|
||||
|
||||
def _new_extractor(output_dir: Path, disable_llm: bool, use_pc_data: bool):
|
||||
extractor = V16AquaExtractor(
|
||||
output_dir=output_dir,
|
||||
use_mineru=True,
|
||||
use_llm_relations=not disable_llm,
|
||||
use_rule_relation_fallback=False,
|
||||
use_vlm=False,
|
||||
use_neo4j=False,
|
||||
use_unified_storage=False,
|
||||
use_llm_definitions=False,
|
||||
use_pc_data=use_pc_data,
|
||||
use_disambiguation=use_pc_data,
|
||||
)
|
||||
# Cost instrumentation: JSONL append at the LLM lane (enabled when LLM on).
|
||||
if not disable_llm and extractor.llm_relation_extractor is not None:
|
||||
llm = extractor.llm_relation_extractor
|
||||
if hasattr(llm, "set_cost_log"):
|
||||
llm.set_cost_log(output_dir / "cost_log.jsonl")
|
||||
return extractor
|
||||
|
||||
|
||||
def _run_pdf(pdf: Path, output_dir: Path, pmid, doi, paper_id, disable_llm, use_pc_data) -> int:
|
||||
extractor = _new_extractor(output_dir, disable_llm, use_pc_data)
|
||||
result = extractor.extract(pdf, pmid=paper_id or pmid, doi=doi)
|
||||
_print_summary(result, pdf, output_dir, kind="pdf")
|
||||
return 0
|
||||
|
||||
|
||||
def _run_md(md_path: Path, output_dir: Path, pmid, doi, paper_id, disable_llm, use_pc_data) -> int:
|
||||
# Markdown entry: skip MinerU and feed the builder's markdown path directly,
|
||||
# mirroring phase1_5 _load_markdown/_from_markdown. We stage the .md as
|
||||
# <stem>.md inside a temp "mineru" dir so the block builder can find it, and
|
||||
# pass a placeholder pdf so downstream phases see the same stem.
|
||||
text = md_path.read_text(encoding="utf-8", errors="ignore")
|
||||
extractor = _new_extractor(output_dir, disable_llm, use_pc_data)
|
||||
builder = extractor.v16_block_builder
|
||||
paper_id = paper_id or pmid or md_path.stem
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="v16aqua_md_") as tmp:
|
||||
tmp_path = Path(tmp)
|
||||
(tmp_path / f"{md_path.stem}.md").write_text(text, encoding="utf-8")
|
||||
blocks = builder.build(tmp_path, md_path.stem, fallback_text=text)
|
||||
reconstructed = builder.reconstruct_text(blocks)
|
||||
if not reconstructed.strip():
|
||||
reconstructed = text
|
||||
|
||||
# Sections from block spans (same as _phase1_document_parser).
|
||||
# MD entry has no page numbers: position/end use block order offsets.
|
||||
from src.rag.extractors.v16_aqua_1_0.phase1_document_parser import Section
|
||||
from src.rag.extractors.v16_aqua_1_0.sectionizer_v16_1_1 import build_section_spans
|
||||
|
||||
spans = build_section_spans(blocks)
|
||||
sections = [
|
||||
Section(
|
||||
title=span.title,
|
||||
level=1,
|
||||
parent=None,
|
||||
start_page=0,
|
||||
end_page=0,
|
||||
section_type=span.section_type,
|
||||
position=span.start_order,
|
||||
end=span.end_order,
|
||||
)
|
||||
for span in spans
|
||||
if span.section_type != "references"
|
||||
]
|
||||
print(f" markdown blocks: {len(blocks)}")
|
||||
print(f" sections: {len(sections)}")
|
||||
print(f" reconstructed text: {len(reconstructed)} chars")
|
||||
|
||||
# Enter the normal pipeline after phase 1, keeping the same phase order as
|
||||
# V16Extractor.extract (supplementary/multiomics/gates are disabled in
|
||||
# practice: no local supp files, no PMID, no dataset context).
|
||||
extractor._v16_blocks = blocks
|
||||
extractor._sections = sections
|
||||
extractor._last_pdf_path = md_path
|
||||
extractor._last_mineru_dir = None
|
||||
extractor._last_md_path = md_path
|
||||
|
||||
supp_relations = extractor._phase2_supplementary(md_path, reconstructed, paper_id, doi=doi)
|
||||
supp_relations = extractor._phase4_5_triple_canonicalization(
|
||||
supp_relations, reconstructed, source_label="supplementary"
|
||||
)
|
||||
supp_relations = extractor._phase4_6_relation_semantic_review(
|
||||
supp_relations, reconstructed, source_label="supplementary"
|
||||
)
|
||||
multiomics_relations = extractor._phase2_5_multiomics(reconstructed, pmid=paper_id, doi=doi, title=None)
|
||||
multiomics_relations = extractor._phase4_5_triple_canonicalization(
|
||||
multiomics_relations, reconstructed, source_label="multiomics"
|
||||
)
|
||||
multiomics_relations = extractor._phase4_6_relation_semantic_review(
|
||||
multiomics_relations, reconstructed, source_label="multiomics"
|
||||
)
|
||||
entities = extractor._phase3_entity_recognition(reconstructed)
|
||||
text_relations = extractor._phase4_relation_extraction(reconstructed, sections, entities, paper_id)
|
||||
text_relations = extractor._phase4_5_triple_canonicalization(
|
||||
text_relations, reconstructed, source_label="text"
|
||||
)
|
||||
text_relations = extractor._split_enumeration_relations(text_relations, reconstructed)
|
||||
text_relations = extractor._phase4_6_relation_semantic_review(
|
||||
text_relations, reconstructed, source_label="text"
|
||||
)
|
||||
validated_relations, chunks = extractor._phase5_post_processing(
|
||||
reconstructed, entities, text_relations
|
||||
)
|
||||
all_relations = validated_relations + supp_relations + multiomics_relations
|
||||
all_relations = extractor._deduplicate_relations(all_relations)
|
||||
if all_relations:
|
||||
entities = extractor._backfill_entities_from_relations(entities, all_relations)
|
||||
|
||||
extractor.result.entities = entities
|
||||
extractor.result.relations = all_relations
|
||||
extractor.result.sections = sections
|
||||
extractor.result.chunks = chunks
|
||||
extractor.result.metadata["assessment"] = {"note": "markdown entry; phase7 review skipped"}
|
||||
|
||||
output = extractor._build_output()
|
||||
result_path = output_dir / "result.json"
|
||||
result_path.write_text(json.dumps(output, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(f" result.json saved: {result_path}")
|
||||
_print_summary(output, md_path, output_dir, kind="md")
|
||||
return 0
|
||||
|
||||
|
||||
def _print_summary(result, source: Path, output_dir: Path, kind: str) -> None:
|
||||
summary = {
|
||||
"version": result.get("metadata", {}).get("version"),
|
||||
"kind": kind,
|
||||
"source": str(source),
|
||||
"output_dir": str(output_dir),
|
||||
"entities": len(result.get("entities", [])),
|
||||
"relations": len(result.get("relations", [])),
|
||||
"relation_units": len(result.get("chunks", [])),
|
||||
"result_json": str(output_dir / "result.json"),
|
||||
}
|
||||
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -33,6 +33,10 @@ set -a
|
||||
. "$ENV_FILE"
|
||||
set +a
|
||||
|
||||
if [[ -f "$AGENT_ROOT/env/active_library.env" ]]; then
|
||||
set -a; . "$AGENT_ROOT/env/active_library.env"; set +a
|
||||
fi
|
||||
|
||||
cd "$AGENT_ROOT"
|
||||
setsid "$PYTHON" -m uvicorn "$APP_MOD" \
|
||||
--host "$BIND" --port "$PORT" \
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
"""V16.aqua.1.0 aquaculture-science SRT extraction package (fork of V16.1.2)."""
|
||||
|
||||
__version__ = "16.aqua.1.0"
|
||||
|
||||
from .v16_extractor import V16AquaExtractor, V16Extractor
|
||||
@@ -0,0 +1,319 @@
|
||||
"""
|
||||
ChunkBuilder — unified section-to-chunk layer (replaces RelationBlockBuilder + window filtering).
|
||||
|
||||
Sections from Phase 1 are cleaned, split into complete sentences, then packed into
|
||||
fixed-size chunks with overlap (= last sentence of previous chunk). Each chunk is
|
||||
a self-contained LLM input with endpoint candidates pre-extracted.
|
||||
|
||||
Flow:
|
||||
Section → trim_refs → filter_noise → split_sentences → pack_chunks(overlap) → Chunk[]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
|
||||
|
||||
|
||||
@dataclass
|
||||
class Chunk:
|
||||
chunk_id: str
|
||||
section_title: str
|
||||
source_label: str
|
||||
text: str
|
||||
char_count: int
|
||||
endpoint_candidates: Dict[str, List[str]] = field(default_factory=dict)
|
||||
|
||||
|
||||
class ChunkBuilder:
|
||||
"""Build LLM input chunks from paper sections."""
|
||||
|
||||
# Patterns for noise/ref detection (copied from RelationBlockBuilder)
|
||||
REFERENCE_HEADING = re.compile(
|
||||
r"^\s*(?:#+\s*)?(?:references?|bibliography|literature cited)\s*$", re.I
|
||||
)
|
||||
URL_ONLY = re.compile(r"^\s*(?:\d{1,3}\s*)?(?:https?://|www\.)\S+\s*$", re.I)
|
||||
SUPPLEMENT_NOTICE = re.compile(
|
||||
r"\b(supplementary material|additional file|available at|data availability|supporting information)\b.*(?:https?://|doi:)",
|
||||
re.I,
|
||||
)
|
||||
|
||||
ABSTRACT_HEADING = re.compile(r"^\s*(abstract|summary)\b", re.I)
|
||||
HIGH_RECALL_SECTION = re.compile(
|
||||
r"\b(abstract|summary|result|results|discussion|conclusion|background|introduction|supplement(?:ary)?|table|figure|caption)\b",
|
||||
re.I,
|
||||
)
|
||||
EXCLUDED_SECTION = re.compile(
|
||||
r"\b(references?|acknowledg(?:e)?ments?|author contributions?|funding|conflicts? of interest)\b",
|
||||
re.I,
|
||||
)
|
||||
|
||||
# Sentence split placeholder patterns
|
||||
SENTENCE_PLACEHOLDERS = {
|
||||
"et al.": "et al<prd>",
|
||||
"Fig.": "Fig<prd>",
|
||||
"Figs.": "Figs<prd>",
|
||||
"Dr.": "Dr<prd>",
|
||||
"Prof.": "Prof<prd>",
|
||||
"e.g.": "e<prd>g<prd>",
|
||||
"i.e.": "i<prd>e<prd>",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_chunk_chars: int = 5200,
|
||||
):
|
||||
self.max_chunk_chars = max_chunk_chars
|
||||
|
||||
def build(
|
||||
self,
|
||||
full_text: str,
|
||||
sections: Optional[Iterable] = None,
|
||||
section_text_getter: Optional[Callable[[str, Any], str]] = None,
|
||||
pmid: str = "",
|
||||
) -> List[Chunk]:
|
||||
"""Build chunks from paper sections."""
|
||||
chunks: List[Chunk] = []
|
||||
|
||||
if not sections:
|
||||
# Fallback: try regex heading splitting, then full_text
|
||||
chunks.extend(self._build_from_headings(full_text, pmid))
|
||||
return chunks
|
||||
|
||||
for section in sections:
|
||||
title = str(getattr(section, "title", "") or "section")
|
||||
if self._is_excluded(title):
|
||||
continue
|
||||
text = ""
|
||||
if section_text_getter:
|
||||
try:
|
||||
text = section_text_getter(full_text, section)
|
||||
except Exception:
|
||||
text = ""
|
||||
if not text:
|
||||
start = int(getattr(section, "position", 0) or 0)
|
||||
end = int(getattr(section, "end", len(full_text)) or len(full_text))
|
||||
text = (full_text or "")[start:end]
|
||||
chunks.extend(self._build_from_text(title, text, pmid))
|
||||
|
||||
return chunks
|
||||
|
||||
# Regex for inline headings (fallback when Phase 1 sections are empty)
|
||||
INLINE_MAJOR_HEADING = re.compile(
|
||||
r"(?:^|\n)\s*(?:\d+(?:\.\d+)*\.?\s*)?"
|
||||
r"(abstract|summary|introduction|background|results?|discussion|"
|
||||
r"conclusions?|materials and methods|methods?|supplementary(?: material| information)?)"
|
||||
r"\b\s*[::]?\s*(?=\n|$)",
|
||||
re.I,
|
||||
)
|
||||
|
||||
def _build_from_headings(self, full_text: str, pmid: str) -> List[Chunk]:
|
||||
"""Fallback: split full_text by inline headings when Phase 1 sections are empty."""
|
||||
text = str(full_text or "")
|
||||
if not text.strip():
|
||||
return []
|
||||
|
||||
matches = list(self.INLINE_MAJOR_HEADING.finditer(text))
|
||||
if not matches:
|
||||
return self._build_from_text("full_text", text, pmid)
|
||||
|
||||
chunks: List[Chunk] = []
|
||||
|
||||
# Front matter before first heading
|
||||
first = matches[0]
|
||||
front = text[: first.start()].strip()
|
||||
if front:
|
||||
chunks.extend(self._build_from_text("introduction", front, pmid))
|
||||
|
||||
for idx, match in enumerate(matches):
|
||||
title = match.group(1).lower().strip()
|
||||
# Normalize: "results" not "results"
|
||||
if title in ("method",):
|
||||
title = "methods"
|
||||
if title == "materials and methods":
|
||||
title = "Methods"
|
||||
body_start = match.end()
|
||||
body_end = matches[idx + 1].start() if idx + 1 < len(matches) else len(text)
|
||||
body = text[body_start:body_end].strip()
|
||||
if body:
|
||||
chunks.extend(self._build_from_text(title, body, pmid))
|
||||
|
||||
return chunks
|
||||
|
||||
def _build_from_text(self, section_title: str, text: str, pmid: str) -> List[Chunk]:
|
||||
"""Build chunks from a single section's text."""
|
||||
# Step 1: Trim references tail
|
||||
text = self._trim_reference_tail(str(text or ""))
|
||||
|
||||
# Step 2: Filter noise lines
|
||||
clean_lines = self._filter_lines(text)
|
||||
clean_text = "\n".join(clean_lines)
|
||||
|
||||
# Step 3: Split into sentences
|
||||
sentences = self._split_sentences(clean_text)
|
||||
if not sentences:
|
||||
return []
|
||||
|
||||
# Step 4: Pack into chunks with overlap
|
||||
source_label = self._source_label(section_title)
|
||||
raw_chunks = self._pack_chunks(sentences)
|
||||
|
||||
# Step 5: Build Chunk objects
|
||||
result: List[Chunk] = []
|
||||
for idx, chunk_text in enumerate(raw_chunks):
|
||||
if not chunk_text.strip():
|
||||
continue
|
||||
section_slug = re.sub(r"[^A-Za-z0-9]+", "_", section_title.lower()).strip("_") or "section"
|
||||
chunk_id = f"{pmid or 'paper'}_{section_slug}_{idx + 1:04d}"
|
||||
result.append(Chunk(
|
||||
chunk_id=chunk_id,
|
||||
section_title=section_title,
|
||||
source_label=source_label,
|
||||
text=chunk_text,
|
||||
char_count=len(chunk_text),
|
||||
))
|
||||
|
||||
return result
|
||||
|
||||
def _trim_reference_tail(self, text: str) -> str:
|
||||
"""Drop everything after the references heading."""
|
||||
value = str(text or "")
|
||||
if not value:
|
||||
return value
|
||||
patterns = [
|
||||
r"(?im)^\s*(?:#+\s*)?(?:references|literature cited|bibliography)\s*$",
|
||||
r"(?im)(?:^|\n)\s*(?:#+\s*)?(?:references|literature cited|bibliography)\s*(?=\n|$)",
|
||||
]
|
||||
best = None
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, value)
|
||||
if match:
|
||||
best = match.start() if best is None else min(best, match.start())
|
||||
if best is None:
|
||||
return value
|
||||
return value[:best]
|
||||
|
||||
CAPTION_START = re.compile(
|
||||
r"^\s*(?:fig(?:ure)?\.?\s*[\d]+|table\s*[\d]+|supplementary\s+(?:fig(?:ure)?|table))\b",
|
||||
re.I,
|
||||
)
|
||||
|
||||
def _filter_lines(self, text: str) -> List[str]:
|
||||
"""Filter out noise lines (URLs, OCR junk, markdown images, captions)."""
|
||||
clean: List[str] = []
|
||||
for line in str(text or "").splitlines():
|
||||
stripped = line.strip()
|
||||
|
||||
# Drop markdown image lines: 
|
||||
if re.match(r"^\s*!\\[.*\\]\\(.*\\)\s*$", stripped):
|
||||
continue
|
||||
|
||||
# Drop figure/table caption lines
|
||||
if self.CAPTION_START.match(stripped):
|
||||
continue
|
||||
|
||||
if self.URL_ONLY.match(stripped):
|
||||
continue
|
||||
if self.SUPPLEMENT_NOTICE.search(stripped) and len(stripped) < 260:
|
||||
continue
|
||||
if self._is_visual_noise(stripped):
|
||||
continue
|
||||
clean.append(line)
|
||||
return clean
|
||||
|
||||
def _is_visual_noise(self, line: str) -> bool:
|
||||
if len(line) < 4:
|
||||
return True
|
||||
tokens = re.findall(r"[A-Za-z0-9_.-]+", line)
|
||||
if not tokens:
|
||||
return True
|
||||
short_tokens = sum(1 for t in tokens if len(t) <= 2)
|
||||
alpha = len(re.findall(r"[A-Za-z]", line))
|
||||
digits = len(re.findall(r"\d", line))
|
||||
if len(tokens) >= 8 and short_tokens / max(1, len(tokens)) > 0.65:
|
||||
return True
|
||||
if re.fullmatch(r"(?:[A-Z]\s*){3,}", line):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _split_sentences(self, text: str) -> List[str]:
|
||||
"""Split text into complete sentences."""
|
||||
if not text:
|
||||
return []
|
||||
cleaned = re.sub(r"\s+", " ", text).strip()
|
||||
if not cleaned:
|
||||
return []
|
||||
# Protect abbreviations
|
||||
for key, value in self.SENTENCE_PLACEHOLDERS.items():
|
||||
cleaned = cleaned.replace(key, value)
|
||||
parts = re.split(r"(?<=[.!?])\s+(?=[A-Z0-9(])", cleaned)
|
||||
sentences = []
|
||||
for part in parts:
|
||||
for key, value in self.SENTENCE_PLACEHOLDERS.items():
|
||||
part = part.replace(value, key)
|
||||
part = part.strip()
|
||||
if part:
|
||||
sentences.append(part)
|
||||
return sentences
|
||||
|
||||
def _pack_chunks(self, sentences: List[str]) -> List[str]:
|
||||
"""Pack sentences into fixed-size chunks with overlap.
|
||||
|
||||
Each chunk is ≤ max_chunk_chars.
|
||||
Overlap = last complete sentence of previous chunk.
|
||||
Chunks always end on a complete sentence.
|
||||
"""
|
||||
chunks: List[str] = []
|
||||
current: List[str] = []
|
||||
current_len = 0
|
||||
|
||||
for sentence in sentences:
|
||||
sentence_len = len(sentence)
|
||||
|
||||
# If adding this sentence would exceed limit, flush current chunk
|
||||
if current and current_len + sentence_len + 1 > self.max_chunk_chars:
|
||||
chunks.append(" ".join(current))
|
||||
|
||||
# Start new chunk with overlap = last sentence
|
||||
overlap = current[-1]
|
||||
current = [overlap]
|
||||
current_len = len(overlap) + 1
|
||||
|
||||
current.append(sentence)
|
||||
current_len += sentence_len + 1
|
||||
|
||||
# Don't create a tiny final chunk if it's just the overlap
|
||||
if current:
|
||||
remaining = " ".join(current)
|
||||
# If this is just the overlap sentence carried over, skip if there
|
||||
# was already a previous chunk (the overlap is already included)
|
||||
if len(current) == 1 and len(chunks) >= 1:
|
||||
# Check if this sentence is identical to the last chunk's last sentence
|
||||
last_chunk = chunks[-1]
|
||||
last_sentence_of_last = last_chunk.rfind(current[0])
|
||||
if last_sentence_of_last != -1 and last_sentence_of_last > len(last_chunk) - len(current[0]) - 5:
|
||||
# It's already at the end of the previous chunk — skip
|
||||
pass
|
||||
else:
|
||||
chunks.append(remaining)
|
||||
else:
|
||||
chunks.append(remaining)
|
||||
|
||||
return chunks
|
||||
|
||||
def _source_label(self, section_title: str) -> str:
|
||||
clean = str(section_title or "").lower()
|
||||
if self.ABSTRACT_HEADING.match(clean):
|
||||
return "abstract"
|
||||
if re.search(r"\b(supplement|additional file|supporting information|appendix)\b", clean):
|
||||
return "supplementary"
|
||||
if re.search(r"\b(table|figure|caption)\b", clean):
|
||||
return "caption"
|
||||
return "text"
|
||||
|
||||
def _is_excluded(self, title: str) -> bool:
|
||||
clean = re.sub(r"\s+", " ", str(title or "")).strip()
|
||||
return bool(self.REFERENCE_HEADING.match(clean) or self.EXCLUDED_SECTION.search(clean))
|
||||
@@ -0,0 +1,184 @@
|
||||
"""ChunkBuilder v15.2.1 - subheading-aware chunking with section standardization + robust abstract detection."""
|
||||
from __future__ import annotations
|
||||
import re
|
||||
from typing import List, Tuple
|
||||
from .chunk_builder import ChunkBuilder, Chunk
|
||||
from .phase1_document_parser_mineru import Section
|
||||
|
||||
|
||||
class ChunkBuilderV15_2_1(ChunkBuilder):
|
||||
"""Extends ChunkBuilder with subheading-aware section splitting."""
|
||||
|
||||
SUBHEADING = re.compile(r"^#+\s+(.+)$", re.MULTILINE)
|
||||
|
||||
SECTION_MAP = {
|
||||
"abstract": "abstract", "summary": "abstract", "overview": "abstract",
|
||||
"introduction": "introduction", "background": "introduction",
|
||||
"results": "results", "findings": "results",
|
||||
"discussion": "discussion", "interpretation": "discussion",
|
||||
"methods": "methods", "methodology": "methods", "materials": "methods",
|
||||
"conclusion": "conclusion",
|
||||
}
|
||||
KEEP_SECTIONS = {"abstract", "introduction", "results", "discussion", "methods", "conclusion"}
|
||||
|
||||
def _standardize_section(self, raw_title: str) -> str:
|
||||
t = raw_title.strip().lower()
|
||||
base = re.sub(r"\d+$", "", t).strip()
|
||||
if base in self.SECTION_MAP:
|
||||
return self.SECTION_MAP[base]
|
||||
if t in self.SECTION_MAP:
|
||||
return self.SECTION_MAP[t]
|
||||
first_word = t.split()[0] if t.split() else ""
|
||||
if first_word in self.SECTION_MAP:
|
||||
return self.SECTION_MAP[first_word]
|
||||
return raw_title
|
||||
|
||||
def _detect_abstract_from_md(self, full_text: str) -> Section | None:
|
||||
"""Detect abstract section from full_text (MinerU markdown).
|
||||
|
||||
1. Look for 'Abstract:' or 'ABSTRACT' heading
|
||||
2. Fallback: look for 'Keywords:', text before it is abstract
|
||||
"""
|
||||
# Method 1: Direct Abstract heading
|
||||
abs_m = re.search(r"(?i)^Abstract[\s:]", full_text, re.MULTILINE)
|
||||
if abs_m:
|
||||
abs_start = abs_m.start()
|
||||
kw_m = re.search(r"(?i)^Keywords?\b", full_text[abs_start:], re.MULTILINE)
|
||||
abs_end = abs_start + kw_m.start() if kw_m else abs_start + 2000
|
||||
# Don't overlap with next section
|
||||
next_h = re.search(r"^##+\s", full_text[abs_end:], re.MULTILINE)
|
||||
if next_h:
|
||||
abs_end = min(abs_end, abs_end + next_h.start())
|
||||
return Section(title="abstract", level=1, parent=None,
|
||||
start_page=1, end_page=1, section_type="heading",
|
||||
position=abs_start, end=abs_end)
|
||||
|
||||
# Method 2: Keywords fallback - text before Keywords: is abstract
|
||||
kw_m = re.search(r"(?i)^Keywords?\b", full_text, re.MULTILINE)
|
||||
if kw_m:
|
||||
kw_pos = kw_m.start()
|
||||
# Find the paragraph ending right before Keywords:
|
||||
# Split by double newlines, take the last paragraph before Keywords:
|
||||
pre_text = full_text[:kw_pos]
|
||||
paras = re.split(r"\n\s*\n", pre_text)
|
||||
if paras:
|
||||
abs_text = paras[-1].strip()
|
||||
if len(abs_text) > 50: # Must be substantial text (not a heading line)
|
||||
abs_start = pre_text.rfind(abs_text)
|
||||
if abs_start >= 0:
|
||||
return Section(title="abstract", level=1, parent=None,
|
||||
start_page=1, end_page=1, section_type="heading",
|
||||
position=abs_start, end=kw_pos)
|
||||
# If last paragraph is too short (>50 chars), check if it's a heading line
|
||||
# and the paragraph before it is the actual abstract
|
||||
elif len(paras) >= 2:
|
||||
abs_text = paras[-2].strip()
|
||||
if len(abs_text) > 100:
|
||||
abs_start = pre_text.rfind(abs_text, 0, kw_pos - len(paras[-1]) - 2)
|
||||
if abs_start >= 0:
|
||||
return Section(title="abstract", level=1, parent=None,
|
||||
start_page=1, end_page=1, section_type="heading",
|
||||
position=abs_start, end=kw_pos)
|
||||
return None
|
||||
|
||||
def _fix_section_positions(self, sections: list, full_text: str) -> list:
|
||||
"""Validate and fix section positions using markdown headings.
|
||||
|
||||
For each section, scan full_text for '## N. SectionName' and use heading position
|
||||
if the content_list position is clearly wrong.
|
||||
"""
|
||||
# Build heading map from markdown: heading_title -> position
|
||||
heading_pat = re.compile(r"^##+\s+(?:\d+(?:\.\d+)*\.?\s+)?(.+)$", re.MULTILINE)
|
||||
md_headings = {}
|
||||
for m in heading_pat.finditer(full_text):
|
||||
title = m.group(1).strip().lower()
|
||||
std = self._standardize_section(title)
|
||||
if std in self.KEEP_SECTIONS and std not in md_headings:
|
||||
md_headings[std] = m.start()
|
||||
|
||||
# For each section, check if position should be corrected
|
||||
for sec in sections:
|
||||
std = self._standardize_section(sec.title)
|
||||
if std in md_headings:
|
||||
correct_pos = md_headings[std]
|
||||
# Only correct if content_list position is significantly off
|
||||
if abs(sec.position - correct_pos) > 200:
|
||||
sec.position = correct_pos
|
||||
|
||||
# Sort by position and recalculate ends
|
||||
sections.sort(key=lambda s: s.position)
|
||||
for i in range(len(sections) - 1):
|
||||
sections[i].end = sections[i + 1].position
|
||||
|
||||
return sections
|
||||
|
||||
def build(self, full_text: str, sections=None, section_text_getter=None, pmid: str = "") -> List[Chunk]:
|
||||
sections = list(sections or [])
|
||||
|
||||
# Step 1: fix section positions using markdown headings
|
||||
sections = self._fix_section_positions(sections, full_text)
|
||||
|
||||
# Step 2: detect abstract if missing
|
||||
has_abstract = any(getattr(s, "title", "") == "abstract" for s in sections)
|
||||
if not has_abstract:
|
||||
abs_sec = self._detect_abstract_from_md(full_text)
|
||||
if abs_sec:
|
||||
sections.insert(0, abs_sec)
|
||||
print(f" \u2713 V15.2.1: detected abstract ({abs_sec.end - abs_sec.position} chars)")
|
||||
|
||||
# Step 3: standardize section names in-place and filter
|
||||
filtered = []
|
||||
for sec in sections:
|
||||
std = self._standardize_section(getattr(sec, "title", "") or "")
|
||||
if std not in self.KEEP_SECTIONS:
|
||||
print(f" - Excluded: {getattr(sec,'title','')} (standardized: {std})")
|
||||
continue
|
||||
sec.title = std
|
||||
filtered.append(sec)
|
||||
|
||||
return super().build(full_text, filtered, section_text_getter, pmid)
|
||||
|
||||
def _split_by_subheadings(self, text: str) -> List[Tuple[str, str]]:
|
||||
matches = list(self.SUBHEADING.finditer(text))
|
||||
if not matches:
|
||||
return [("", text)]
|
||||
parts = []
|
||||
for i, m in enumerate(matches):
|
||||
title = m.group(1).strip()
|
||||
start = m.end()
|
||||
end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
|
||||
body = text[start:end].strip()
|
||||
if body:
|
||||
parts.append((title, body))
|
||||
return parts
|
||||
|
||||
def _build_from_text(self, section_title: str, text: str, pmid: str) -> List[Chunk]:
|
||||
std_title = self._standardize_section(section_title)
|
||||
if std_title not in self.KEEP_SECTIONS:
|
||||
return []
|
||||
text = self._trim_reference_tail(text)
|
||||
clean_lines = self._filter_lines(text)
|
||||
clean_text = "\n".join(clean_lines)
|
||||
sub_parts = self._split_by_subheadings(clean_text)
|
||||
chunks = []
|
||||
prev_sentences: List[str] = []
|
||||
for sub_title, sub_body in sub_parts:
|
||||
full_title = "%s: %s" % (std_title, sub_title) if sub_title else std_title
|
||||
sentences = self._split_sentences(sub_body)
|
||||
if prev_sentences:
|
||||
overlap = prev_sentences[-1]
|
||||
if overlap and (not sentences or sentences[0] != overlap):
|
||||
sentences.insert(0, overlap)
|
||||
sub_chunks = self._pack_chunks(sentences)
|
||||
for idx, chunk_text in enumerate(sub_chunks):
|
||||
if not chunk_text.strip():
|
||||
continue
|
||||
slug = re.sub(r"[^A-Za-z0-9]+", "_", sub_title.lower()).strip("_") or "section"
|
||||
chunk_id = "%s_%s_%s_%04d" % (pmid, slug, std_title[:10] if sub_title else "main", idx + 1)
|
||||
chunks.append(Chunk(
|
||||
chunk_id=chunk_id, section_title=full_title,
|
||||
source_label=self._source_label(section_title),
|
||||
text=chunk_text, char_count=len(chunk_text)))
|
||||
if sentences:
|
||||
prev_sentences = sentences
|
||||
return chunks if chunks else super()._build_from_text(std_title, text, pmid)
|
||||
@@ -0,0 +1,228 @@
|
||||
"""ChunkBuilder v15.2.2 - line-based section detection, char-limit chunking."""
|
||||
from __future__ import annotations
|
||||
import re
|
||||
from typing import List
|
||||
from .chunk_builder import ChunkBuilder, Chunk
|
||||
|
||||
|
||||
class ChunkBuilderV15_2_2(ChunkBuilder):
|
||||
|
||||
SECTION_MAP = {
|
||||
"abstract": "abstract", "summary": "abstract", "overview": "abstract",
|
||||
"introduction": "introduction", "background": "introduction",
|
||||
"results": "results", "findings": "results",
|
||||
"discussion": "discussion", "interpretation": "discussion",
|
||||
"methods": "methods", "methodology": "methods", "materials": "methods",
|
||||
"conclusion": "conclusion",
|
||||
}
|
||||
KEEP_SECTIONS = {"abstract", "introduction", "results", "discussion", "methods", "conclusion"}
|
||||
|
||||
# Keywords to scan in first 50 chars of each line
|
||||
SECTION_KEYWORDS = [
|
||||
"abstract", "summary", "overview",
|
||||
"introduction", "background",
|
||||
"results", "findings",
|
||||
"discussion", "interpretation",
|
||||
"methods", "methodology", "materials and methods",
|
||||
"conclusion",
|
||||
"references", "bibliography",
|
||||
]
|
||||
|
||||
def _match_keyword(self, text50: str):
|
||||
"""Check first 50 chars for any section keyword. Returns standardized name or None."""
|
||||
for kw in self.SECTION_KEYWORDS:
|
||||
if kw in text50:
|
||||
if kw in ("references", "bibliography"):
|
||||
return "references"
|
||||
if kw in self.SECTION_MAP:
|
||||
return self.SECTION_MAP[kw]
|
||||
# Check substring match
|
||||
for sk, sv in self.SECTION_MAP.items():
|
||||
if kw in sk or sk in kw:
|
||||
return sv
|
||||
return kw
|
||||
return None
|
||||
|
||||
def _detect_sections(self, full_text: str):
|
||||
"""Scan md text line by line, first 50 chars per line, match keywords."""
|
||||
lines = full_text.split("\n")
|
||||
boundaries = [] # [(std_name, line_start_pos), ...]
|
||||
pos = 0
|
||||
for line in lines:
|
||||
t50 = line[:50].strip().lower()
|
||||
matched = self._match_keyword(t50)
|
||||
if matched:
|
||||
boundaries.append((matched, pos))
|
||||
pos += len(line) + 1 # +1 for newline
|
||||
|
||||
sections = []
|
||||
|
||||
# Abstract: find "keywords" in first 50 chars
|
||||
# Text before that line is abstract
|
||||
kw_line_pos = None
|
||||
for line in lines:
|
||||
if "keywords" in line[:50].lower().strip():
|
||||
kw_line_pos = full_text.find(line)
|
||||
break
|
||||
|
||||
if kw_line_pos is not None and kw_line_pos > 0:
|
||||
# Take text before Keywords line as abstract
|
||||
# Ensure it's substantial (>200 chars)
|
||||
abstract_end = kw_line_pos
|
||||
# Find the start: go backwards to find the paragraph before keywords
|
||||
pre = full_text[:abstract_end]
|
||||
# Split by double newlines to find paragraphs
|
||||
paras = re.split(r"\n\s*\n", pre)
|
||||
if paras and len(paras[-1].strip()) > 50:
|
||||
abs_text = paras[-1].strip()
|
||||
abs_start = pre.rfind(abs_text)
|
||||
if abs_start >= 0:
|
||||
sections.append(("abstract", abs_start, abstract_end))
|
||||
|
||||
# Build sections from boundaries
|
||||
for i, (name, start) in enumerate(boundaries):
|
||||
if name == "references":
|
||||
break
|
||||
# Skip if within abstract
|
||||
if sections and start < sections[-1][2]:
|
||||
continue
|
||||
# Next boundary or end of text
|
||||
if i + 1 < len(boundaries):
|
||||
end = boundaries[i + 1][1]
|
||||
else:
|
||||
end = len(full_text)
|
||||
sections.append((name, start, end))
|
||||
|
||||
return sections
|
||||
|
||||
def _is_figure_line(self, line: str) -> bool:
|
||||
stripped = line.strip()
|
||||
if not stripped:
|
||||
return False
|
||||
if re.match(r"^!\[.*\]\(.*\)$", stripped):
|
||||
return True
|
||||
if re.match(r"(?i)^(FIG(URE)?\.?\s*\d+|TABLE\s*\d+)", stripped):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _extract_figures_tables(self, text: str):
|
||||
lines = text.splitlines()
|
||||
remove_indices = set()
|
||||
for i, line in enumerate(lines):
|
||||
if self._is_figure_line(line):
|
||||
remove_indices.add(i)
|
||||
if i > 0 and lines[i - 1].strip():
|
||||
remove_indices.add(i - 1)
|
||||
if i + 1 < len(lines) and lines[i + 1].strip():
|
||||
remove_indices.add(i + 1)
|
||||
removed = [lines[i] for i in sorted(remove_indices)]
|
||||
kept = [lines[i] for i in range(len(lines)) if i not in remove_indices]
|
||||
return ("\n".join(kept), "\n".join(removed))
|
||||
|
||||
def _standardize_section(self, raw_title: str) -> str:
|
||||
t = raw_title.strip().lower()
|
||||
base = re.sub(r"\d+$", "", t).strip()
|
||||
if base in self.SECTION_MAP:
|
||||
return self.SECTION_MAP[base]
|
||||
if t in self.SECTION_MAP:
|
||||
return self.SECTION_MAP[t]
|
||||
first_word = t.split()[0] if t.split() else ""
|
||||
if first_word in self.SECTION_MAP:
|
||||
return self.SECTION_MAP[first_word]
|
||||
return raw_title
|
||||
|
||||
SENTENCE_PLACEHOLDERS = {
|
||||
"et al.": "et al<prd>",
|
||||
"Fig.": "Fig<prd>",
|
||||
"Figs.": "Figs<prd>",
|
||||
"Dr.": "Dr<prd>",
|
||||
"Prof.": "Prof<prd>",
|
||||
"e.g.": "e<prd>g<prd>",
|
||||
"i.e.": "i<prd>e<prd>",
|
||||
}
|
||||
|
||||
def _split_sentences(self, text: str) -> List[str]:
|
||||
if not text:
|
||||
return []
|
||||
cleaned = re.sub(r"\s+", " ", str(text)).strip()
|
||||
if not cleaned:
|
||||
return []
|
||||
for key, val in self.SENTENCE_PLACEHOLDERS.items():
|
||||
cleaned = cleaned.replace(key, val)
|
||||
parts = re.split(r"(?<=[.!?])\s+(?=[A-Z0-9(q])", cleaned)
|
||||
sentences = []
|
||||
for part in parts:
|
||||
for key, val in self.SENTENCE_PLACEHOLDERS.items():
|
||||
part = part.replace(val, key)
|
||||
part = part.strip()
|
||||
if part:
|
||||
sentences.append(part)
|
||||
return sentences
|
||||
|
||||
MAX_CHUNK_CHARS = 5200
|
||||
|
||||
def _pack_chunks(self, sentences: List[str]) -> List[str]:
|
||||
if not sentences:
|
||||
return []
|
||||
maxc = self.max_chunk_chars if hasattr(self, 'max_chunk_chars') else self.MAX_CHUNK_CHARS
|
||||
chunks = []
|
||||
current = []
|
||||
cur_len = 0
|
||||
for sent in sentences:
|
||||
slen = len(sent)
|
||||
if current and cur_len + slen + 1 > maxc:
|
||||
chunks.append(" ".join(current))
|
||||
overlap = current[-1]
|
||||
current = [overlap]
|
||||
cur_len = len(overlap) + 1
|
||||
current.append(sent)
|
||||
cur_len += slen + 1
|
||||
if current:
|
||||
chunks.append(" ".join(current))
|
||||
return chunks
|
||||
|
||||
def _build_from_text(self, section_title: str, text: str, pmid: str) -> List[Chunk]:
|
||||
std_title = self._standardize_section(section_title)
|
||||
if std_title not in self.KEEP_SECTIONS:
|
||||
return []
|
||||
text = self._trim_reference_tail(text)
|
||||
clean_lines = self._filter_lines(text)
|
||||
clean_text = "\n".join(clean_lines)
|
||||
clean_text, figures_text = self._extract_figures_tables(clean_text)
|
||||
sentences = self._split_sentences(clean_text)
|
||||
if not sentences:
|
||||
return []
|
||||
raw_chunks = self._pack_chunks(sentences)
|
||||
source_label = self._source_label(section_title)
|
||||
result = []
|
||||
for idx, chunk_text in enumerate(raw_chunks):
|
||||
if not chunk_text.strip():
|
||||
continue
|
||||
slug = re.sub(r"[^A-Za-z0-9]+", "_", std_title.lower()).strip("_") or "section"
|
||||
cid = f"{pmid}_{slug}_{idx + 1:04d}"
|
||||
result.append(Chunk(chunk_id=cid, section_title=std_title,
|
||||
source_label=source_label, text=chunk_text,
|
||||
char_count=len(chunk_text)))
|
||||
if figures_text.strip():
|
||||
fig_sents = self._split_sentences(figures_text)
|
||||
if fig_sents:
|
||||
for idx, ft in enumerate(self._pack_chunks(fig_sents)):
|
||||
if not ft.strip():
|
||||
continue
|
||||
cid = f"{pmid}_figures_and_tables_{idx + 1:04d}"
|
||||
result.append(Chunk(chunk_id=cid, section_title="figures_and_tables",
|
||||
source_label="caption", text=ft,
|
||||
char_count=len(ft)))
|
||||
return result if result else super()._build_from_text(std_title, text, pmid)
|
||||
|
||||
def build(self, full_text: str, sections=None,
|
||||
section_text_getter=None, pmid: str = "") -> List[Chunk]:
|
||||
md_sections = self._detect_sections(full_text)
|
||||
all_chunks = []
|
||||
for sec_type, start, end in md_sections:
|
||||
text = full_text[start:end].strip()
|
||||
if not text:
|
||||
continue
|
||||
chunks = self._build_from_text(sec_type, text, pmid)
|
||||
all_chunks.extend(chunks)
|
||||
return all_chunks
|
||||
@@ -0,0 +1,198 @@
|
||||
"""ChunkBuilder v15.2.3 - strict major-section detection from MinerU markdown.
|
||||
|
||||
Only major paper sections are accepted as section boundaries. Subsection headings
|
||||
such as "2.1 Gene structure analysis" stay inside the parent section text and do
|
||||
not split chunks.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import List
|
||||
|
||||
from .chunk_builder import Chunk
|
||||
from .chunk_builder_v15_2_2 import ChunkBuilderV15_2_2
|
||||
|
||||
|
||||
class ChunkBuilderV15_2_3(ChunkBuilderV15_2_2):
|
||||
"""Use MinerU markdown major headings for section reconstruction."""
|
||||
|
||||
SECTION_MAP = {
|
||||
"abstract": "abstract",
|
||||
"summary": "abstract",
|
||||
"overview": "abstract",
|
||||
"introduction": "introduction",
|
||||
"background": "introduction",
|
||||
"results": "results",
|
||||
"result": "results",
|
||||
"findings": "results",
|
||||
"results and discussion": "results",
|
||||
"result and discussion": "results",
|
||||
"discussion": "discussion",
|
||||
"interpretation": "discussion",
|
||||
"methods": "methods",
|
||||
"method": "methods",
|
||||
"methodology": "methods",
|
||||
"materials": "methods",
|
||||
"materials & methods": "methods",
|
||||
"materials and methods": "methods",
|
||||
"conclusion": "conclusion",
|
||||
"conclusions": "conclusion",
|
||||
"final remarks": "conclusion",
|
||||
"full_text": "full_text",
|
||||
}
|
||||
KEEP_SECTIONS = {
|
||||
"abstract",
|
||||
"introduction",
|
||||
"results",
|
||||
"discussion",
|
||||
"methods",
|
||||
"conclusion",
|
||||
"full_text",
|
||||
}
|
||||
|
||||
MAJOR_SECTION_TERMS = (
|
||||
r"abstract|summary|overview|introduction|background|results?|findings|"
|
||||
r"results and discussion|result and discussion|discussion|interpretation|"
|
||||
r"materials\s*(?:&|and)\s*methods|methods?|methodology|conclusions?|"
|
||||
r"final remarks|references?|bibliography|literature cited"
|
||||
)
|
||||
MAJOR_SECTION_HEADING = re.compile(
|
||||
r"^\s*(?:#+\s*)?(?:\d+\.?\s*)?"
|
||||
rf"({MAJOR_SECTION_TERMS})"
|
||||
r"\s*[::]?\s*$",
|
||||
re.I,
|
||||
)
|
||||
MARKDOWN_INLINE_MAJOR_HEADING = re.compile(
|
||||
r"^\s*#+\s*(?:\d+\.?\s*)?"
|
||||
rf"({MAJOR_SECTION_TERMS})"
|
||||
r"\b",
|
||||
re.I,
|
||||
)
|
||||
INLINE_ABSTRACT = re.compile(r"^\s*(?:#+\s*)?abstract\s*[::]\s+\S+", re.I)
|
||||
REFERENCE_NAME = {"reference", "references", "bibliography", "literature cited"}
|
||||
|
||||
def _standardize_section(self, raw_title: str) -> str:
|
||||
t = re.sub(r"\s+", " ", str(raw_title or "").strip().lower())
|
||||
t = re.sub(r"^#+\s*", "", t).strip()
|
||||
t = re.sub(r"^\d+\.?\s*", "", t).strip()
|
||||
t = t.rstrip("::").strip()
|
||||
t = re.sub(r"\s*&\s*", " & ", t)
|
||||
return self.SECTION_MAP.get(t, t)
|
||||
|
||||
def _match_major_heading(self, line: str):
|
||||
clean = re.sub(r"\s+", " ", str(line or "").strip())
|
||||
if not clean:
|
||||
return None
|
||||
match = self.MAJOR_SECTION_HEADING.match(clean)
|
||||
if not match:
|
||||
match = self.MARKDOWN_INLINE_MAJOR_HEADING.match(clean)
|
||||
if match:
|
||||
raw = match.group(1).lower().strip()
|
||||
if raw in self.REFERENCE_NAME:
|
||||
return "references"
|
||||
return self._standardize_section(raw)
|
||||
# Allow inline abstract lines like "Abstract: this study..." as an
|
||||
# abstract boundary. Other inline headings are not accepted because they
|
||||
# are too easy to confuse with prose.
|
||||
if self.INLINE_ABSTRACT.match(clean):
|
||||
return "abstract"
|
||||
return None
|
||||
|
||||
def _keyword_abstract_section(self, full_text: str):
|
||||
lines = full_text.split("\n")
|
||||
pos = 0
|
||||
keyword_pos = None
|
||||
for line in lines:
|
||||
if re.match(r"(?i)^\s*keywords?\b", line.strip()):
|
||||
keyword_pos = pos
|
||||
break
|
||||
pos += len(line) + 1
|
||||
if keyword_pos is None or keyword_pos <= 0:
|
||||
return None
|
||||
pre = full_text[:keyword_pos]
|
||||
paras = [p.strip() for p in re.split(r"\n\s*\n", pre) if p.strip()]
|
||||
for paragraph in reversed(paras):
|
||||
if len(paragraph) >= 80 and not self._match_major_heading(paragraph):
|
||||
start = pre.rfind(paragraph)
|
||||
if start >= 0:
|
||||
return ("abstract", start, keyword_pos)
|
||||
return None
|
||||
|
||||
def _detect_sections(self, full_text: str):
|
||||
"""Detect only major paper sections from MinerU markdown text."""
|
||||
text = str(full_text or "")
|
||||
if not text.strip():
|
||||
return []
|
||||
|
||||
headings = []
|
||||
pos = 0
|
||||
for line in text.split("\n"):
|
||||
name = self._match_major_heading(line)
|
||||
if name:
|
||||
headings.append((name, pos))
|
||||
pos += len(line) + 1
|
||||
|
||||
sections = []
|
||||
has_abstract_heading = any(name == "abstract" for name, _ in headings)
|
||||
if not has_abstract_heading:
|
||||
abs_sec = self._keyword_abstract_section(text)
|
||||
if abs_sec:
|
||||
sections.append(abs_sec)
|
||||
|
||||
for idx, (name, start) in enumerate(headings):
|
||||
if name == "references":
|
||||
break
|
||||
end = len(text)
|
||||
for _next_name, next_start in headings[idx + 1 :]:
|
||||
end = next_start
|
||||
break
|
||||
if sections and start < sections[-1][2]:
|
||||
continue
|
||||
if end > start:
|
||||
sections.append((name, start, end))
|
||||
|
||||
sections.sort(key=lambda item: item[1])
|
||||
cleaned = []
|
||||
for name, start, end in sections:
|
||||
if cleaned and start < cleaned[-1][2]:
|
||||
prev_name, prev_start, _prev_end = cleaned[-1]
|
||||
if name != prev_name:
|
||||
cleaned[-1] = (prev_name, prev_start, start)
|
||||
cleaned.append((name, start, end))
|
||||
continue
|
||||
cleaned.append((name, start, end))
|
||||
|
||||
if not cleaned:
|
||||
return [("full_text", 0, len(text))]
|
||||
return cleaned
|
||||
|
||||
def _build_from_text(self, section_title: str, text: str, pmid: str) -> List[Chunk]:
|
||||
"""Build chunks without re-adding removed figure/table captions."""
|
||||
std_title = self._standardize_section(section_title)
|
||||
if std_title not in self.KEEP_SECTIONS:
|
||||
return []
|
||||
text = self._trim_reference_tail(text)
|
||||
clean_lines = self._filter_lines(text)
|
||||
clean_text = "\n".join(clean_lines)
|
||||
clean_text, _figures_text = self._extract_figures_tables(clean_text)
|
||||
sentences = self._split_sentences(clean_text)
|
||||
if not sentences:
|
||||
return []
|
||||
raw_chunks = self._pack_chunks(sentences)
|
||||
source_label = self._source_label(std_title)
|
||||
result = []
|
||||
for idx, chunk_text in enumerate(raw_chunks):
|
||||
if not chunk_text.strip():
|
||||
continue
|
||||
slug = re.sub(r"[^A-Za-z0-9]+", "_", std_title.lower()).strip("_") or "section"
|
||||
cid = f"{pmid}_{slug}_{idx + 1:04d}"
|
||||
result.append(
|
||||
Chunk(
|
||||
chunk_id=cid,
|
||||
section_title=std_title,
|
||||
source_label=source_label,
|
||||
text=chunk_text,
|
||||
char_count=len(chunk_text),
|
||||
)
|
||||
)
|
||||
return result
|
||||
@@ -0,0 +1,53 @@
|
||||
"""ChunkBuilder v15.2.4 - v15.2.3 major-section chunks with safe visual cleanup.
|
||||
|
||||
This version keeps the v15.2.3 major-section policy, but fixes the image/table
|
||||
cleanup stage: figure/table image lines and explicit captions are removed without
|
||||
removing adjacent prose. MinerU frequently places valid Results/Discussion text
|
||||
immediately before or after image/caption lines, so neighbor deletion causes
|
||||
coverage loss.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
from .chunk_builder_v15_2_3 import ChunkBuilderV15_2_3
|
||||
|
||||
|
||||
class ChunkBuilderV15_2_4(ChunkBuilderV15_2_3):
|
||||
"""Major-section chunker with non-destructive figure/table cleanup."""
|
||||
|
||||
MARKDOWN_IMAGE_LINE = re.compile(r"^\s*!\[.*?\]\(.*?\)\s*$")
|
||||
FIGURE_TABLE_CAPTION_LINE = re.compile(
|
||||
r"^\s*(?:"
|
||||
r"fig(?:ure)?\.?\s*\d+[A-Za-z]?"
|
||||
r"|table\s*\d+[A-Za-z]?"
|
||||
r"|supplementary\s+(?:fig(?:ure)?|fig|table)\s*\w+"
|
||||
r")\b",
|
||||
re.I,
|
||||
)
|
||||
|
||||
def _is_v15_2_4_visual_line(self, line: str) -> bool:
|
||||
stripped = str(line or "").strip()
|
||||
if not stripped:
|
||||
return False
|
||||
if self.MARKDOWN_IMAGE_LINE.match(stripped):
|
||||
return True
|
||||
if self.FIGURE_TABLE_CAPTION_LINE.match(stripped):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _extract_figures_tables(self, text: str):
|
||||
"""Remove only visual/caption lines, never neighboring prose lines.
|
||||
|
||||
v15.2.2 removed the line before and after every image/caption. That is
|
||||
unsafe for MinerU markdown because valid section prose can be adjacent to
|
||||
figure images, especially in two-column papers and around page breaks.
|
||||
"""
|
||||
kept = []
|
||||
removed = []
|
||||
for line in str(text or "").splitlines():
|
||||
if self._is_v15_2_4_visual_line(line):
|
||||
removed.append(line)
|
||||
else:
|
||||
kept.append(line)
|
||||
return "\n".join(kept), "\n".join(removed)
|
||||
@@ -0,0 +1,90 @@
|
||||
from __future__ import annotations
|
||||
import csv,json,re
|
||||
from pathlib import Path
|
||||
from typing import Any,Dict,Iterable,List,Optional
|
||||
from .chunk_builder import Chunk
|
||||
from .sectionizer_v16_1_1 import build_section_spans, classify_heading, is_major_heading, is_subheading
|
||||
REF=re.compile(r"^\s*(?:#+\s*)?(?:references?|bibliography|literature cited)\s*$",re.I)
|
||||
FIG=re.compile(r"^\s*(?:#+\s*)?(?:fig(?:ure)?\.?\s*\d+|supplementary\s+fig(?:ure)?\.?\s*S?\d+)[\.:\s-]",re.I)
|
||||
TAB=re.compile(r"^\s*(?:#+\s*)?(?:table\s+\d+|supplementary\s+table\s*S?\d+)[\.:\s-]",re.I)
|
||||
IMG=re.compile(r"^\s*(?:!\[[^\]]*\]\([^)]*\)|\[image\]|<image[^>]*>|image\s*placeholder)\s*$",re.I)
|
||||
HEAD=re.compile(r"^\s{0,3}#{1,6}\s+(.+?)\s*$")
|
||||
HTML=re.compile(r"<\s*(?:table|tr|td|th)\b",re.I); PIPE=re.compile(r"^\s*\|.+\|\s*$")
|
||||
FORM=re.compile(r"\$([^$]{1,180})\$"); HYPH=re.compile(r"(?<=[A-Za-z])-\s+\n\s*(?=[a-z])"); LB=re.compile(r"(?<![.!?:;])\n(?=\S)")
|
||||
class ChunkBuilderV15_2_5:
|
||||
def __init__(self,max_chunk_chars:int=5200,output_dir:Optional[Path]=None,overlap_sentences:int=1):
|
||||
self.max_chunk_chars=max_chunk_chars; self.output_dir=Path(output_dir) if output_dir else None; self.overlap_sentences=overlap_sentences
|
||||
def build(self,full_text:str,sections:Optional[Iterable]=None,section_text_getter=None,pmid:str="")->List[Chunk]:
|
||||
paper_id=pmid or "paper"; blocks=self._blocks(full_text or "",paper_id); units=self._units(blocks,paper_id); chunks=[]
|
||||
for u in units:
|
||||
c=Chunk(chunk_id=u["relation_unit_id"],section_title=u["section_type"],source_label=u["source_label"],text=u["text"],char_count=len(u["text"]),endpoint_candidates={})
|
||||
c.block_types=u.get("block_types",[]); c.source_block_ids=u.get("block_ids",[]); c.relation_unit_id=u["relation_unit_id"]; chunks.append(c)
|
||||
self._write(blocks,units,paper_id); return chunks
|
||||
def _clean(self,t:str)->str:
|
||||
t=str(t or "").replace("\r\n","\n").replace("\r","\n"); t=HYPH.sub("",t); t=LB.sub(" ",t)
|
||||
t=FORM.sub(lambda m:"[FORMULA: "+re.sub(r"\s+"," ",m.group(1)).strip()+"]",t)
|
||||
return re.sub(r"\s+"," ",t).strip()
|
||||
def _blocks(self,text:str,paper_id:str)->List[Dict[str,Any]]:
|
||||
out=[]; para=[]; in_refs=False; order=0
|
||||
def emit(raw:str):
|
||||
nonlocal order,in_refs
|
||||
raw=raw.rstrip("\n");
|
||||
if not raw: return
|
||||
order+=1; s=raw.strip(); clean=self._clean(raw); mh=HEAD.match(s); htxt=mh.group(1).strip() if mh else s; sec=classify_heading(htxt)
|
||||
is_ref=bool(REF.match(s) or sec=="references" or in_refs)
|
||||
if REF.match(s) or sec=="references": in_refs=True; is_ref=True
|
||||
is_img=bool(IMG.match(s)); is_fig=bool(FIG.match(s)); is_tab=bool(TAB.match(s) or HTML.search(s) or PIPE.match(s))
|
||||
is_major=bool(mh and is_major_heading(htxt)) or (sec not in {None,"references","subheading_inside_parent"} and len(s)<=80)
|
||||
is_sub=(not is_major) and (bool(mh) or is_subheading(htxt))
|
||||
vtype="figure_caption" if is_fig else "table_block" if is_tab else "image_placeholder" if is_img else "heading" if is_major else "subheading" if is_sub else "paragraph"
|
||||
keep=not is_ref and not is_img and bool(clean) and vtype!="heading"
|
||||
reason="references_after_heading" if is_ref else "image_placeholder" if is_img else "empty_after_cleaning" if not clean else "heading_metadata_only" if vtype=="heading" else ""
|
||||
out.append({"block_id":f"{paper_id}_b{order:06d}","paper_id":paper_id,"page":None,"order":order,"mineru_type":"text","v16_block_type":vtype,"raw_text":raw,"clean_text":clean,"char_count":len(clean),"is_major_heading":is_major,"is_subheading":is_sub,"section_type":sec if is_major else "","is_reference":is_ref,"is_figure_caption":is_fig,"is_table_block":is_tab,"is_image_placeholder":is_img,"is_formula_heavy":clean.count("[FORMULA:")>=3,"keep_for_llm":keep,"drop_reason":reason})
|
||||
for line in text.splitlines():
|
||||
if not line.strip():
|
||||
if para: emit("\n".join(para)); para=[]
|
||||
continue
|
||||
s=line.strip(); special=HEAD.match(s) or REF.match(s) or FIG.match(s) or TAB.match(s) or IMG.match(s) or HTML.search(s) or PIPE.match(s)
|
||||
if special:
|
||||
if para: emit("\n".join(para)); para=[]
|
||||
emit(line)
|
||||
else: para.append(line)
|
||||
if para: emit("\n".join(para))
|
||||
return out
|
||||
def _sentences(self,t:str)->List[str]:
|
||||
t=re.sub(r"\s+"," ",t).strip(); return [p.strip() for p in re.split(r"(?<=[.!?])\s+(?=[A-Z0-9\[])",t) if p.strip()] if t else []
|
||||
def _units(self,blocks:List[Dict[str,Any]],paper_id:str)->List[Dict[str,Any]]:
|
||||
units=[]; sec="full_text"; buf=[]; bids=[]; btypes=[]; no=0
|
||||
def flush(label="text"):
|
||||
nonlocal buf,bids,btypes,no
|
||||
text=" ".join(x for x in buf if x).strip()
|
||||
if not text: buf=[]; bids=[]; btypes=[]; return
|
||||
no+=1; units.append({"relation_unit_id":f"{paper_id}_{sec}_{no:04d}","paper_id":paper_id,"section_type":sec,"source_label":label,"block_types":sorted(set(btypes)) or ["paragraph"],"block_ids":list(bids),"char_count":len(text),"text":text}); buf=[]; bids=[]; btypes=[]
|
||||
for b in blocks:
|
||||
if b.get("is_reference"): break
|
||||
if b.get("is_major_heading"):
|
||||
flush(); sec=b.get("section_type") or "full_text"; continue
|
||||
if not b.get("keep_for_llm"): continue
|
||||
vt=b.get("v16_block_type") or "paragraph"
|
||||
if vt in {"figure_caption","table_block"}:
|
||||
flush(); label="figure_caption" if vt=="figure_caption" else "table"; ssec="figure_caption" if vt=="figure_caption" else "table"
|
||||
units.append({"relation_unit_id":f"{paper_id}_{ssec}_{len(units)+1:04d}","paper_id":paper_id,"section_type":ssec,"source_label":label,"block_types":[vt],"block_ids":[b["block_id"]],"char_count":len(b.get("clean_text") or ""),"text":b.get("clean_text") or ""}); continue
|
||||
for sent in self._sentences(b.get("clean_text") or "") or [b.get("clean_text") or ""]:
|
||||
if sum(len(x) for x in buf)+len(sent)+1>self.max_chunk_chars and buf: flush()
|
||||
buf.append(sent); bids.append(b["block_id"]); btypes.append(vt)
|
||||
flush(); return units
|
||||
def _write(self,blocks,units,paper_id):
|
||||
if not self.output_dir: return
|
||||
out=self.output_dir; out.mkdir(parents=True,exist_ok=True)
|
||||
def jsonl(name,rows):
|
||||
with open(out/name,"w",encoding="utf-8") as f:
|
||||
for r in rows: f.write(json.dumps(r,ensure_ascii=False)+"\n")
|
||||
jsonl("01_mineru_blocks_raw.jsonl",blocks); jsonl("02_mineru_blocks_clean.jsonl",blocks); jsonl("04_relation_units.jsonl",units)
|
||||
(out/"03_section_spans.json").write_text(json.dumps([s.__dict__ for s in build_section_spans(blocks)],ensure_ascii=False,indent=2),encoding="utf-8")
|
||||
denom=sum(b.get("char_count",0) for b in blocks if b.get("keep_for_llm") and not b.get("is_reference") and not b.get("is_image_placeholder")); numer=min(denom,sum(u.get("char_count",0) for u in units)) if denom else 0
|
||||
with open(out/"05_relation_units_coverage.csv","w",newline="",encoding="utf-8") as f:
|
||||
w=csv.DictWriter(f,fieldnames=["paper_id","covered_chars","eligible_chars","coverage_ratio","unit_count"]); w.writeheader(); w.writerow({"paper_id":paper_id,"covered_chars":numer,"eligible_chars":denom,"coverage_ratio":round(numer/denom,4) if denom else 0,"unit_count":len(units)})
|
||||
with open(out/"dropped_block_audit.csv","w",newline="",encoding="utf-8") as f:
|
||||
fields=["block_id","order","v16_block_type","drop_reason","raw_text"]; w=csv.DictWriter(f,fieldnames=fields); w.writeheader(); [w.writerow({k:b.get(k,"") for k in fields}) for b in blocks if not b.get("keep_for_llm")]
|
||||
summary={"paper_id":paper_id,"block_count":len(blocks),"kept_block_count":sum(1 for b in blocks if b.get("keep_for_llm")),"relation_unit_count":len(units),"figure_caption_unit_count":sum(1 for u in units if u.get("source_label")=="figure_caption"),"table_unit_count":sum(1 for u in units if u.get("source_label")=="table"),"coverage_ratio":round(numer/denom,4) if denom else 0}
|
||||
(out/"mineru_block_summary.json").write_text(json.dumps(summary,ensure_ascii=False,indent=2),encoding="utf-8")
|
||||
@@ -0,0 +1,325 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
V14 实体链接优化器
|
||||
实现模糊匹配和实体标准化
|
||||
"""
|
||||
|
||||
import re
|
||||
from difflib import SequenceMatcher
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class EntityLinker:
|
||||
"""
|
||||
实体链接优化器
|
||||
|
||||
功能:
|
||||
1. 精确匹配
|
||||
2. 别名匹配 (scientific_chunks.json)
|
||||
3. 模糊匹配 (编辑距离)
|
||||
4. 子串匹配
|
||||
5. 基因别名扩展
|
||||
6. 物种前缀保护(防止跨物种错误链接)
|
||||
"""
|
||||
|
||||
# 物种前缀映射表
|
||||
SPECIES_PREFIXES = {
|
||||
'Ci': 'Carya illinoinensis', # 山核桃/pecan
|
||||
'At': 'Arabidopsis thaliana', # 拟南芥
|
||||
'Os': 'Oryza sativa', # 水稻
|
||||
'Zm': 'Zea mays', # 玉米
|
||||
'Sl': 'Solanum lycopersicum', # 番茄
|
||||
'Gm': 'Glycine max', # 大豆
|
||||
'Pt': 'Populus trichocarpa', # 杨树
|
||||
'Vv': 'Vitis vinifera', # 葡萄
|
||||
'St': 'Solanum tuberosum', # 土豆
|
||||
'Bd': 'Brachypodium distachyon', # 短柄草
|
||||
'Hv': 'Hordeum vulgare', # 大麦
|
||||
'Ta': 'Triticum aestivum', # 小麦
|
||||
'Nb': 'Nicotiana benthamiana', # 本氏烟草
|
||||
}
|
||||
|
||||
def __init__(self, pc_data_loader=None):
|
||||
self.pc_loader = pc_data_loader
|
||||
self.alias_map: Dict[str, str] = {} # 小写别名 -> 标准名
|
||||
self.gene_aliases: Dict[str, List[str]] = {} # 基因 -> 别名列表
|
||||
self.entity_dict: set = set() # 有效实体集合
|
||||
|
||||
# 加载PC数据
|
||||
if pc_data_loader:
|
||||
self._load_pc_data()
|
||||
|
||||
# 匹配缓存
|
||||
self._match_cache: Dict[str, Tuple[str, float, str]] = {} # name -> (resolved, score, method)
|
||||
|
||||
def _load_pc_data(self):
|
||||
"""加载PC数据"""
|
||||
try:
|
||||
self.alias_map = self.pc_loader.load_scientific_chunks()
|
||||
print(f"✓ 加载别名映射: {len(self.alias_map)} 条")
|
||||
except Exception as e:
|
||||
print(f"⚠️ 别名映射加载失败: {e}")
|
||||
|
||||
try:
|
||||
self.gene_aliases = self.pc_loader.load_gene_aliases()
|
||||
print(f"✓ 加载基因别名: {len(self.gene_aliases)} 个基因")
|
||||
except Exception as e:
|
||||
print(f"⚠️ 基因别名加载失败: {e}")
|
||||
|
||||
try:
|
||||
self.entity_dict = self.pc_loader.load_entity_dict()
|
||||
print(f"✓ 加载实体目录: {len(self.entity_dict)} 个")
|
||||
except Exception as e:
|
||||
print(f"⚠️ 实体目录加载失败: {e}")
|
||||
|
||||
def link(self, entity_name: str, entity_type: str = "") -> Dict:
|
||||
"""
|
||||
实体链接主函数
|
||||
|
||||
Returns:
|
||||
{
|
||||
'original': 原始名,
|
||||
'resolved': 标准化名,
|
||||
'score': 匹配分数 (0-1),
|
||||
'method': 匹配方法,
|
||||
'aliases': 别名列表,
|
||||
'is_valid': 是否有效实体
|
||||
}
|
||||
"""
|
||||
# 检查缓存
|
||||
cache_key = f"{entity_name.lower()}:{entity_type}"
|
||||
if cache_key in self._match_cache:
|
||||
resolved, score, method = self._match_cache[cache_key]
|
||||
return self._build_result(entity_name, resolved, score, method)
|
||||
|
||||
# 步骤0: 检测物种前缀
|
||||
source_species = self._detect_species_prefix(entity_name)
|
||||
|
||||
# 1. 精确匹配 (忽略大小写)
|
||||
resolved, score, method = self._exact_match(entity_name)
|
||||
if score >= 1.0:
|
||||
# 检查是否跨物种
|
||||
if source_species and not self._is_same_species(resolved, source_species):
|
||||
# 跨物种,保持原样
|
||||
self._match_cache[cache_key] = (entity_name, 0.0, 'cross_species_blocked')
|
||||
return self._build_result(entity_name, entity_name, 0.0, 'cross_species_blocked')
|
||||
self._match_cache[cache_key] = (resolved, score, method)
|
||||
return self._build_result(entity_name, resolved, score, method)
|
||||
|
||||
# 2. 别名匹配
|
||||
resolved, score, method = self._alias_match(entity_name)
|
||||
if score >= 0.95:
|
||||
# 检查是否跨物种
|
||||
if source_species and not self._is_same_species(resolved, source_species):
|
||||
self._match_cache[cache_key] = (entity_name, 0.0, 'cross_species_blocked')
|
||||
return self._build_result(entity_name, entity_name, 0.0, 'cross_species_blocked')
|
||||
self._match_cache[cache_key] = (resolved, score, method)
|
||||
return self._build_result(entity_name, resolved, score, method)
|
||||
|
||||
# 3. 基因别名匹配
|
||||
if entity_type == 'gene' or not entity_type:
|
||||
resolved, score, method = self._gene_alias_match(entity_name)
|
||||
if score >= 0.95:
|
||||
# 检查是否跨物种
|
||||
if source_species and not self._is_same_species(resolved, source_species):
|
||||
self._match_cache[cache_key] = (entity_name, 0.0, 'cross_species_blocked')
|
||||
return self._build_result(entity_name, entity_name, 0.0, 'cross_species_blocked')
|
||||
self._match_cache[cache_key] = (resolved, score, method)
|
||||
return self._build_result(entity_name, resolved, score, method)
|
||||
|
||||
# 4. 模糊匹配 (编辑距离) - 添加物种约束
|
||||
resolved, score, method = self._fuzzy_match(entity_name, source_species=source_species)
|
||||
if score >= 0.85:
|
||||
self._match_cache[cache_key] = (resolved, score, method)
|
||||
return self._build_result(entity_name, resolved, score, method)
|
||||
|
||||
# 5. 子串匹配 - 添加物种约束
|
||||
resolved, score, method = self._substring_match(entity_name, source_species=source_species)
|
||||
if score >= 0.8:
|
||||
self._match_cache[cache_key] = (resolved, score, method)
|
||||
return self._build_result(entity_name, resolved, score, method)
|
||||
|
||||
# 6. 无匹配,返回原样
|
||||
self._match_cache[cache_key] = (entity_name, 0.0, 'no_match')
|
||||
return self._build_result(entity_name, entity_name, 0.0, 'no_match')
|
||||
|
||||
def _exact_match(self, name: str) -> Tuple[str, float, str]:
|
||||
"""精确匹配 (忽略大小写)"""
|
||||
lower_name = name.lower()
|
||||
|
||||
# 检查别名映射
|
||||
if lower_name in self.alias_map:
|
||||
return self.alias_map[lower_name], 1.0, 'exact_alias'
|
||||
|
||||
# 检查是否在实体目录中
|
||||
if name in self.entity_dict:
|
||||
return name, 1.0, 'exact_entity'
|
||||
|
||||
return name, 0.0, 'no_match'
|
||||
|
||||
def _alias_match(self, name: str) -> Tuple[str, float, str]:
|
||||
"""别名匹配 - V14.2优化: 添加缓存避免重复计算"""
|
||||
lower_name = name.lower()
|
||||
|
||||
# 直接查找
|
||||
if lower_name in self.alias_map:
|
||||
return self.alias_map[lower_name], 1.0, 'alias_direct'
|
||||
|
||||
# V14.2优化: 只对短名称进行清理后查找(基因名通常较短)
|
||||
if len(lower_name) <= 20:
|
||||
cleaned = re.sub(r'[\s\-_.]+', '', lower_name)
|
||||
# 限制搜索范围:只检查前1000个别名
|
||||
for alias, standard in list(self.alias_map.items())[:1000]:
|
||||
cleaned_alias = re.sub(r'[\s\-_.]+', '', alias)
|
||||
if cleaned == cleaned_alias:
|
||||
return standard, 0.98, 'alias_cleaned'
|
||||
|
||||
return name, 0.0, 'no_match'
|
||||
|
||||
def _gene_alias_match(self, name: str) -> Tuple[str, float, str]:
|
||||
"""基因别名匹配 - V14.2优化: 使用直接查找"""
|
||||
upper_name = name.upper()
|
||||
|
||||
# V14.2优化: 先直接查找基因名
|
||||
if upper_name in self.gene_aliases:
|
||||
return upper_name, 1.0, 'gene_exact'
|
||||
|
||||
# 检查是否匹配基因名(大小写不敏感)
|
||||
for gene in self.gene_aliases:
|
||||
if upper_name == gene.upper():
|
||||
return gene, 1.0, 'gene_exact'
|
||||
|
||||
# V14.2优化: 限制别名搜索范围
|
||||
for gene, aliases in list(self.gene_aliases.items())[:100]:
|
||||
for alias in aliases:
|
||||
if upper_name == alias.upper():
|
||||
return gene, 0.98, 'gene_alias'
|
||||
|
||||
return name, 0.0, 'no_match'
|
||||
|
||||
def _detect_species_prefix(self, name: str) -> Optional[str]:
|
||||
"""检测基因名的物种前缀(支持大小写变体)"""
|
||||
for prefix, species in self.SPECIES_PREFIXES.items():
|
||||
# 支持 Ci/CISPL, At/ATSPL, Os/OSSPL 等变体
|
||||
if name.startswith(prefix) or name.startswith(prefix.upper()):
|
||||
return species
|
||||
return None
|
||||
|
||||
def _is_same_species(self, resolved_name: str, source_species: Optional[str]) -> bool:
|
||||
"""检查解析后的名称是否与源物种一致"""
|
||||
# 如果源没有物种前缀,允许任何匹配
|
||||
if source_species is None:
|
||||
return True
|
||||
|
||||
# 获取resolved_name的物种前缀(支持大小写变体)
|
||||
resolved_species = None
|
||||
for prefix, species in self.SPECIES_PREFIXES.items():
|
||||
if resolved_name.startswith(prefix) or resolved_name.startswith(prefix.upper()):
|
||||
resolved_species = species
|
||||
break
|
||||
|
||||
# 如果resolved_name没有物种前缀,无法判断,默认允许
|
||||
if resolved_species is None:
|
||||
return True
|
||||
|
||||
# 核心逻辑:如果源有物种前缀,目标必须有相同前缀才允许
|
||||
# 例如:CiSPL6a(山核桃)只能链接到Ci开头的基因,不能链接到AT开头的基因
|
||||
return resolved_species == source_species
|
||||
|
||||
def _fuzzy_match(self, name: str, threshold: float = 0.85,
|
||||
source_species: Optional[str] = None) -> Tuple[str, float, str]:
|
||||
"""模糊匹配 (编辑距离),带物种约束 - V14.2优化: 完全禁用模糊匹配以提高性能"""
|
||||
# V14.2优化: 模糊匹配太慢,完全禁用
|
||||
# 如果需要模糊匹配,应该使用更高效的算法(如倒排索引)
|
||||
return name, 0.0, 'no_match'
|
||||
|
||||
def _substring_match(self, name: str,
|
||||
source_species: Optional[str] = None) -> Tuple[str, float, str]:
|
||||
"""子串匹配,带物种约束 - V14.2优化: 完全禁用子串匹配以提高性能"""
|
||||
# V14.2优化: 子串匹配太慢,完全禁用
|
||||
return name, 0.0, 'no_match'
|
||||
|
||||
def _build_result(self, original: str, resolved: str, score: float,
|
||||
method: str) -> Dict:
|
||||
"""构建结果"""
|
||||
# 获取别名
|
||||
aliases = []
|
||||
if resolved in self.gene_aliases:
|
||||
aliases = self.gene_aliases[resolved]
|
||||
|
||||
# 验证有效性
|
||||
is_valid = resolved in self.entity_dict or score > 0.8
|
||||
|
||||
return {
|
||||
'original': original,
|
||||
'resolved': resolved,
|
||||
'score': round(score, 3),
|
||||
'method': method,
|
||||
'aliases': aliases,
|
||||
'is_valid': is_valid
|
||||
}
|
||||
|
||||
def batch_link(self, entity_names: List[str],
|
||||
entity_types: List[str] = None) -> List[Dict]:
|
||||
"""批量实体链接"""
|
||||
if entity_types is None:
|
||||
entity_types = [''] * len(entity_names)
|
||||
|
||||
results = []
|
||||
for name, etype in zip(entity_names, entity_types):
|
||||
result = self.link(name, etype)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
def get_stats(self) -> Dict:
|
||||
"""获取统计信息"""
|
||||
return {
|
||||
'alias_map_size': len(self.alias_map),
|
||||
'gene_aliases_size': len(self.gene_aliases),
|
||||
'entity_dict_size': len(self.entity_dict),
|
||||
'cache_size': len(self._match_cache)
|
||||
}
|
||||
|
||||
|
||||
# 全局单例
|
||||
_entity_linker: Optional[EntityLinker] = None
|
||||
|
||||
def get_entity_linker(pc_data_loader=None) -> EntityLinker:
|
||||
"""获取实体链接器单例"""
|
||||
global _entity_linker
|
||||
if _entity_linker is None:
|
||||
_entity_linker = EntityLinker(pc_data_loader)
|
||||
return _entity_linker
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 测试
|
||||
print("="*80)
|
||||
print("🧪 实体链接优化器测试")
|
||||
print("="*80)
|
||||
|
||||
from pc_data_loader import get_pc_data_loader
|
||||
|
||||
loader = get_pc_data_loader()
|
||||
linker = EntityLinker(loader)
|
||||
|
||||
# 测试用例
|
||||
test_cases = [
|
||||
"SUC5",
|
||||
"sucrose transporter 5",
|
||||
"Suc transporter 5",
|
||||
"AT1G71890",
|
||||
"sucrose transporter",
|
||||
"SUCROSE TRANSPORTER 5",
|
||||
"UnknownGene123",
|
||||
]
|
||||
|
||||
print("\n📋 实体链接测试:")
|
||||
for name in test_cases:
|
||||
result = linker.link(name, 'gene')
|
||||
print(f"\n '{name}'")
|
||||
print(f" → '{result['resolved']}' (score: {result['score']}, method: {result['method']})")
|
||||
if result['aliases']:
|
||||
print(f" 别名: {result['aliases'][:3]}...")
|
||||
@@ -0,0 +1,107 @@
|
||||
from __future__ import annotations
|
||||
import re
|
||||
from typing import List, Tuple
|
||||
|
||||
ALLOWED_ENTITY_TYPES = {
|
||||
"gene","gene_family","gene_group","protein","enzyme","transcription_factor","miRNA","lncRNA","transcript","promoter","cis_element","motif","domain","sequence_feature","genetic_variant","QTL","marker","SNP","haplotype","allele","fish_species","mollusk","crustacean","algae","seaweed","coral","plankton","species","cultivar","ecotype","population","genotype","water_quality_parameter","pathogen","parasite","virus","disease","microbe","tissue","organ","cell_type","growth_stage","life_stage","developmental_stage","treatment_or_stress","environmental_factor","feed_ingredient","feed_additive","aquaculture_system","hormone","metabolite","compound","nutrient","phenotype","trait","process","pathway","compartment","method","tool_or_database","dataset_accession","experimental_material","geographic_region","genome","chromosome","genomic_region"}
|
||||
ALLOWED_RELATION_CATEGORIES = {"catalog_count","family_membership","classification","chromosomal_location","gene_structure","motif_domain","cis_element","duplication_synteny","orthology_homology","growth_performance","feed_conversion_efficiency","survival_rate","disease_resistance","immune_response","water_quality_effect","expression_tissue","expression_condition","expression_time","stress_response","subcellular_localization","phenotype_effect","protein_property","variant_trait_association","QTL_trait_association","marker_trait_association","pathway_participation","metabolic_change","hormone_response","pathogen_response","developmental_regulation","reproduction_related","nutritional_value","physiological_measurement","stocking_density_effect","compound_effect","enzyme_activity","protein_interaction","regulatory_interaction","method_tool","experimental_material","dataset_accession","background_prior_work","figure_caption_observation","table_structured_fact","other_supported"}
|
||||
MAIN_SOURCE_TYPES = {"gene","gene_family","gene_group","protein","enzyme","transcription_factor","miRNA","lncRNA","transcript","promoter","cis_element","motif","domain","sequence_feature","genetic_variant","QTL","marker","SNP","haplotype","allele","fish_species","mollusk","crustacean","algae","seaweed","coral","plankton","species","cultivar","genotype","water_quality_parameter","pathogen","parasite","virus","disease","metabolite","compound","hormone","feed_ingredient","feed_additive","phenotype","trait","process","pathway","genome","chromosome","genomic_region"}
|
||||
BAD_GENERIC = {"various","both","one","two","three","four","five","several","many","levels","treatments","treatment","control","response","responses","expression","process","result","results","analysis","data","group","groups","type","types","figure","table"}
|
||||
QA_WEIGHTS={"main_kg":1.0,"dataset_metadata":0.95,"method_relations":0.9,"background_relations":0.7,"candidate_review":0.45,"caption_relations":0.35,"rejected":0.0}
|
||||
STARTING_FRAGMENT=re.compile(r"^(?:in|of|to|by|with|and|or|for|from|under|during|between|among)\b",re.I)
|
||||
URL_RE=re.compile(r"https?://|www\.",re.I)
|
||||
NUMERIC_ONLY_RE=re.compile(r"^[\s\d.,+\-??%/()]+$")
|
||||
REF_RE=re.compile(r"\bet\s+al\.?\b|^\[?\d+\]?$",re.I)
|
||||
UNIT_RE=re.compile(r"^(?:mg|g|kg|ml|l|mm|cm|h|hr|hrs|min|s|bp|kb|mb|aa|nt|percent|degc|celsius|ppt|psu|dof)$",re.I)
|
||||
ALIASES={"software":"tool_or_database","database":"tool_or_database","tool":"tool_or_database","stress":"treatment_or_stress","treatment":"treatment_or_stress","variant":"genetic_variant","cis-element":"cis_element","tf":"transcription_factor","chromosomal_location":"chromosome","chromosome_location":"chromosome","fish":"fish_species","shrimp":"crustacean","prawn":"crustacean","crab":"crustacean","lobster":"crustacean","bivalve":"mollusk","oyster":"mollusk","mussel":"mollusk","clam":"mollusk","scallop":"mollusk","snail":"mollusk","abalone":"mollusk","microalgae":"algae","macroalgae":"seaweed","kelp":"seaweed","bacteria":"pathogen","bacterium":"pathogen"}
|
||||
REL_NORM=[(re.compile(r"^(?:decreased_in|reduced_in|repressed_in|down[-_ ]?regulated_in)$",re.I),"downregulated_in","expression synonym"),(re.compile(r"^(?:increased_in|induced_in|enhanced_expression_in|up[-_ ]?regulated_in)$",re.I),"upregulated_in","expression synonym"),(re.compile(r"^(?:expression_changed_in|changed_in|differentially_expressed_in)$",re.I),"differentially_expressed_under","expression synonym"),(re.compile(r"^(?:located_in|localized_in|localised_in|localized_to|localised_to)$",re.I),"localizes_to","localization synonym"),(re.compile(r"^(?:belongs_to|classified_into|grouped_into|separated_into)$",re.I),"member_of","classification synonym"),(re.compile(r"^(?:associated_with_trait|linked_to_trait)$",re.I),"associated_with","trait association synonym"),(re.compile(r"^(?:confers_resistance_to|improves_resistance_to)$",re.I),"increases_resistance_to","resistance synonym"),(re.compile(r"^(?:responds_to|induced_by|affected_by)$",re.I),"responds_to","response synonym")]
|
||||
REL_NORM += [
|
||||
(re.compile(r"^expression_level_measured_in$", re.I), "expression_measured_in", "expression measurement synonym"),
|
||||
(re.compile(r"^up[-_ ]?regulated_post_infection$", re.I), "upregulated_after_infection", "post-infection expression synonym"),
|
||||
(re.compile(r"^no_difference_in_expression$", re.I), "not_differentially_expressed_between", "negative expression comparison"),
|
||||
(re.compile(r"^regulates_.+_resistance_through$", re.I), "regulates_resistance_via", "resistance mechanism synonym"),
|
||||
(re.compile(r"^is_(?:a_)?sa[-_ ]related_defense_gene$", re.I), "associated_with_sa_mediated_defense", "SA defense synonym"),
|
||||
(re.compile(r"^is_(?:a_)?ja[-_ ]related_defense_gene$", re.I), "associated_with_ja_mediated_defense", "JA defense synonym"),
|
||||
(re.compile(r"^is_(?:an?_)?et[-_ ]related_defense_gene$", re.I), "associated_with_et_mediated_defense", "ET defense synonym"),
|
||||
# V16.aqua: aquaculture-specific normalization rules (kept after the plant rules).
|
||||
(re.compile(r"^improved_fcr$", re.I), "improved_feed_conversion", "FCR synonym"),
|
||||
(re.compile(r"^improved_sgr$", re.I), "improved_specific_growth_rate", "SGR synonym"),
|
||||
(re.compile(r"^reduced_fcr$", re.I), "improved_feed_conversion", "FCR direction synonym"),
|
||||
(re.compile(r"^increased_wg$", re.I), "increased_weight_gain", "WG synonym"),
|
||||
(re.compile(r"^cumulative_mortality_of$", re.I), "mortality_of", "mortality synonym"),
|
||||
(re.compile(r"^challenge_test(?:ing|ed)?_with$", re.I), "challenged_with", "challenge synonym"),
|
||||
(re.compile(r"^showed_tolerance_to$", re.I), "tolerant_to", "tolerance synonym"),
|
||||
]
|
||||
WEAK_RELATIONS = {"is", "are", "has", "have", "supported", "provided"}
|
||||
ADMIN_EVIDENCE_RE = re.compile(
|
||||
r"\b(?:this work was supported by|funded by|funding was provided by|"
|
||||
r"authors? declare(?:s)? (?:that )?(?:they have )?no conflict|"
|
||||
r"declaration of competing interest|we thank)\b",
|
||||
re.I,
|
||||
)
|
||||
|
||||
def clean_text(v:object)->str:
|
||||
return re.sub(r"\s+"," ",str(v or "")).strip().strip(" \t\r\n\"'`")
|
||||
|
||||
def normalize_entity_type(v:object)->str:
|
||||
t=clean_text(v)
|
||||
return ALIASES.get(t,ALIASES.get(t.lower(),t or "unknown"))
|
||||
|
||||
def entity_issues(name:object,entity_type:object="",role:str="endpoint")->List[str]:
|
||||
name=clean_text(name); et=normalize_entity_type(entity_type); low=name.lower(); out=[]
|
||||
if not name: return [f"{role}_missing"]
|
||||
if NUMERIC_ONLY_RE.match(name): out.append(f"{role}_numeric_only")
|
||||
if len(name)==1 and et not in {"allele","genetic_variant"}: out.append(f"{role}_single_character")
|
||||
if low in BAD_GENERIC: out.append(f"{role}_generic_fragment")
|
||||
if STARTING_FRAGMENT.search(name): out.append(f"{role}_starts_with_preposition")
|
||||
if URL_RE.search(name): out.append(f"{role}_url")
|
||||
if REF_RE.search(name): out.append(f"{role}_reference_fragment")
|
||||
if UNIT_RE.match(name): out.append(f"{role}_unit_only")
|
||||
if et and et!="unknown" and et not in ALLOWED_ENTITY_TYPES: out.append(f"{role}_unsupported_type:{et}")
|
||||
return out
|
||||
|
||||
def hard_reject_reasons(source,source_type,relation,target,target_type,evidence)->List[str]:
|
||||
out=[]
|
||||
if not clean_text(relation): out.append("relation_missing")
|
||||
if not clean_text(evidence): out.append("evidence_missing")
|
||||
out+=entity_issues(source,source_type,"source")+entity_issues(target,target_type,"target")
|
||||
return [x for x in out if "unsupported_type:" not in x]
|
||||
|
||||
def normalize_relation(raw_relation,relation_category="",target_type="")->Tuple[str,str]:
|
||||
raw=clean_text(raw_relation); safe=re.sub(r"\s+","_",raw.lower())
|
||||
if not safe: return "","empty_relation"
|
||||
tgt=normalize_entity_type(target_type); cat=clean_text(relation_category).lower()
|
||||
if cat=="classification" and safe in {"is","are","is_a","type_of"}:
|
||||
return "is_a","category-based classification"
|
||||
if safe in {"contains","harbors","possesses"} and tgt in {"motif","domain","cis_element","sequence_feature"}:
|
||||
return ({"cis_element":"contains_cis_element","domain":"contains_domain"}.get(tgt,"contains_motif"),f"contains + {tgt} target")
|
||||
for pat,norm,reason in REL_NORM:
|
||||
if pat.match(safe): return norm,reason
|
||||
if cat=="subcellular_localization" and safe in {"located","localized","localised"}: return "localizes_to","category-based localization"
|
||||
return safe,"kept_raw_relation"
|
||||
|
||||
def is_weak_relation(relation: object) -> bool:
|
||||
return clean_text(relation).lower().replace(" ", "_") in WEAK_RELATIONS
|
||||
|
||||
def is_administrative_relation(relation: object, evidence: object) -> bool:
|
||||
return bool(
|
||||
ADMIN_EVIDENCE_RE.search(clean_text(evidence))
|
||||
and clean_text(relation).lower().replace(" ", "_") in {
|
||||
"supported", "funded_by", "provided", "acknowledged"
|
||||
}
|
||||
)
|
||||
|
||||
def infer_layer(source_type,relation_category,layer_hint="",source_label="",claim_scope="",confidence=0.0,endpoint_problem=False)->str:
|
||||
hint=clean_text(layer_hint); label=clean_text(source_label).lower(); cat=clean_text(relation_category).lower(); scope=clean_text(claim_scope).lower(); st=normalize_entity_type(source_type)
|
||||
if endpoint_problem: return "candidate_review"
|
||||
if hint in QA_WEIGHTS and hint!="rejected": return hint
|
||||
if label=="figure_caption" or cat=="figure_caption_observation": return "caption_relations"
|
||||
if cat in {"method_tool","experimental_material"} or st in {"method","tool_or_database","experimental_material"}: return "method_relations"
|
||||
if cat=="dataset_accession" or st=="dataset_accession": return "dataset_metadata"
|
||||
if cat=="background_prior_work" or scope in {"background","prior_work","previous_study"}: return "background_relations"
|
||||
if float(confidence or 0)<0.58 or cat in {"other_supported",""}: return "candidate_review"
|
||||
if st not in MAIN_SOURCE_TYPES: return "candidate_review"
|
||||
return "main_kg"
|
||||
|
||||
def qa_weight(layer)->float:
|
||||
return QA_WEIGHTS.get(clean_text(layer),0.5)
|
||||
+205
@@ -0,0 +1,205 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
V14 LLM定义生成器
|
||||
基于检索到的文本上下文生成实体定义
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import Optional, Dict
|
||||
|
||||
try:
|
||||
from ...project_config import env
|
||||
except ImportError:
|
||||
import os
|
||||
|
||||
def env(name: str, default: str = "") -> str:
|
||||
return os.environ.get(name, default).strip()
|
||||
|
||||
|
||||
class LLMDefinitionGenerator:
|
||||
"""
|
||||
LLM定义生成器
|
||||
|
||||
功能:
|
||||
1. 从文本上下文提取实体定义
|
||||
2. 使用LLM生成生物学意义的定义
|
||||
3. 基于检索而非预训练知识
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None):
|
||||
self.api_key = api_key
|
||||
self._client = None
|
||||
|
||||
# 尝试初始化 DeepSeek OpenAI-compatible 客户端
|
||||
try:
|
||||
import openai
|
||||
effective_api_key = api_key or env("DEEPSEEK_API_KEY")
|
||||
base_url = env("DEEPSEEK_BASE_URL", "https://api.deepseek.com")
|
||||
timeout = float(env("V15_RELATION_LLM_TIMEOUT", "45") or 45)
|
||||
max_retries = int(env("V15_RELATION_LLM_MAX_RETRIES", "0") or 0)
|
||||
if effective_api_key:
|
||||
self._client = openai.OpenAI(
|
||||
api_key=effective_api_key,
|
||||
base_url=base_url,
|
||||
timeout=timeout,
|
||||
max_retries=max_retries,
|
||||
)
|
||||
else:
|
||||
# 尝试从配置文件读取
|
||||
self._client = self._load_api_config()
|
||||
except Exception as e:
|
||||
print(f"⚠️ LLM客户端初始化失败: {e}")
|
||||
|
||||
def _load_api_config(self):
|
||||
"""从.env文件加载API设置"""
|
||||
import openai
|
||||
api_key = env("DEEPSEEK_API_KEY")
|
||||
if api_key:
|
||||
print(" ✓ 从环境变量加载DEEPSEEK_API_KEY")
|
||||
base_url = env("DEEPSEEK_BASE_URL", "https://api.deepseek.com")
|
||||
timeout = float(env("V15_RELATION_LLM_TIMEOUT", "45") or 45)
|
||||
max_retries = int(env("V15_RELATION_LLM_MAX_RETRIES", "0") or 0)
|
||||
return openai.OpenAI(
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
timeout=timeout,
|
||||
max_retries=max_retries,
|
||||
)
|
||||
|
||||
print(f" ⚠️ .env文件未找到或无DEEPSEEK_API_KEY,使用模板回退")
|
||||
return None
|
||||
|
||||
def generate_definition(self, entity_name: str, entity_type: str,
|
||||
context: str, method: str = "hybrid") -> str:
|
||||
"""
|
||||
生成实体定义
|
||||
|
||||
Args:
|
||||
entity_name: 实体名
|
||||
entity_type: 实体类型
|
||||
context: 文本上下文
|
||||
method: 生成方法 - "extract"(提取), "llm"(LLM生成), "hybrid"(混合)
|
||||
|
||||
Returns:
|
||||
实体定义
|
||||
"""
|
||||
if method == "extract":
|
||||
return self._extract_definition(entity_name, context)
|
||||
elif method == "llm":
|
||||
return self._llm_generate_definition(entity_name, entity_type, context)
|
||||
else: # hybrid
|
||||
extracted = self._extract_definition(entity_name, context)
|
||||
if extracted and len(extracted) > 20:
|
||||
return extracted
|
||||
return self._llm_generate_definition(entity_name, entity_type, context)
|
||||
|
||||
def _extract_definition(self, entity_name: str, context: str) -> str:
|
||||
"""从上下文提取定义"""
|
||||
# 查找包含实体的句子
|
||||
sentences = re.split(r'[.!?]+', context)
|
||||
|
||||
for sent in sentences:
|
||||
if entity_name.lower() in sent.lower():
|
||||
# 清理句子
|
||||
sent = sent.strip()
|
||||
if len(sent) > 20 and len(sent) < 300:
|
||||
return sent
|
||||
|
||||
return ""
|
||||
|
||||
def _llm_generate_definition(self, entity_name: str, entity_type: str,
|
||||
context: str) -> str:
|
||||
"""使用LLM生成定义"""
|
||||
if not self._client:
|
||||
# 回退到模板
|
||||
return self._template_definition(entity_name, entity_type)
|
||||
|
||||
try:
|
||||
prompt = f"""Based on the following text, generate a concise biological definition for "{entity_name}" ({entity_type}).
|
||||
|
||||
Text context:
|
||||
{context[:1000]}
|
||||
|
||||
Requirements:
|
||||
1. Use ONLY information from the provided text
|
||||
2. Be specific and accurate
|
||||
3. Keep it under 100 words
|
||||
4. Focus on biological function and significance
|
||||
|
||||
Definition:"""
|
||||
|
||||
response = self._client.chat.completions.create(
|
||||
model=self._normalize_model_name(
|
||||
env("V15_RELATION_LLM_MODEL")
|
||||
or env("DEEPSEEK_MODEL")
|
||||
or "deepseek-v4-flash"
|
||||
),
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a precise biological knowledge extractor. Only use provided context."},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0.3,
|
||||
max_tokens=200
|
||||
)
|
||||
|
||||
definition = response.choices[0].message.content.strip()
|
||||
|
||||
# 验证定义质量
|
||||
if len(definition) < 10 or entity_name.lower() not in definition.lower():
|
||||
return self._template_definition(entity_name, entity_type)
|
||||
|
||||
return definition
|
||||
|
||||
except Exception as e:
|
||||
print(f"⚠️ LLM生成失败 ({entity_name}): {str(e)[:80]}")
|
||||
return self._template_definition(entity_name, entity_type)
|
||||
|
||||
def _normalize_model_name(self, model: str) -> str:
|
||||
"""Normalize legacy V15 definition model names to DeepSeek v4 flash."""
|
||||
normalized = (model or "").strip()
|
||||
if normalized in {"", "kimi-k2.5", "bailian/kimi-k2.5"}:
|
||||
return "deepseek-v4-flash"
|
||||
return normalized
|
||||
|
||||
def _template_definition(self, entity_name: str, entity_type: str) -> str:
|
||||
"""模板定义(回退)"""
|
||||
templates = {
|
||||
'gene': f"{entity_name} is a gene involved in biological processes.",
|
||||
'protein': f"{entity_name} is a protein encoded by a gene.",
|
||||
'tissue': f"{entity_name} is a plant tissue with specific functions.",
|
||||
'metabolite': f"{entity_name} is a metabolic compound.",
|
||||
'process': f"{entity_name} is a biological process.",
|
||||
'compartment': f"{entity_name} is a subcellular compartment.",
|
||||
}
|
||||
return templates.get(entity_type, f"{entity_name} is a biological entity.")
|
||||
|
||||
|
||||
# 全局单例
|
||||
_llm_generator: Optional[LLMDefinitionGenerator] = None
|
||||
|
||||
def get_llm_generator(api_key: Optional[str] = None) -> LLMDefinitionGenerator:
|
||||
"""获取LLM定义生成器单例"""
|
||||
global _llm_generator
|
||||
if _llm_generator is None:
|
||||
_llm_generator = LLMDefinitionGenerator(api_key)
|
||||
return _llm_generator
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 测试
|
||||
print("="*80)
|
||||
print("🧪 LLM定义生成器测试")
|
||||
print("="*80)
|
||||
|
||||
generator = LLMDefinitionGenerator()
|
||||
|
||||
# 测试提取
|
||||
context = "SUC5 encodes a sucrose/H+ symporter that facilitates sucrose transport in the phloem. It plays a crucial role in phloem loading."
|
||||
|
||||
print("\n📝 测试提取定义:")
|
||||
definition = generator.generate_definition("SUC5", "gene", context, method="extract")
|
||||
print(f" 结果: {definition}")
|
||||
|
||||
print("\n🤖 测试LLM定义 (如果有API):")
|
||||
definition = generator.generate_definition("SUC5", "gene", context, method="llm")
|
||||
print(f" 结果: {definition}")
|
||||
@@ -0,0 +1,514 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
V15 方法学参数审核模块
|
||||
整合 hybrid extraction 的审核机制到 V15
|
||||
"""
|
||||
|
||||
import re
|
||||
import json
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class ReviewSeverity(Enum):
|
||||
ERROR = "error" # 必须修正
|
||||
WARNING = "warning" # 建议确认
|
||||
INFO = "info" # 仅供参考
|
||||
|
||||
|
||||
@dataclass
|
||||
class MethodologyIssue:
|
||||
field: str
|
||||
severity: ReviewSeverity
|
||||
message: str
|
||||
current_value: any
|
||||
suggested_value: any
|
||||
source: str
|
||||
reason: str
|
||||
stage: str # identify/phylogeny/motif等
|
||||
|
||||
|
||||
class MethodologyReviewer:
|
||||
"""
|
||||
方法学参数审核器
|
||||
|
||||
审核维度:
|
||||
1. 值合理性(范围、格式、枚举)
|
||||
2. 正则 vs LLM 冲突
|
||||
3. 上下文一致性
|
||||
4. 缺失值与默认值
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# 参数验证规则
|
||||
self.validation_rules = {
|
||||
# 数值范围
|
||||
'bootstrap': {'type': 'int', 'min': 100, 'max': 10000, 'common': [500, 1000, 2000]},
|
||||
'e_value': {'type': 'float', 'format': r'\d+(\.\d+)?[eE][+-]?\d+', 'common': ['1e-5', '1e-10', '0.001']},
|
||||
'max_motifs': {'type': 'int', 'min': 1, 'max': 50, 'common': [10, 15, 20]},
|
||||
'motif_width_min': {'type': 'int', 'min': 3, 'max': 50},
|
||||
'motif_width_max': {'type': 'int', 'min': 6, 'max': 500},
|
||||
'identity': {'type': 'int', 'min': 30, 'max': 100, 'default': 40},
|
||||
'coverage': {'type': 'int', 'min': 30, 'max': 100, 'default': 70},
|
||||
|
||||
# 枚举值
|
||||
'tree_method': {'type': 'enum', 'values': ['NJ', 'ML', 'Bayesian', 'UPGMA', 'MP']},
|
||||
'substitution_model': {'type': 'enum', 'values': ['Poisson', 'JTT', 'WAG', 'LG', 'Dayhoff', 'JC']},
|
||||
'motif_mode': {'type': 'enum', 'values': ['optimum', 'zoops', 'anr']},
|
||||
'de_tool': {'type': 'enum', 'values': ['DESeq2', 'edgeR', 'limma']},
|
||||
'aligner': {'type': 'enum', 'values': ['HISAT2', 'STAR', 'TopHat', 'Bowtie2']},
|
||||
|
||||
# 格式验证
|
||||
'pfam_id': {'type': 'regex', 'pattern': r'^PF\d{5}$'},
|
||||
'geo_accession': {'type': 'regex', 'pattern': r'^GSE\d+$'},
|
||||
'sra_accession': {'type': 'regex', 'pattern': r'^PRJNA\d+$'},
|
||||
}
|
||||
|
||||
# 常见错误自动修正
|
||||
self.auto_corrections = {
|
||||
'tree_method': {
|
||||
'N-J': 'NJ',
|
||||
'n-j': 'NJ',
|
||||
'neighbor-joining': 'NJ',
|
||||
'Neighbor-Joining': 'NJ',
|
||||
'neighbor joining': 'NJ',
|
||||
'maximum likelihood': 'ML',
|
||||
'Maximum Likelihood': 'ML',
|
||||
},
|
||||
'substitution_model': {
|
||||
'lg': 'LG',
|
||||
'poisson': 'Poisson',
|
||||
'jtt': 'JTT',
|
||||
'wag': 'WAG',
|
||||
'dayhoff': 'Dayhoff',
|
||||
},
|
||||
'e_value': {
|
||||
'0.00001': '1e-5',
|
||||
'0.0000000001': '1e-10',
|
||||
'1e-05': '1e-5',
|
||||
'10e-6': '1e-5',
|
||||
}
|
||||
}
|
||||
|
||||
# 阶段间依赖关系
|
||||
self.stage_dependencies = {
|
||||
'phylogeny': ['identify'], # 建树需要鉴定出的基因
|
||||
'motif': ['identify'], # motif分析需要基因序列
|
||||
'structure': ['identify'], # 结构分析需要基因ID
|
||||
'promoter': ['identify'], # 启动子分析需要基因
|
||||
'synteny': ['identify'], # 共线性需要基因位置
|
||||
'expression': ['identify'], # 表达分析需要基因
|
||||
'rnaseq': ['identify'],
|
||||
}
|
||||
|
||||
def review(self, extraction_result: Dict) -> Dict:
|
||||
"""
|
||||
审核提取结果
|
||||
|
||||
Args:
|
||||
extraction_result: 混合提取器输出
|
||||
|
||||
Returns:
|
||||
{
|
||||
'issues': [MethodologyIssue, ...],
|
||||
'corrected': {...}, # 自动修正后的结果
|
||||
'requires_confirmation': [...], # 需要人工确认的字段
|
||||
'confidence_score': float,
|
||||
'stage_completeness': {...}, # 各阶段完整度
|
||||
'recommendations': [...] # 建议
|
||||
}
|
||||
"""
|
||||
issues = []
|
||||
corrected = {}
|
||||
requires_confirmation = []
|
||||
|
||||
final_results = extraction_result.get('final_results', {})
|
||||
|
||||
# 按阶段分组审核
|
||||
stage_params = self._group_by_stage(final_results)
|
||||
stage_completeness = {}
|
||||
|
||||
for stage, params in stage_params.items():
|
||||
stage_issues, stage_corrected, stage_needs_review = self._review_stage(stage, params)
|
||||
issues.extend(stage_issues)
|
||||
corrected.update(stage_corrected)
|
||||
requires_confirmation.extend(stage_needs_review)
|
||||
|
||||
# 计算阶段完整度
|
||||
stage_completeness[stage] = self._calculate_stage_completeness(params, stage_issues)
|
||||
|
||||
# 检查阶段依赖
|
||||
dependency_issues = self._check_stage_dependencies(stage_params)
|
||||
issues.extend(dependency_issues)
|
||||
|
||||
# 计算整体置信度
|
||||
confidence_score = self._calculate_confidence(issues, len(final_results))
|
||||
|
||||
# 生成建议
|
||||
recommendations = self._generate_recommendations(issues, stage_completeness)
|
||||
|
||||
return {
|
||||
'issues': [
|
||||
{
|
||||
'field': i.field,
|
||||
'severity': i.severity.value,
|
||||
'message': i.message,
|
||||
'current': i.current_value,
|
||||
'suggested': i.suggested_value,
|
||||
'source': i.source,
|
||||
'reason': i.reason,
|
||||
'stage': i.stage
|
||||
}
|
||||
for i in issues
|
||||
],
|
||||
'corrected': corrected,
|
||||
'requires_confirmation': requires_confirmation,
|
||||
'confidence_score': confidence_score,
|
||||
'stage_completeness': stage_completeness,
|
||||
'recommendations': recommendations,
|
||||
'review_summary': {
|
||||
'errors': len([i for i in issues if i.severity == ReviewSeverity.ERROR]),
|
||||
'warnings': len([i for i in issues if i.severity == ReviewSeverity.WARNING]),
|
||||
'infos': len([i for i in issues if i.severity == ReviewSeverity.INFO]),
|
||||
'total_fields': len(final_results),
|
||||
'needs_review': len(requires_confirmation),
|
||||
'stages_reviewed': len(stage_params)
|
||||
}
|
||||
}
|
||||
|
||||
def _group_by_stage(self, final_results: Dict) -> Dict[str, Dict]:
|
||||
"""按阶段分组参数"""
|
||||
stages = {}
|
||||
for field, data in final_results.items():
|
||||
# 从字段名解析阶段
|
||||
if '.' in field:
|
||||
stage = field.split('.')[0]
|
||||
else:
|
||||
# 根据字段名推断阶段
|
||||
stage = self._infer_stage(field)
|
||||
|
||||
if stage not in stages:
|
||||
stages[stage] = {}
|
||||
stages[stage][field] = data
|
||||
return stages
|
||||
|
||||
def _infer_stage(self, field: str) -> str:
|
||||
"""从字段名推断阶段"""
|
||||
stage_keywords = {
|
||||
'identify': ['pfam', 'hmm', 'blast', 'e_value', 'identity', 'coverage'],
|
||||
'phylogeny': ['bootstrap', 'tree_method', 'substitution', 'alignment'],
|
||||
'motif': ['motif', 'meme', 'width', 'max_motifs'],
|
||||
'structure': ['gsds', 'intron', 'exon', 'structure'],
|
||||
'promoter': ['plantcare', 'promoter', 'upstream'],
|
||||
'synteny': ['mcscanx', 'synteny', 'collinear'],
|
||||
'expression': ['expression', 'geo', 'fpkm', 'tissue'],
|
||||
'rnaseq': ['rnaseq', 'hisat', 'star', 'deseq'],
|
||||
}
|
||||
field_lower = field.lower()
|
||||
for stage, keywords in stage_keywords.items():
|
||||
if any(kw in field_lower for kw in keywords):
|
||||
return stage
|
||||
return 'general'
|
||||
|
||||
def _review_stage(self, stage: str, params: Dict) -> Tuple[List[MethodologyIssue], Dict, List[str]]:
|
||||
"""审核单个阶段的参数"""
|
||||
issues = []
|
||||
corrected = {}
|
||||
requires_confirmation = []
|
||||
|
||||
for field, data in params.items():
|
||||
value = data.get('value')
|
||||
confidence = data.get('confidence')
|
||||
source = data.get('source', '')
|
||||
llm_alt = data.get('llm_alternative')
|
||||
|
||||
# 1. 检查正则 vs LLM 冲突
|
||||
if 'regex' in source and llm_alt and value != llm_alt:
|
||||
resolved = self._resolve_conflict(field, value, llm_alt)
|
||||
if resolved != value:
|
||||
issues.append(MethodologyIssue(
|
||||
field=field,
|
||||
severity=ReviewSeverity.WARNING,
|
||||
message=f"正则与LLM结果冲突",
|
||||
current_value=value,
|
||||
suggested_value=resolved,
|
||||
source="conflict_resolution",
|
||||
reason=f"正则提取'{value}',LLM提取'{llm_alt}',根据规则选择'{resolved}'",
|
||||
stage=stage
|
||||
))
|
||||
requires_confirmation.append(field)
|
||||
value = resolved
|
||||
|
||||
# 2. 自动修正常见错误
|
||||
corrected_value = self._auto_correct(field, value)
|
||||
if corrected_value != value:
|
||||
issues.append(MethodologyIssue(
|
||||
field=field,
|
||||
severity=ReviewSeverity.INFO,
|
||||
message=f"自动修正格式",
|
||||
current_value=value,
|
||||
suggested_value=corrected_value,
|
||||
source="auto_correction",
|
||||
reason=f"'{value}' -> '{corrected_value}'",
|
||||
stage=stage
|
||||
))
|
||||
value = corrected_value
|
||||
|
||||
# 3. 验证值合理性
|
||||
validation = self._validate_value(field, value)
|
||||
if not validation['valid']:
|
||||
issues.append(MethodologyIssue(
|
||||
field=field,
|
||||
severity=validation['severity'],
|
||||
message=validation['message'],
|
||||
current_value=value,
|
||||
suggested_value=validation['suggested'],
|
||||
source="validation",
|
||||
reason=validation['reason'],
|
||||
stage=stage
|
||||
))
|
||||
if validation['severity'] == ReviewSeverity.ERROR:
|
||||
value = validation['suggested']
|
||||
elif validation['severity'] == ReviewSeverity.WARNING:
|
||||
requires_confirmation.append(field)
|
||||
|
||||
# 4. 缺失值处理
|
||||
if value is None or value == '':
|
||||
default = self._get_default_value(field)
|
||||
if default is not None:
|
||||
issues.append(MethodologyIssue(
|
||||
field=field,
|
||||
severity=ReviewSeverity.WARNING,
|
||||
message=f"参数缺失,建议使用默认值",
|
||||
current_value=None,
|
||||
suggested_value=default,
|
||||
source="default_suggestion",
|
||||
reason=f"原文未明确,领域默认值为{default}",
|
||||
stage=stage
|
||||
))
|
||||
requires_confirmation.append(field)
|
||||
|
||||
corrected[field] = {
|
||||
'value': value,
|
||||
'original': data.get('value'),
|
||||
'confidence': confidence,
|
||||
'source': source
|
||||
}
|
||||
|
||||
return issues, corrected, requires_confirmation
|
||||
|
||||
def _resolve_conflict(self, field: str, regex_val, llm_val) -> any:
|
||||
"""解决正则与LLM的冲突"""
|
||||
# 检查哪个值符合枚举
|
||||
if field in self.validation_rules:
|
||||
rule = self.validation_rules[field]
|
||||
if rule['type'] == 'enum':
|
||||
valid_values = [v.lower() for v in rule['values']]
|
||||
regex_lower = str(regex_val).lower()
|
||||
llm_lower = str(llm_val).lower()
|
||||
|
||||
regex_valid = regex_lower in valid_values
|
||||
llm_valid = llm_lower in valid_values
|
||||
|
||||
if regex_valid and not llm_valid:
|
||||
return regex_val
|
||||
elif llm_valid and not regex_valid:
|
||||
return llm_val
|
||||
elif regex_valid and llm_valid:
|
||||
# 都有效,优先LLM(通常更准确)
|
||||
return llm_val
|
||||
|
||||
# 默认优先LLM
|
||||
return llm_val
|
||||
|
||||
def _auto_correct(self, field: str, value: any) -> any:
|
||||
"""自动修正常见错误"""
|
||||
if field in self.auto_corrections and value in self.auto_corrections[field]:
|
||||
return self.auto_corrections[field][value]
|
||||
return value
|
||||
|
||||
def _validate_value(self, field: str, value: any) -> Dict:
|
||||
"""验证值的合理性"""
|
||||
if field not in self.validation_rules:
|
||||
return {'valid': True}
|
||||
|
||||
rule = self.validation_rules[field]
|
||||
|
||||
if value is None:
|
||||
return {'valid': True} # 缺失值单独处理
|
||||
|
||||
# 数值范围验证
|
||||
if rule['type'] == 'int':
|
||||
try:
|
||||
val = int(value)
|
||||
if val < rule['min'] or val > rule['max']:
|
||||
return {
|
||||
'valid': False,
|
||||
'severity': ReviewSeverity.ERROR,
|
||||
'message': f"值{val}超出范围[{rule['min']}, {rule['max']}]",
|
||||
'suggested': rule.get('default', rule['min']),
|
||||
'reason': f"数值超出合理范围"
|
||||
}
|
||||
except:
|
||||
return {
|
||||
'valid': False,
|
||||
'severity': ReviewSeverity.ERROR,
|
||||
'message': f"'{value}'不是有效整数",
|
||||
'suggested': rule.get('default', rule['min']),
|
||||
'reason': "类型错误"
|
||||
}
|
||||
|
||||
# 枚举值验证
|
||||
if rule['type'] == 'enum':
|
||||
if str(value).lower() not in [v.lower() for v in rule['values']]:
|
||||
return {
|
||||
'valid': False,
|
||||
'severity': ReviewSeverity.WARNING,
|
||||
'message': f"'{value}'不是标准枚举值",
|
||||
'suggested': rule['values'][0],
|
||||
'reason': f"有效值为: {', '.join(rule['values'])}"
|
||||
}
|
||||
|
||||
# 正则验证
|
||||
if rule['type'] == 'regex':
|
||||
if not re.match(rule['pattern'], str(value)):
|
||||
return {
|
||||
'valid': False,
|
||||
'severity': ReviewSeverity.ERROR,
|
||||
'message': f"'{value}'格式不符合{field}标准",
|
||||
'suggested': None,
|
||||
'reason': f"应匹配模式: {rule['pattern']}"
|
||||
}
|
||||
|
||||
return {'valid': True}
|
||||
|
||||
def _get_default_value(self, field: str) -> Optional[any]:
|
||||
"""获取字段默认值"""
|
||||
defaults = {
|
||||
'e_value': '1e-5',
|
||||
'identity': 40,
|
||||
'coverage': 70,
|
||||
'bootstrap': 1000,
|
||||
'tree_method': 'NJ',
|
||||
'substitution_model': 'Poisson',
|
||||
'max_motifs': 10,
|
||||
'motif_mode': 'optimum',
|
||||
'de_tool': 'DESeq2',
|
||||
'aligner': 'HISAT2',
|
||||
}
|
||||
return defaults.get(field)
|
||||
|
||||
def _calculate_stage_completeness(self, params: Dict, issues: List[MethodologyIssue]) -> Dict:
|
||||
"""计算阶段完整度"""
|
||||
total = len(params)
|
||||
if total == 0:
|
||||
return {'score': 0, 'critical_missing': []}
|
||||
|
||||
# 关键参数
|
||||
critical_params = ['e_value', 'bootstrap', 'tree_method', 'pfam_id']
|
||||
critical_missing = [p for p in critical_params if p not in params or params[p].get('value') is None]
|
||||
|
||||
# 计算得分
|
||||
error_count = len([i for i in issues if i.severity == ReviewSeverity.ERROR])
|
||||
warning_count = len([i for i in issues if i.severity == ReviewSeverity.WARNING])
|
||||
|
||||
score = 100 - (error_count * 20) - (warning_count * 10) - (len(critical_missing) * 15)
|
||||
score = max(0, score)
|
||||
|
||||
return {
|
||||
'score': score,
|
||||
'total_params': total,
|
||||
'error_count': error_count,
|
||||
'warning_count': warning_count,
|
||||
'critical_missing': critical_missing,
|
||||
'grade': 'A' if score >= 90 else 'B' if score >= 75 else 'C' if score >= 60 else 'D'
|
||||
}
|
||||
|
||||
def _check_stage_dependencies(self, stage_params: Dict) -> List[MethodologyIssue]:
|
||||
"""检查阶段依赖关系"""
|
||||
issues = []
|
||||
|
||||
for stage, deps in self.stage_dependencies.items():
|
||||
if stage in stage_params:
|
||||
for dep in deps:
|
||||
if dep not in stage_params:
|
||||
issues.append(MethodologyIssue(
|
||||
field=f"{stage}_dependency",
|
||||
severity=ReviewSeverity.WARNING,
|
||||
message=f"{stage}阶段依赖的{dep}阶段未配置",
|
||||
current_value=None,
|
||||
suggested_value=dep,
|
||||
source="dependency_check",
|
||||
reason=f"执行{stage}前需要先完成{dep}",
|
||||
stage=stage
|
||||
))
|
||||
|
||||
return issues
|
||||
|
||||
def _calculate_confidence(self, issues: List[MethodologyIssue], total_fields: int) -> float:
|
||||
"""计算整体置信度"""
|
||||
if total_fields == 0:
|
||||
return 0.0
|
||||
|
||||
error_weight = 20
|
||||
warning_weight = 5
|
||||
info_weight = 1
|
||||
|
||||
errors = len([i for i in issues if i.severity == ReviewSeverity.ERROR])
|
||||
warnings = len([i for i in issues if i.severity == ReviewSeverity.WARNING])
|
||||
infos = len([i for i in issues if i.severity == ReviewSeverity.INFO])
|
||||
|
||||
penalty = (errors * error_weight) + (warnings * warning_weight) + (infos * info_weight)
|
||||
score = max(0, 100 - (penalty / total_fields * 100))
|
||||
|
||||
return round(score, 1)
|
||||
|
||||
def _generate_recommendations(self, issues: List[MethodologyIssue], stage_completeness: Dict) -> List[str]:
|
||||
"""生成建议"""
|
||||
recommendations = []
|
||||
|
||||
# 基于问题生成建议
|
||||
error_fields = [i.field for i in issues if i.severity == ReviewSeverity.ERROR]
|
||||
warning_fields = [i.field for i in issues if i.severity == ReviewSeverity.WARNING]
|
||||
|
||||
if error_fields:
|
||||
recommendations.append(f"修正错误参数: {', '.join(error_fields[:3])}")
|
||||
|
||||
if warning_fields:
|
||||
recommendations.append(f"确认警告参数: {', '.join(warning_fields[:3])}")
|
||||
|
||||
# 基于阶段完整度生成建议
|
||||
low_score_stages = [s for s, c in stage_completeness.items() if c.get('score', 0) < 60]
|
||||
if low_score_stages:
|
||||
recommendations.append(f"完善阶段配置: {', '.join(low_score_stages)}")
|
||||
|
||||
# 通用建议
|
||||
if not recommendations:
|
||||
recommendations.append("参数配置完整,可直接执行")
|
||||
|
||||
return recommendations
|
||||
|
||||
|
||||
# 便捷函数
|
||||
def review_methodology(extraction_result: Dict) -> Dict:
|
||||
"""审核方法学参数"""
|
||||
reviewer = MethodologyReviewer()
|
||||
return reviewer.review(extraction_result)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# 测试
|
||||
test_result = {
|
||||
'final_results': {
|
||||
'identify.pfam_id': {'value': 'PF03110', 'confidence': 'high', 'source': 'regex'},
|
||||
'identify.e_value': {'value': None, 'confidence': 'missing', 'source': 'none'},
|
||||
'phylogeny.bootstrap': {'value': '1000', 'confidence': 'high', 'source': 'regex'},
|
||||
'phylogeny.tree_method': {'value': 'N-J', 'confidence': 'medium', 'source': 'regex', 'llm_alternative': 'NJ'},
|
||||
'phylogeny.substitution_model': {'value': 'lg', 'confidence': 'medium', 'source': 'regex', 'llm_alternative': 'Poisson'},
|
||||
'motif.max_motifs': {'value': '10', 'confidence': 'high', 'source': 'llm_only'},
|
||||
}
|
||||
}
|
||||
|
||||
result = review_methodology(test_result)
|
||||
print(json.dumps(result, indent=2, ensure_ascii=False))
|
||||
@@ -0,0 +1,478 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
V14 Neo4j 知识图谱存储
|
||||
将V14提取结果持久化到Neo4j
|
||||
"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Optional, Any
|
||||
from dataclasses import dataclass, asdict, field
|
||||
from datetime import datetime
|
||||
from contextlib import contextmanager
|
||||
|
||||
try:
|
||||
from neo4j import GraphDatabase, Driver
|
||||
NEO4J_AVAILABLE = True
|
||||
except ImportError:
|
||||
NEO4J_AVAILABLE = False
|
||||
GraphDatabase = None
|
||||
Driver = None
|
||||
|
||||
# 导入V14类型
|
||||
try:
|
||||
from .phase3_entity_recognition import Entity
|
||||
from .phase4_relation_extraction import Relation
|
||||
except ImportError:
|
||||
from phase3_entity_recognition import Entity
|
||||
from phase4_relation_extraction import Relation
|
||||
|
||||
try:
|
||||
from ...project_config import env
|
||||
except ImportError:
|
||||
import os
|
||||
|
||||
def env(name: str, default: str = "") -> str:
|
||||
return os.environ.get(name, default).strip()
|
||||
|
||||
|
||||
@dataclass
|
||||
class Neo4jConfig:
|
||||
"""Neo4j配置"""
|
||||
uri: str = field(default_factory=lambda: env("NEO4J_URI", "bolt://localhost:7687"))
|
||||
user: str = field(default_factory=lambda: env("NEO4J_USER", "neo4j"))
|
||||
password: str = field(default_factory=lambda: env("NEO4J_PASSWORD", ""))
|
||||
database: str = field(default_factory=lambda: env("NEO4J_DATABASE", "neo4j"))
|
||||
|
||||
def to_dict(self) -> Dict:
|
||||
return {
|
||||
'uri': self.uri,
|
||||
'user': self.user,
|
||||
'password': self.password,
|
||||
'database': self.database
|
||||
}
|
||||
|
||||
|
||||
class V15Neo4jStore:
|
||||
"""
|
||||
V14 Neo4j 存储
|
||||
|
||||
功能:
|
||||
1. 实体/关系的CRUD
|
||||
2. 支持V14完整21列格式
|
||||
3. 知识图谱查询
|
||||
4. 与V14提取器无缝集成
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[Neo4jConfig] = None):
|
||||
if not NEO4J_AVAILABLE:
|
||||
raise ImportError("neo4j驱动未安装,请运行: pip install neo4j")
|
||||
|
||||
self.config = config or Neo4jConfig()
|
||||
self._driver: Optional[Driver] = None
|
||||
self._connect()
|
||||
self._init_schema()
|
||||
|
||||
def _connect(self):
|
||||
"""连接Neo4j数据库"""
|
||||
try:
|
||||
self._driver = GraphDatabase.driver(
|
||||
self.config.uri,
|
||||
auth=(self.config.user, self.config.password)
|
||||
)
|
||||
# 测试连接
|
||||
with self._driver.session(database=self.config.database) as session:
|
||||
session.run("RETURN 1")
|
||||
print(f"[V15-Neo4j] 成功连接到 {self.config.uri}")
|
||||
except Exception as e:
|
||||
print(f"[V15-Neo4j] 连接失败: {e}")
|
||||
raise
|
||||
|
||||
def _init_schema(self):
|
||||
"""初始化Neo4j schema (约束和索引)"""
|
||||
with self._get_session() as session:
|
||||
# 创建实体ID唯一约束
|
||||
try:
|
||||
session.run("""
|
||||
CREATE CONSTRAINT entity_id_constraint IF NOT EXISTS
|
||||
FOR (e:Entity) REQUIRE e.entity_id IS UNIQUE
|
||||
""")
|
||||
except Exception as e:
|
||||
print(f"[V15-Neo4j] 约束创建失败(可能已存在): {e}")
|
||||
|
||||
# 创建索引
|
||||
try:
|
||||
session.run("""
|
||||
CREATE INDEX entity_type_index IF NOT EXISTS
|
||||
FOR (e:Entity) ON (e.entity_type)
|
||||
""")
|
||||
session.run("""
|
||||
CREATE INDEX entity_name_index IF NOT EXISTS
|
||||
FOR (e:Entity) ON (e.name)
|
||||
""")
|
||||
session.run("""
|
||||
CREATE INDEX relation_type_index IF NOT EXISTS
|
||||
FOR ()-[r:RELATION]-() ON (r.relation_type)
|
||||
""")
|
||||
print("[V15-Neo4j] 索引创建完成")
|
||||
except Exception as e:
|
||||
print(f"[V15-Neo4j] 索引创建失败: {e}")
|
||||
|
||||
@contextmanager
|
||||
def _get_session(self):
|
||||
"""获取Neo4j会话"""
|
||||
if self._driver is None:
|
||||
raise RuntimeError("Neo4j未连接")
|
||||
session = self._driver.session(database=self.config.database)
|
||||
try:
|
||||
yield session
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
def close(self):
|
||||
"""关闭连接"""
|
||||
if self._driver:
|
||||
self._driver.close()
|
||||
self._driver = None
|
||||
print("[V15-Neo4j] 连接已关闭")
|
||||
|
||||
# ========== 实体操作 ==========
|
||||
|
||||
def save_entity(self, entity: Entity) -> bool:
|
||||
"""保存V14实体到Neo4j"""
|
||||
with self._get_session() as session:
|
||||
try:
|
||||
# 准备属性
|
||||
props = {
|
||||
'entity_id': entity.id,
|
||||
'name': entity.name,
|
||||
'entity_type': entity.type,
|
||||
'name_resolved': getattr(entity, 'name_resolved', entity.name),
|
||||
'aliases': json.dumps(getattr(entity, 'aliases', [])),
|
||||
'confidence': getattr(entity, 'confidence', 1.0),
|
||||
'source': getattr(entity, 'source', ''),
|
||||
}
|
||||
|
||||
# 添加properties中的额外属性
|
||||
if hasattr(entity, 'properties') and entity.properties:
|
||||
for k, v in entity.properties.items():
|
||||
if k not in props:
|
||||
props[f'prop_{k}'] = json.dumps(v) if isinstance(v, (dict, list)) else str(v)
|
||||
|
||||
# 创建实体节点
|
||||
session.run(f"""
|
||||
MERGE (e:Entity {{entity_id: $entity_id}})
|
||||
SET e += $props
|
||||
""", entity_id=entity.id, props=props)
|
||||
|
||||
# 添加类型标签
|
||||
session.run(f"""
|
||||
MATCH (e:Entity {{entity_id: $entity_id}})
|
||||
SET e:{entity.type}
|
||||
""", entity_id=entity.id)
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"[V15-Neo4j] 保存实体失败 {entity.name}: {e}")
|
||||
return False
|
||||
|
||||
def save_entities(self, entities: List[Entity]) -> Dict[str, int]:
|
||||
"""批量保存实体"""
|
||||
success = 0
|
||||
failed = 0
|
||||
|
||||
for entity in entities:
|
||||
if self.save_entity(entity):
|
||||
success += 1
|
||||
else:
|
||||
failed += 1
|
||||
|
||||
print(f"[V15-Neo4j] 实体保存: {success}成功, {failed}失败")
|
||||
return {'success': success, 'failed': failed}
|
||||
|
||||
# ========== 关系操作 ==========
|
||||
|
||||
def save_relation(self, relation: Relation) -> bool:
|
||||
"""保存V14关系到Neo4j (完整21列)"""
|
||||
with self._get_session() as session:
|
||||
try:
|
||||
# 确保源和目标实体存在 (使用entity_id作为主键)
|
||||
source_id = getattr(relation, 'source_resolved', relation.source) or relation.source
|
||||
target_id = getattr(relation, 'target_resolved', relation.target) or relation.target
|
||||
|
||||
# V14.2修复: 使用MERGE和entity_id确保实体存在,避免约束冲突
|
||||
session.run("""
|
||||
MERGE (s:Entity {entity_id: $source_id})
|
||||
ON CREATE SET s.name = $source_name, s.type = $source_type
|
||||
ON MATCH SET s.name = $source_name
|
||||
MERGE (t:Entity {entity_id: $target_id})
|
||||
ON CREATE SET t.name = $target_name, t.type = $target_type
|
||||
ON MATCH SET t.name = $target_name
|
||||
""",
|
||||
source_id=source_id,
|
||||
source_name=relation.source,
|
||||
source_type=relation.source_type,
|
||||
target_id=target_id,
|
||||
target_name=relation.target,
|
||||
target_type=relation.target_type
|
||||
)
|
||||
|
||||
# 准备关系属性 (完整21列)
|
||||
props = {
|
||||
# 原始字段
|
||||
'source': relation.source,
|
||||
'source_type': relation.source_type,
|
||||
'relation': relation.relation,
|
||||
'target': relation.target,
|
||||
'target_type': relation.target_type,
|
||||
|
||||
# 元数据
|
||||
'gene': getattr(relation, 'gene', ''),
|
||||
'pubmed_id': getattr(relation, 'pubmed_id', ''),
|
||||
'p_source': getattr(relation, 'p_source', ''),
|
||||
'species': getattr(relation, 'species', ''),
|
||||
'basis': getattr(relation, 'basis', ''),
|
||||
|
||||
# 定义
|
||||
'source_extracted_definition': getattr(relation, 'source_extracted_definition', ''),
|
||||
'source_generated_definition': getattr(relation, 'source_generated_definition', ''),
|
||||
'target_extracted_definition': getattr(relation, 'target_extracted_definition', ''),
|
||||
'target_generated_definition': getattr(relation, 'target_generated_definition', ''),
|
||||
|
||||
# 消歧层
|
||||
'source_type_resolved': getattr(relation, 'source_type_resolved', ''),
|
||||
'target_type_resolved': getattr(relation, 'target_type_resolved', ''),
|
||||
'source_resolved': getattr(relation, 'source_resolved', ''),
|
||||
'target_resolved': getattr(relation, 'target_resolved', ''),
|
||||
'relationship_resolved': getattr(relation, 'relationship_resolved', ''),
|
||||
|
||||
# 别名
|
||||
'source_gene_alias': getattr(relation, 'source_gene_alias', ''),
|
||||
'target_gene_alias': getattr(relation, 'target_gene_alias', ''),
|
||||
|
||||
# 其他
|
||||
'evidence': getattr(relation, 'evidence', ''),
|
||||
'section': getattr(relation, 'section', ''),
|
||||
'confidence': getattr(relation, 'confidence', 0.9),
|
||||
'method': getattr(relation, 'method', 'rule'),
|
||||
}
|
||||
|
||||
# 创建关系
|
||||
rel_type = getattr(relation, 'relationship_resolved', relation.relation) or relation.relation
|
||||
rel_type = rel_type.replace(' ', '_').replace('-', '_').upper()
|
||||
|
||||
session.run(f"""
|
||||
MATCH (s:Entity {{entity_id: $source_id}})
|
||||
MATCH (t:Entity {{entity_id: $target_id}})
|
||||
MERGE (s)-[r:RELATION {{relation: $relation}}]->(t)
|
||||
SET r += $props
|
||||
""",
|
||||
source_id=source_id,
|
||||
target_id=target_id,
|
||||
relation=relation.relation,
|
||||
props=props
|
||||
)
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"[V15-Neo4j] 保存关系失败 {relation.source}-{relation.relation}-{relation.target}: {e}")
|
||||
return False
|
||||
|
||||
def save_relations(self, relations: List[Relation]) -> Dict[str, int]:
|
||||
"""批量保存关系"""
|
||||
success = 0
|
||||
failed = 0
|
||||
|
||||
for relation in relations:
|
||||
if self.save_relation(relation):
|
||||
success += 1
|
||||
else:
|
||||
failed += 1
|
||||
|
||||
print(f"[V15-Neo4j] 关系保存: {success}成功, {failed}失败")
|
||||
return {'success': success, 'failed': failed}
|
||||
|
||||
# ========== 批量导入 ==========
|
||||
|
||||
def save_v14_result(self, entities: List[Entity], relations: List[Relation],
|
||||
clear_existing: bool = False) -> Dict:
|
||||
"""保存完整的V14提取结果"""
|
||||
print("\n" + "="*80)
|
||||
print("🧬 V14 → Neo4j 批量导入")
|
||||
print("="*80)
|
||||
|
||||
# 可选:清空现有数据
|
||||
if clear_existing:
|
||||
with self._get_session() as session:
|
||||
session.run("MATCH (n) DETACH DELETE n")
|
||||
print("[V15-Neo4j] 已清空现有数据")
|
||||
|
||||
# 导入实体
|
||||
print(f"\n📦 导入实体: {len(entities)} 个")
|
||||
entity_stats = self.save_entities(entities)
|
||||
|
||||
# 导入关系
|
||||
print(f"\n🔗 导入关系: {len(relations)} 个")
|
||||
relation_stats = self.save_relations(relations)
|
||||
|
||||
# 创建索引
|
||||
self._create_indexes()
|
||||
|
||||
print("\n" + "="*80)
|
||||
print("✅ V14导入完成!")
|
||||
print("="*80)
|
||||
print(f" 实体: {entity_stats['success']} 成功")
|
||||
print(f" 关系: {relation_stats['success']} 成功")
|
||||
|
||||
return {
|
||||
'entities': entity_stats,
|
||||
'relations': relation_stats
|
||||
}
|
||||
|
||||
def _create_indexes(self):
|
||||
"""创建额外索引"""
|
||||
with self._get_session() as session:
|
||||
try:
|
||||
# 为常用查询字段创建索引
|
||||
session.run("""
|
||||
CREATE INDEX entity_name_resolved_index IF NOT EXISTS
|
||||
FOR (e:Entity) ON (e.name_resolved)
|
||||
""")
|
||||
session.run("""
|
||||
CREATE INDEX relation_pubmed_index IF NOT EXISTS
|
||||
FOR ()-[r:RELATION]-() ON (r.pubmed_id)
|
||||
""")
|
||||
print("[V15-Neo4j] 额外索引创建完成")
|
||||
except Exception as e:
|
||||
print(f"[V15-Neo4j] 索引创建失败: {e}")
|
||||
|
||||
# ========== 查询功能 ==========
|
||||
|
||||
def query_entity(self, name: str) -> Optional[Dict]:
|
||||
"""查询实体"""
|
||||
with self._get_session() as session:
|
||||
result = session.run("""
|
||||
MATCH (e:Entity)
|
||||
WHERE e.name = $name OR e.name_resolved = $name
|
||||
RETURN e
|
||||
""", name=name)
|
||||
|
||||
record = result.single()
|
||||
if record:
|
||||
node = record["e"]
|
||||
return dict(node)
|
||||
return None
|
||||
|
||||
def query_relations_by_gene(self, gene_name: str) -> List[Dict]:
|
||||
"""查询基因相关的所有关系"""
|
||||
with self._get_session() as session:
|
||||
result = session.run("""
|
||||
MATCH (s:Entity)-[r:RELATION]->(t:Entity)
|
||||
WHERE s.name = $gene_name OR s.name_resolved = $gene_name
|
||||
OR t.name = $gene_name OR t.name_resolved = $gene_name
|
||||
OR r.gene = $gene_name
|
||||
RETURN s.name as source, r.relation as relation, t.name as target,
|
||||
r.source_resolved as source_resolved,
|
||||
r.target_resolved as target_resolved,
|
||||
r.relationship_resolved as relationship_resolved
|
||||
""", gene_name=gene_name)
|
||||
|
||||
return [dict(record) for record in result]
|
||||
|
||||
def query_by_pubmed(self, pubmed_id: str) -> Dict[str, List]:
|
||||
"""按PubMed ID查询"""
|
||||
with self._get_session() as session:
|
||||
# 查询实体
|
||||
entity_result = session.run("""
|
||||
MATCH (e:Entity)
|
||||
WHERE e.pubmed_id = $pubmed_id
|
||||
RETURN e
|
||||
""", pubmed_id=pubmed_id)
|
||||
|
||||
entities = [dict(record["e"]) for record in entity_result]
|
||||
|
||||
# 查询关系
|
||||
relation_result = session.run("""
|
||||
MATCH (s:Entity)-[r:RELATION]->(t:Entity)
|
||||
WHERE r.pubmed_id = $pubmed_id
|
||||
RETURN s.name as source, r.relation as relation, t.name as target, r
|
||||
""", pubmed_id=pubmed_id)
|
||||
|
||||
relations = [
|
||||
{**dict(record["r"]),
|
||||
'source': record['source'],
|
||||
'target': record['target']}
|
||||
for record in relation_result
|
||||
]
|
||||
|
||||
return {
|
||||
'entities': entities,
|
||||
'relations': relations
|
||||
}
|
||||
|
||||
def export_to_csv(self, output_path: Path):
|
||||
"""导出到CSV (兼容PC格式)"""
|
||||
import csv
|
||||
|
||||
with self._get_session() as session:
|
||||
result = session.run("""
|
||||
MATCH (s:Entity)-[r:RELATION]->(t:Entity)
|
||||
RETURN s.name as source, s.entity_type as source_type,
|
||||
r.relation as relationship, t.name as target, t.entity_type as target_type,
|
||||
r.gene as gene, r.pubmed_id as pubmedID, r.p_source as p_source,
|
||||
r.species as species, r.basis as basis,
|
||||
r.source_extracted_definition as source_extracted_definition,
|
||||
r.source_generated_definition as source_generated_definition,
|
||||
r.target_extracted_definition as target_extracted_definition,
|
||||
r.target_generated_definition as target_generated_definition,
|
||||
r.source_type_resolved as source_type_resolved,
|
||||
r.target_type_resolved as target_type_resolved,
|
||||
r.source_resolved as source_resolved,
|
||||
r.target_resolved as target_resolved,
|
||||
r.relationship_resolved as relationship_resolved,
|
||||
r.source_gene_alias as source_gene_alias,
|
||||
r.target_gene_alias as target_gene_alias
|
||||
""")
|
||||
|
||||
# 21列PC格式
|
||||
headers = [
|
||||
'source', 'source_type', 'relationship', 'target', 'target_type',
|
||||
'gene', 'pubmedID', 'p_source', 'species', 'basis',
|
||||
'source_extracted_definition', 'source_generated_definition',
|
||||
'target_extracted_definition', 'target_generated_definition',
|
||||
'source_type_resolved', 'target_type_resolved',
|
||||
'source_resolved', 'target_resolved', 'relationship_resolved',
|
||||
'source_gene_alias', 'target_gene_alias'
|
||||
]
|
||||
|
||||
with open(output_path, 'w', newline='', encoding='utf-8') as f:
|
||||
writer = csv.writer(f)
|
||||
writer.writerow(headers)
|
||||
|
||||
for record in result:
|
||||
writer.writerow([
|
||||
record['source'], record['source_type'], record['relationship'],
|
||||
record['target'], record['target_type'], record['gene'],
|
||||
record['pubmedID'], record['p_source'], record['species'],
|
||||
record['basis'], record['source_extracted_definition'],
|
||||
record['source_generated_definition'], record['target_extracted_definition'],
|
||||
record['target_generated_definition'], record['source_type_resolved'],
|
||||
record['target_type_resolved'], record['source_resolved'],
|
||||
record['target_resolved'], record['relationship_resolved'],
|
||||
record['source_gene_alias'], record['target_gene_alias']
|
||||
])
|
||||
|
||||
print(f"[V15-Neo4j] CSV导出完成: {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 测试
|
||||
print("V14 Neo4j Store 测试")
|
||||
|
||||
try:
|
||||
store = V15Neo4jStore()
|
||||
print("✅ 连接成功")
|
||||
store.close()
|
||||
except Exception as e:
|
||||
print(f"❌ 连接失败: {e}")
|
||||
@@ -0,0 +1,385 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, Optional
|
||||
from urllib.parse import unquote
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
EUROPE_PMC_SEARCH = "https://www.ebi.ac.uk/europepmc/webservices/rest/search"
|
||||
NCBI_ID_CONVERTER = "https://www.ncbi.nlm.nih.gov/pmc/utils/idconv/v1.0/"
|
||||
CROSSREF_WORKS = "https://api.crossref.org/works"
|
||||
NCBI_TOOL_NAME = "plantai-rag"
|
||||
USER_AGENT = "plantai-rag-v16/16.1.2 (paper metadata enrichment)"
|
||||
|
||||
|
||||
def normalize_doi(value: Optional[str]) -> str:
|
||||
text = unquote(str(value or "")).strip()
|
||||
text = re.sub(r"^(?:https?://(?:dx\.)?doi\.org/|doi\s*:\s*)", "", text, flags=re.I)
|
||||
text = text.strip().strip("()[]{}<>.,;")
|
||||
match = re.search(r"10\.\d{4,9}/\S+", text, flags=re.I)
|
||||
if not match:
|
||||
return ""
|
||||
return match.group(0).rstrip(".,;:)]}").lower()
|
||||
|
||||
|
||||
def normalize_pmid(value: Optional[str]) -> str:
|
||||
match = re.search(r"\d{5,10}", str(value or ""))
|
||||
return match.group(0) if match else ""
|
||||
|
||||
|
||||
def normalize_pmcid(value: Optional[str]) -> str:
|
||||
match = re.search(r"PMC\s*(\d+)", str(value or ""), flags=re.I)
|
||||
return f"PMC{match.group(1)}" if match else ""
|
||||
|
||||
|
||||
def doi_from_pdf_name(pdf_path: Optional[Path]) -> str:
|
||||
if not pdf_path:
|
||||
return ""
|
||||
stem = Path(pdf_path).stem
|
||||
direct = normalize_doi(stem)
|
||||
if direct:
|
||||
return direct
|
||||
match = re.match(r"^(10\.\d{4,9})_(.+)$", stem, flags=re.I)
|
||||
if match:
|
||||
return normalize_doi(f"{match.group(1)}/{match.group(2)}")
|
||||
return ""
|
||||
|
||||
|
||||
def doi_candidates_from_pdf_name(pdf_path: Optional[Path]) -> list[str]:
|
||||
if not pdf_path:
|
||||
return []
|
||||
stem = Path(pdf_path).stem
|
||||
match = re.match(r"^(10\.\d{4,9})_(.+)$", stem, flags=re.I)
|
||||
if not match:
|
||||
direct = normalize_doi(stem)
|
||||
return [direct] if direct else []
|
||||
prefix, suffix = match.groups()
|
||||
parts = suffix.split("_")
|
||||
candidates = []
|
||||
# The first filename underscore always represents the DOI's first slash.
|
||||
# Later underscores may be literal underscores or additional slashes.
|
||||
combinations = 1 << max(0, len(parts) - 1)
|
||||
for mask in range(combinations):
|
||||
rebuilt = parts[0]
|
||||
for index, part in enumerate(parts[1:]):
|
||||
separator = "/" if mask & (1 << index) else "_"
|
||||
rebuilt += separator + part
|
||||
candidate = normalize_doi(f"{prefix}/{rebuilt}")
|
||||
if candidate and candidate not in candidates:
|
||||
candidates.append(candidate)
|
||||
# Prefer slash-rich candidates after the simple one only when an external
|
||||
# metadata source validates them.
|
||||
return candidates
|
||||
|
||||
|
||||
def sha256_file(path: Optional[Path]) -> str:
|
||||
if not path or not Path(path).exists():
|
||||
return ""
|
||||
digest = hashlib.sha256()
|
||||
with Path(path).open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
class PaperMetadataEnricherV16:
|
||||
def __init__(self, timeout: int = 30):
|
||||
self.timeout = timeout
|
||||
self.session = requests.Session()
|
||||
self.session.headers.update({"User-Agent": USER_AGENT})
|
||||
self.audit: list[dict[str, Any]] = []
|
||||
|
||||
def _record(self, source: str, action: str, detail: str = "") -> None:
|
||||
self.audit.append(
|
||||
{
|
||||
"source": source,
|
||||
"action": action,
|
||||
"detail": detail,
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
)
|
||||
|
||||
def _europe_pmc(self, *, doi: str = "", pmid: str = "") -> Dict[str, Any]:
|
||||
if doi:
|
||||
query = f'DOI:"{doi}"'
|
||||
elif pmid:
|
||||
query = f"EXT_ID:{pmid} AND SRC:MED"
|
||||
else:
|
||||
return {}
|
||||
try:
|
||||
response = self.session.get(
|
||||
EUROPE_PMC_SEARCH,
|
||||
params={"query": query, "format": "json", "pageSize": 1},
|
||||
timeout=self.timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
rows = response.json().get("resultList", {}).get("result", [])
|
||||
if not rows:
|
||||
self._record("europe_pmc", "not_found", query)
|
||||
return {}
|
||||
row = rows[0]
|
||||
authors = []
|
||||
for author in (row.get("authorList") or {}).get("author", []) or []:
|
||||
name = author.get("fullName") or author.get("firstName")
|
||||
if name:
|
||||
authors.append(name)
|
||||
self._record("europe_pmc", "resolved", query)
|
||||
return {
|
||||
"doi": normalize_doi(row.get("doi") or doi),
|
||||
"pmid": normalize_pmid(row.get("pmid") or pmid),
|
||||
"pmcid": normalize_pmcid(row.get("pmcid")),
|
||||
"title": str(row.get("title") or "").strip(),
|
||||
"authors": authors,
|
||||
"author_string": str(row.get("authorString") or "").strip(),
|
||||
"journal": str(row.get("journalTitle") or "").strip(),
|
||||
"publication_year": str(row.get("pubYear") or "").strip(),
|
||||
"publication_date": str(
|
||||
row.get("firstPublicationDate")
|
||||
or row.get("journalInfo", {}).get("printPublicationDate")
|
||||
or ""
|
||||
).strip(),
|
||||
"has_supplementary": row.get("hasSuppl") == "Y",
|
||||
"is_open_access": row.get("isOpenAccess") == "Y",
|
||||
}
|
||||
except Exception as exc:
|
||||
self._record("europe_pmc", "error", str(exc))
|
||||
return {}
|
||||
|
||||
def _ncbi_idconv(self, identifier: str) -> Dict[str, Any]:
|
||||
if not identifier:
|
||||
return {}
|
||||
try:
|
||||
response = self.session.get(
|
||||
NCBI_ID_CONVERTER,
|
||||
params={"ids": identifier, "format": "json", "tool": NCBI_TOOL_NAME},
|
||||
timeout=self.timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
records = response.json().get("records", [])
|
||||
if not records:
|
||||
self._record("ncbi_idconv", "not_found", identifier)
|
||||
return {}
|
||||
row = records[0]
|
||||
self._record("ncbi_idconv", "resolved", identifier)
|
||||
return {
|
||||
"doi": normalize_doi(row.get("doi")),
|
||||
"pmid": normalize_pmid(row.get("pmid")),
|
||||
"pmcid": normalize_pmcid(row.get("pmcid")),
|
||||
}
|
||||
except Exception as exc:
|
||||
self._record("ncbi_idconv", "error", str(exc))
|
||||
return {}
|
||||
|
||||
def _crossref(self, doi: str) -> Dict[str, Any]:
|
||||
if not doi:
|
||||
return {}
|
||||
try:
|
||||
response = self.session.get(
|
||||
f"{CROSSREF_WORKS}/{doi}",
|
||||
timeout=self.timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
row = response.json().get("message", {})
|
||||
authors = []
|
||||
for author in row.get("author", []) or []:
|
||||
name = " ".join(
|
||||
x for x in [author.get("given", ""), author.get("family", "")] if x
|
||||
).strip()
|
||||
if name:
|
||||
authors.append(name)
|
||||
dates = (
|
||||
row.get("published-print")
|
||||
or row.get("published-online")
|
||||
or row.get("issued")
|
||||
or {}
|
||||
).get("date-parts", [])
|
||||
date_parts = dates[0] if dates else []
|
||||
publication_date = "-".join(str(x) for x in date_parts)
|
||||
self._record("crossref", "resolved", doi)
|
||||
return {
|
||||
"title": str((row.get("title") or [""])[0]).strip(),
|
||||
"authors": authors,
|
||||
"author_string": ", ".join(authors),
|
||||
"journal": str((row.get("container-title") or [""])[0]).strip(),
|
||||
"publication_year": str(date_parts[0]) if date_parts else "",
|
||||
"publication_date": publication_date,
|
||||
"publisher": str(row.get("publisher") or "").strip(),
|
||||
}
|
||||
except Exception as exc:
|
||||
self._record("crossref", "error", str(exc))
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
def _merge(target: Dict[str, Any], source: Dict[str, Any]) -> None:
|
||||
for key, value in source.items():
|
||||
if value not in (None, "", [], {}):
|
||||
if target.get(key) in (None, "", [], {}):
|
||||
target[key] = value
|
||||
|
||||
def resolve(
|
||||
self,
|
||||
*,
|
||||
pdf_path: Optional[Path] = None,
|
||||
doi: Optional[str] = None,
|
||||
pmid: Optional[str] = None,
|
||||
pmcid: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
explicit_doi = normalize_doi(doi)
|
||||
metadata: Dict[str, Any] = {
|
||||
"doi": explicit_doi,
|
||||
"pmid": normalize_pmid(pmid),
|
||||
"pmcid": normalize_pmcid(pmcid),
|
||||
"title": "",
|
||||
"authors": [],
|
||||
"author_string": "",
|
||||
"journal": "",
|
||||
"publication_year": "",
|
||||
"publication_date": "",
|
||||
"publisher": "",
|
||||
"has_supplementary": False,
|
||||
"is_open_access": False,
|
||||
}
|
||||
if explicit_doi:
|
||||
self._record("explicit_input", "doi", metadata["doi"])
|
||||
if metadata["pmid"]:
|
||||
self._record("explicit_input", "pmid", metadata["pmid"])
|
||||
if metadata["pmcid"]:
|
||||
self._record("explicit_input", "pmcid", metadata["pmcid"])
|
||||
|
||||
if not metadata["doi"] and not metadata["pmid"] and pdf_path:
|
||||
candidates = doi_candidates_from_pdf_name(pdf_path)
|
||||
for candidate in candidates:
|
||||
europe = self._europe_pmc(doi=candidate)
|
||||
if europe.get("title") or europe.get("pmid") or europe.get("pmcid"):
|
||||
metadata["doi"] = normalize_doi(europe.get("doi")) or candidate
|
||||
self._merge(metadata, europe)
|
||||
self._record("pdf_filename", "doi_validated", metadata["doi"])
|
||||
break
|
||||
crossref = self._crossref(candidate)
|
||||
if crossref.get("title"):
|
||||
metadata["doi"] = candidate
|
||||
self._merge(metadata, crossref)
|
||||
self._record("pdf_filename", "doi_validated", metadata["doi"])
|
||||
break
|
||||
if not metadata["doi"] and candidates:
|
||||
metadata["doi"] = candidates[0]
|
||||
self._record("pdf_filename", "doi_unverified", metadata["doi"])
|
||||
|
||||
europe_metadata = self._europe_pmc(
|
||||
doi=metadata["doi"],
|
||||
pmid=metadata["pmid"],
|
||||
)
|
||||
self._merge(metadata, europe_metadata)
|
||||
identifier = metadata["doi"] or metadata["pmid"] or metadata["pmcid"]
|
||||
self._merge(metadata, self._ncbi_idconv(identifier))
|
||||
self._merge(metadata, self._crossref(metadata["doi"]))
|
||||
|
||||
if metadata["doi"]:
|
||||
paper_id = f"doi:{metadata['doi']}"
|
||||
elif metadata["pmid"]:
|
||||
paper_id = f"pmid:{metadata['pmid']}"
|
||||
elif metadata["pmcid"]:
|
||||
paper_id = f"pmcid:{metadata['pmcid']}"
|
||||
else:
|
||||
pdf_hash = sha256_file(pdf_path)
|
||||
paper_id = f"sha256:{pdf_hash}" if pdf_hash else ""
|
||||
|
||||
metadata.update(
|
||||
{
|
||||
"paper_id": paper_id,
|
||||
"source_pdf": Path(pdf_path).name if pdf_path else "",
|
||||
"source_pdf_sha256": sha256_file(pdf_path),
|
||||
"doi_url": f"https://doi.org/{metadata['doi']}" if metadata["doi"] else "",
|
||||
"pubmed_url": (
|
||||
f"https://pubmed.ncbi.nlm.nih.gov/{metadata['pmid']}/"
|
||||
if metadata["pmid"]
|
||||
else ""
|
||||
),
|
||||
"pmc_url": (
|
||||
f"https://pmc.ncbi.nlm.nih.gov/articles/{metadata['pmcid']}/"
|
||||
if metadata["pmcid"]
|
||||
else ""
|
||||
),
|
||||
"metadata_sources": sorted(
|
||||
{row["source"] for row in self.audit if row["action"] == "resolved"}
|
||||
| {
|
||||
row["source"]
|
||||
for row in self.audit
|
||||
if row["source"] in {"explicit_input", "pdf_filename"}
|
||||
}
|
||||
),
|
||||
"resolved_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
)
|
||||
return metadata
|
||||
|
||||
def enrich_output(
|
||||
self,
|
||||
output: Dict[str, Any],
|
||||
paper: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
output["paper"] = paper
|
||||
for relation in output.get("relations", []) or []:
|
||||
locator = (
|
||||
relation.get("relation_unit_id")
|
||||
or relation.get("section")
|
||||
or str(relation.get("p_source") or "").split(":", 1)[-1]
|
||||
)
|
||||
relation.update(
|
||||
{
|
||||
"paper_id": paper.get("paper_id", ""),
|
||||
"doi": paper.get("doi", ""),
|
||||
"pmid": paper.get("pmid", ""),
|
||||
"pmcid": paper.get("pmcid", ""),
|
||||
"pubmed_id": paper.get("pmid", ""),
|
||||
"paper_title": paper.get("title", ""),
|
||||
"journal": paper.get("journal", ""),
|
||||
"publication_year": paper.get("publication_year", ""),
|
||||
"doi_url": paper.get("doi_url", ""),
|
||||
"pubmed_url": paper.get("pubmed_url", ""),
|
||||
"pmc_url": paper.get("pmc_url", ""),
|
||||
"p_source": (
|
||||
f"{paper.get('paper_id', '')}:{locator}"
|
||||
if paper.get("paper_id") and locator
|
||||
else relation.get("p_source", "")
|
||||
),
|
||||
}
|
||||
)
|
||||
for chunk in output.get("chunks", []) or []:
|
||||
chunk.update(
|
||||
{
|
||||
"paper_id": paper.get("paper_id", ""),
|
||||
"doi": paper.get("doi", ""),
|
||||
"pmid": paper.get("pmid", ""),
|
||||
"pmcid": paper.get("pmcid", ""),
|
||||
}
|
||||
)
|
||||
metadata = output.setdefault("metadata", {})
|
||||
metadata["paper_metadata_enrichment"] = {
|
||||
"enabled": True,
|
||||
"paper_id": paper.get("paper_id", ""),
|
||||
"doi": paper.get("doi", ""),
|
||||
"pmid": paper.get("pmid", ""),
|
||||
"pmcid": paper.get("pmcid", ""),
|
||||
"metadata_sources": paper.get("metadata_sources", []),
|
||||
"relation_count_enriched": len(output.get("relations", []) or []),
|
||||
"chunk_count_enriched": len(output.get("chunks", []) or []),
|
||||
}
|
||||
return output
|
||||
|
||||
def save_audit(self, path: Path, paper: Dict[str, Any]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(
|
||||
json.dumps(
|
||||
{"paper": paper, "queries": self.audit},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
@@ -0,0 +1,238 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
PC数据层加载器
|
||||
加载PC的4个核心数据文件并提供查询接口
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Set
|
||||
|
||||
try:
|
||||
from ...project_config import pc_data_root
|
||||
except Exception: # pragma: no cover - supports direct script execution
|
||||
def pc_data_root() -> Path:
|
||||
value = os.environ.get("PC_DATA_ROOT", "").strip()
|
||||
if value:
|
||||
return Path(value)
|
||||
return Path(__file__).resolve().parents[4] / "data" / "pc"
|
||||
|
||||
|
||||
class PCDataLoader:
|
||||
"""加载PC数据文件并提供查询接口"""
|
||||
|
||||
def __init__(self, pc_data_dir: str | Path | None = None):
|
||||
if pc_data_dir is None:
|
||||
pc_data_dir = pc_data_root()
|
||||
|
||||
self.pc_data_dir = Path(pc_data_dir)
|
||||
|
||||
# 数据缓存
|
||||
self._alias_map: Optional[Dict[str, str]] = None # scientific_chunks.json
|
||||
self._gene_aliases: Optional[Dict[str, List[str]]] = None # gene_alias.json
|
||||
self._relationship_patterns: Optional[Dict] = None # relationships_by_category_final.json
|
||||
self._entity_dict: Optional[Set[str]] = None # all_dic.json
|
||||
|
||||
def load_scientific_chunks(self) -> Dict[str, str]:
|
||||
"""
|
||||
加载scientific_chunks.json
|
||||
格式: {"aan": "AAN, AT5G46760", "aar1": "AAR1, AT5G46760", ...}
|
||||
用途: 小写别名→标准化名称映射
|
||||
"""
|
||||
if self._alias_map is not None:
|
||||
return self._alias_map
|
||||
|
||||
file_path = self.pc_data_dir / 'scientific_chunks.json'
|
||||
if not file_path.exists():
|
||||
print(f"⚠️ 未找到PC别名映射文件: {file_path}")
|
||||
return {}
|
||||
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
# 解析: "aan": "AAN, AT5G46760" → {"aan": "AAN"}
|
||||
self._alias_map = {}
|
||||
for alias, standard_names in data.items():
|
||||
# 取第一个逗号前的名称作为主名称
|
||||
primary_name = standard_names.split(',')[0].strip()
|
||||
self._alias_map[alias.lower()] = primary_name
|
||||
|
||||
print(f"✓ 加载PC别名映射: {len(self._alias_map)} 条")
|
||||
return self._alias_map
|
||||
|
||||
def load_gene_aliases(self) -> Dict[str, List[str]]:
|
||||
"""
|
||||
加载gene_alias.json
|
||||
格式: [{"gene": "SUC5", "aliases": ["ATSUC5", "SUC5", "AT1G71890"]}, ...]
|
||||
用途: 基因→多别名映射
|
||||
"""
|
||||
if self._gene_aliases is not None:
|
||||
return self._gene_aliases
|
||||
|
||||
file_path = self.pc_data_dir / 'gene_alias.json'
|
||||
if not file_path.exists():
|
||||
print(f"⚠️ 未找到PC基因别名文件: {file_path}")
|
||||
return {}
|
||||
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
self._gene_aliases = {item['gene']: item['aliases'] for item in data}
|
||||
print(f"✓ 加载PC基因别名: {len(self._gene_aliases)} 个基因")
|
||||
return self._gene_aliases
|
||||
|
||||
def load_relationship_patterns(self) -> Dict:
|
||||
"""
|
||||
加载relationships_by_category_final.json
|
||||
格式: {"ACTIVATION/INDUCTION/...": ["AFFECT", "INDUCE", ...], ...}
|
||||
用途: 788种关系模式
|
||||
"""
|
||||
if self._relationship_patterns is not None:
|
||||
return self._relationship_patterns
|
||||
|
||||
file_path = self.pc_data_dir / 'relationships_by_category_final.json'
|
||||
if not file_path.exists():
|
||||
print(f"⚠️ 未找到PC关系模式文件: {file_path}")
|
||||
return {}
|
||||
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
self._relationship_patterns = json.load(f)
|
||||
|
||||
total_patterns = sum(len(v) for v in self._relationship_patterns.values())
|
||||
print(f"✓ 加载PC关系模式: {len(self._relationship_patterns)} 类, {total_patterns} 种")
|
||||
return self._relationship_patterns
|
||||
|
||||
def load_entity_dict(self) -> Set[str]:
|
||||
"""
|
||||
加载all_dic.json
|
||||
格式: 复杂的嵌套JSON,包含238万个实体
|
||||
用途: 实体目录验证
|
||||
V14.2优化: 由于文件太大(95MB, 238万行),采用惰性加载策略
|
||||
"""
|
||||
if self._entity_dict is not None:
|
||||
return self._entity_dict
|
||||
|
||||
file_path = self.pc_data_dir / 'all_dic.json'
|
||||
if not file_path.exists():
|
||||
print(f"⚠️ 未找到PC实体目录文件: {file_path}")
|
||||
return set()
|
||||
|
||||
# V14.2优化: 由于文件太大,只加载前10000个实体作为样本
|
||||
# 完整的验证应该使用别名映射(alias_map)代替
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
# 处理复杂的嵌套格式
|
||||
entities = set()
|
||||
if isinstance(data, list):
|
||||
# 如果是列表,展平第一层
|
||||
for item in data[:100]: # 只处理前100个元素
|
||||
if isinstance(item, str):
|
||||
entities.add(item)
|
||||
elif isinstance(item, list):
|
||||
entities.update(item[:1000]) # 每个子列表最多1000个
|
||||
elif isinstance(item, dict):
|
||||
# 如果是字典,提取所有值
|
||||
for v in list(item.values())[:1000]:
|
||||
if isinstance(v, list):
|
||||
entities.update(v[:100])
|
||||
elif isinstance(v, str):
|
||||
entities.add(v)
|
||||
|
||||
self._entity_dict = entities
|
||||
print(f"✓ 加载PC实体目录: {len(self._entity_dict)} 个实体 (样本)")
|
||||
print(f" ⚠️ 完整目录太大(238万),使用别名映射进行验证")
|
||||
|
||||
except Exception as e:
|
||||
print(f"⚠️ 加载PC实体目录失败: {e}")
|
||||
self._entity_dict = set()
|
||||
|
||||
return self._entity_dict
|
||||
|
||||
def normalize_entity(self, name: str) -> str:
|
||||
"""使用PC别名映射标准化实体名"""
|
||||
alias_map = self.load_scientific_chunks()
|
||||
lower_name = name.lower()
|
||||
if lower_name in alias_map:
|
||||
return alias_map[lower_name]
|
||||
return name
|
||||
|
||||
def expand_gene_aliases(self, gene: str) -> List[str]:
|
||||
"""扩展基因别名"""
|
||||
gene_aliases = self.load_gene_aliases()
|
||||
if gene in gene_aliases:
|
||||
return gene_aliases[gene]
|
||||
return [gene]
|
||||
|
||||
def is_valid_entity(self, name: str) -> bool:
|
||||
"""验证实体是否在PC目录中"""
|
||||
entity_dict = self.load_entity_dict()
|
||||
return name in entity_dict
|
||||
|
||||
def get_pc_patterns_as_regex(self) -> List[tuple]:
|
||||
"""
|
||||
将PC关系模式转换为正则表达式
|
||||
返回: [(pattern, relation, source_type, target_type), ...]
|
||||
"""
|
||||
import re
|
||||
|
||||
patterns = self.load_relationship_patterns()
|
||||
regex_patterns = []
|
||||
|
||||
# 基因名模式
|
||||
gene_pattern = r'[A-Z][a-z]{1,2}[A-Z]{1,4}\d{1,4}[a-z]?'
|
||||
|
||||
for category, rel_list in patterns.items():
|
||||
for rel in rel_list:
|
||||
# 将关系词转换为正则
|
||||
# 例如: "UP-REGULATED" → r'up[-\s]?regulated'
|
||||
rel_regex = rel.lower().replace(' ', r'\s+').replace('-', r'[-\s]?')
|
||||
|
||||
# 构建完整模式
|
||||
# 示例: (\bAtSPL1\b).*?up[-\s]?regulated.*?in\s+(.+)
|
||||
pattern = (
|
||||
rf'(\b{gene_pattern}\b).*?{rel_regex}.*?in\s+(.+)',
|
||||
rel.lower().replace('-', '_').replace(' ', '_'),
|
||||
'gene',
|
||||
'target'
|
||||
)
|
||||
regex_patterns.append(pattern)
|
||||
|
||||
return regex_patterns
|
||||
|
||||
|
||||
# 全局单例
|
||||
_pc_data_loader: Optional[PCDataLoader] = None
|
||||
|
||||
def get_pc_data_loader() -> PCDataLoader:
|
||||
"""获取PC数据加载器单例"""
|
||||
global _pc_data_loader
|
||||
if _pc_data_loader is None:
|
||||
_pc_data_loader = PCDataLoader()
|
||||
return _pc_data_loader
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 测试加载
|
||||
loader = PCDataLoader()
|
||||
|
||||
# 测试别名映射
|
||||
alias_map = loader.load_scientific_chunks()
|
||||
print(f"\n别名映射示例:")
|
||||
for k, v in list(alias_map.items())[:5]:
|
||||
print(f" {k} → {v}")
|
||||
|
||||
# 测试基因别名
|
||||
gene_aliases = loader.load_gene_aliases()
|
||||
print(f"\n基因别名示例:")
|
||||
for k, v in list(gene_aliases.items())[:3]:
|
||||
print(f" {k}: {v}")
|
||||
|
||||
# 测试关系模式
|
||||
patterns = loader.load_relationship_patterns()
|
||||
print(f"\n关系模式类别:")
|
||||
for cat, rels in list(patterns.items())[:3]:
|
||||
print(f" {cat}: {len(rels)} 种")
|
||||
print(f" 示例: {rels[:3]}")
|
||||
@@ -0,0 +1,180 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
PC 模式导入器
|
||||
从 PC 数据库自动生成关系模式
|
||||
"""
|
||||
|
||||
import re
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Tuple
|
||||
from collections import defaultdict
|
||||
|
||||
|
||||
class PCPatternImporter:
|
||||
"""从 PC 数据库导入关系模式"""
|
||||
|
||||
def __init__(self, pc_data_path: str):
|
||||
self.pc_data_path = Path(pc_data_path)
|
||||
self.pc_relations = []
|
||||
self.pattern_templates = {}
|
||||
|
||||
def load_pc_data(self):
|
||||
"""加载 PC 数据"""
|
||||
with open(self.pc_data_path, 'r', encoding='utf-8') as f:
|
||||
self.pc_relations = json.load(f)
|
||||
print(f"✓ 加载 PC 数据: {len(self.pc_relations)} 个关系")
|
||||
|
||||
def analyze_patterns(self):
|
||||
"""分析 PC 关系模式"""
|
||||
# 按关系类型分组
|
||||
rel_by_type = defaultdict(list)
|
||||
for r in self.pc_relations:
|
||||
rel_by_type[r['relation']].append(r)
|
||||
|
||||
print(f"✓ 分析关系类型: {len(rel_by_type)} 种")
|
||||
|
||||
# 为每种关系类型生成模式模板
|
||||
for rel_type, relations in rel_by_type.items():
|
||||
# 分析源实体和目标实体的类型
|
||||
src_types = defaultdict(int)
|
||||
tgt_types = defaultdict(int)
|
||||
|
||||
for r in relations:
|
||||
src_types[r['source_type']] += 1
|
||||
tgt_types[r['target_type']] += 1
|
||||
|
||||
# 取最常见的类型
|
||||
src_type = max(src_types, key=src_types.get) if src_types else 'entity'
|
||||
tgt_type = max(tgt_types, key=tgt_types.get) if tgt_types else 'entity'
|
||||
|
||||
self.pattern_templates[rel_type] = {
|
||||
'count': len(relations),
|
||||
'src_type': src_type,
|
||||
'tgt_type': tgt_type,
|
||||
'examples': relations[:3]
|
||||
}
|
||||
|
||||
return self.pattern_templates
|
||||
|
||||
def generate_regex_patterns(self) -> List[Tuple[str, str, str, str]]:
|
||||
"""生成正则表达式模式"""
|
||||
patterns = []
|
||||
|
||||
# 基因名模式
|
||||
gene_pattern = r'[A-Z][a-z]{1,2}[A-Z]{1,4}\d{1,4}[a-z]?'
|
||||
|
||||
# 为高频关系类型生成模式
|
||||
for rel_type, info in self.pattern_templates.items():
|
||||
if info['count'] < 2: # 跳过低频关系
|
||||
continue
|
||||
|
||||
# 根据关系类型生成不同的正则
|
||||
rel_lower = rel_type.lower()
|
||||
|
||||
# 表达量变化模式
|
||||
if 'up-regulated' in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?(?:was|were|is|are)\s+up[-\s]?regulated\s+(?:during|under|in|at|after)\s+(.+)',
|
||||
'up-regulated', 'gene', 'condition'
|
||||
))
|
||||
elif 'down-regulated' in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?(?:was|were|is|are)\s+down[-\s]?regulated\s+(?:during|under|in|at|after)\s+(.+)',
|
||||
'down-regulated', 'gene', 'condition'
|
||||
))
|
||||
elif 'induced' in rel_lower and 'not' not in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?(?:was|were|is|are)\s+induced\s+(?:by|during|under|at|after)\s+(.+)',
|
||||
'induced', 'gene', 'condition'
|
||||
))
|
||||
elif 'collinear' in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?(?:is|was|were)\s+collinear\s+with\s+(\b{gene_pattern}\b)',
|
||||
'collinear_with', 'gene', 'gene'
|
||||
))
|
||||
elif 'orthologous' in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?(?:is|was|were)\s+(?:highly\s+)?orthologous\s+to\s+(\b{gene_pattern}\b)',
|
||||
'orthologous_to', 'gene', 'gene'
|
||||
))
|
||||
elif 'duplicated' in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?(?:were|was|is|are)\s+segmentally\s+duplicated\s+(?:in|within)\s+(.+)',
|
||||
'segmentally_duplicated_in', 'gene', 'organism'
|
||||
))
|
||||
elif 'regulate' in rel_lower and 'by' not in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?regulates?\s+(.+)',
|
||||
'regulate', 'gene', 'process'
|
||||
))
|
||||
elif 'involved' in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?involved\s+in\s+(.+)',
|
||||
'involved_in', 'gene', 'process'
|
||||
))
|
||||
elif 'identified' in rel_lower and 'in' in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?identified\s+in\s+(.+)',
|
||||
'identified_in', 'gene', 'organism'
|
||||
))
|
||||
elif 'respond' in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?respond\s+to\s+(.+)',
|
||||
'respond_to', 'gene', 'condition'
|
||||
))
|
||||
elif 'expressed' in rel_lower and 'in' in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?expressed\s+in\s+(.+)',
|
||||
'expressed_in', 'gene', 'tissue'
|
||||
))
|
||||
elif 'show' in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?show(?:ed|s)?\s+(.+)',
|
||||
'show', 'gene', 'phenotype'
|
||||
))
|
||||
elif 'had' in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?had\s+(.+)',
|
||||
'had', 'gene', 'attribute'
|
||||
))
|
||||
elif 'play' in rel_lower:
|
||||
patterns.append((
|
||||
rf'(\b{gene_pattern}\b).*?play(?:s|\s+a\s+key\s+role)?\s+(?:in|during)\s+(.+)',
|
||||
'play_role_in', 'gene', 'process'
|
||||
))
|
||||
|
||||
return patterns
|
||||
|
||||
def export_patterns(self, output_path: str):
|
||||
"""导出模式到文件"""
|
||||
patterns = self.generate_regex_patterns()
|
||||
|
||||
# 去重
|
||||
unique_patterns = []
|
||||
seen = set()
|
||||
for pattern, rel, src, tgt in patterns:
|
||||
key = (rel, src, tgt)
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
unique_patterns.append({
|
||||
'pattern': pattern,
|
||||
'relation': rel,
|
||||
'source_type': src,
|
||||
'target_type': tgt
|
||||
})
|
||||
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(unique_patterns, f, ensure_ascii=False, indent=2)
|
||||
|
||||
print(f"✓ 导出模式: {len(unique_patterns)} 个 -> {output_path}")
|
||||
return unique_patterns
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 测试
|
||||
importer = PCPatternImporter('./pc_34966577_relations.json')
|
||||
importer.load_pc_data()
|
||||
templates = importer.analyze_patterns()
|
||||
patterns = importer.export_patterns('./generated_pc_patterns.json')
|
||||
print(f"生成完成: {len(patterns)} 个模式")
|
||||
@@ -0,0 +1,372 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
PC模式加载器 + 自动学习扩展
|
||||
整合PC的788种模式 + 用户学习到的模式
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Tuple, Optional, Set
|
||||
from dataclasses import dataclass, asdict
|
||||
from datetime import datetime
|
||||
|
||||
try:
|
||||
from ...project_config import learning_dir as configured_learning_dir
|
||||
from ...project_config import pc_data_root
|
||||
except Exception: # pragma: no cover - supports direct script execution
|
||||
def pc_data_root() -> Path:
|
||||
value = os.environ.get("PC_DATA_ROOT", "").strip()
|
||||
if value:
|
||||
return Path(value)
|
||||
return Path(__file__).resolve().parents[4] / "data" / "pc"
|
||||
|
||||
def configured_learning_dir() -> Path:
|
||||
value = os.environ.get("RAG_LEARNING_DIR", "").strip()
|
||||
if value:
|
||||
return Path(value)
|
||||
return Path(__file__).resolve().parents[4] / "v14_learning"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RelationPattern:
|
||||
"""关系模式定义"""
|
||||
relation: str
|
||||
pattern: str # 正则表达式
|
||||
source_type: str
|
||||
target_type: str
|
||||
category: str # PC分类,如"ACTIVATION/INDUCTION/..."
|
||||
confidence: float = 0.9
|
||||
source: str = "pc" # "pc" 或 "learned"
|
||||
learned_from: Optional[str] = None # 反馈ID(如果是学习的)
|
||||
timestamp: Optional[str] = None
|
||||
usage_count: int = 0 # 使用次数
|
||||
|
||||
|
||||
class PCPatternsLoader:
|
||||
"""
|
||||
PC模式加载器 + 自动学习
|
||||
|
||||
功能:
|
||||
1. 加载PC的788种关系模式
|
||||
2. 加载用户学习到的模式
|
||||
3. 合并使用(PC模式 + 学习模式)
|
||||
4. 支持模式优先级(PC模式优先或学习模式优先)
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
pc_data_dir: str | Path | None = None,
|
||||
learning_dir: str | Path | None = None):
|
||||
if pc_data_dir is None:
|
||||
pc_data_dir = pc_data_root()
|
||||
if learning_dir is None:
|
||||
learning_dir = configured_learning_dir()
|
||||
|
||||
self.pc_data_dir = Path(pc_data_dir)
|
||||
self.learning_dir = Path(learning_dir)
|
||||
self.learning_dir.mkdir(exist_ok=True)
|
||||
|
||||
# 模式存储
|
||||
self.pc_patterns: Dict[str, RelationPattern] = {} # PC的788种
|
||||
self.learned_patterns: Dict[str, RelationPattern] = {} # 学习的模式
|
||||
self.all_patterns: Dict[str, RelationPattern] = {} # 合并后的所有模式
|
||||
|
||||
# 学习记录文件
|
||||
self.learned_patterns_file = self.learning_dir / 'learned_patterns.json'
|
||||
|
||||
# 加载所有模式
|
||||
self._load_pc_patterns()
|
||||
self._load_learned_patterns()
|
||||
self._merge_patterns()
|
||||
|
||||
print(f"✅ PC模式加载器初始化完成")
|
||||
print(f" PC模式: {len(self.pc_patterns)} 种")
|
||||
print(f" 学习模式: {len(self.learned_patterns)} 种")
|
||||
print(f" 总模式: {len(self.all_patterns)} 种")
|
||||
|
||||
def _load_pc_patterns(self):
|
||||
"""加载PC的788种关系模式"""
|
||||
pattern_file = self.pc_data_dir / 'relationships_by_category_final.json'
|
||||
|
||||
if not pattern_file.exists():
|
||||
print(f"⚠️ PC模式文件不存在: {pattern_file}")
|
||||
return
|
||||
|
||||
with open(pattern_file, 'r', encoding='utf-8') as f:
|
||||
categories = json.load(f)
|
||||
|
||||
# 基因名模式
|
||||
gene_pattern = r'[A-Z][a-z]{1,2}[A-Z]{1,4}\d{1,4}[a-z]?'
|
||||
|
||||
for category, relations in categories.items():
|
||||
for rel in relations:
|
||||
# 标准化关系名
|
||||
relation_name = rel.lower().replace('-', '_').replace(' ', '_')
|
||||
|
||||
# 生成正则表达式
|
||||
rel_regex = re.escape(rel.lower()).replace(' ', r'\s+').replace('-', r'[-\s]?')
|
||||
pattern = rf'(\b{gene_pattern}\b).*?{rel_regex}.*?(.+?)\.'
|
||||
|
||||
self.pc_patterns[relation_name] = RelationPattern(
|
||||
relation=relation_name,
|
||||
pattern=pattern,
|
||||
source_type='gene',
|
||||
target_type='target',
|
||||
category=category,
|
||||
confidence=0.9,
|
||||
source='pc'
|
||||
)
|
||||
|
||||
print(f"✓ 加载PC模式: {len(self.pc_patterns)} 种")
|
||||
|
||||
def _load_learned_patterns(self):
|
||||
"""加载用户学习到的模式"""
|
||||
if not self.learned_patterns_file.exists():
|
||||
print(f"✓ 无学习模式文件,将创建新的")
|
||||
return
|
||||
|
||||
with open(self.learned_patterns_file, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
for rel_name, pattern_data in data.items():
|
||||
self.learned_patterns[rel_name] = RelationPattern(**pattern_data)
|
||||
|
||||
print(f"✓ 加载学习模式: {len(self.learned_patterns)} 种")
|
||||
|
||||
def _save_learned_patterns(self):
|
||||
"""保存学习到的模式"""
|
||||
data = {k: asdict(v) for k, v in self.learned_patterns.items()}
|
||||
with open(self.learned_patterns_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
def _merge_patterns(self):
|
||||
"""合并PC模式和学习模式"""
|
||||
# PC模式优先(基础)
|
||||
self.all_patterns = dict(self.pc_patterns)
|
||||
|
||||
# 添加学习模式(如果与PC不冲突)
|
||||
for rel_name, pattern in self.learned_patterns.items():
|
||||
if rel_name not in self.all_patterns:
|
||||
# 新学习的模式,直接添加
|
||||
self.all_patterns[rel_name] = pattern
|
||||
else:
|
||||
# 冲突:比较置信度,高的保留
|
||||
if pattern.confidence > self.all_patterns[rel_name].confidence:
|
||||
self.all_patterns[rel_name] = pattern
|
||||
print(f" 学习模式覆盖PC模式: {rel_name}")
|
||||
|
||||
def get_all_patterns(self) -> List[Tuple[str, str, str, str]]:
|
||||
"""
|
||||
获取所有模式用于匹配
|
||||
返回: [(pattern, relation, source_type, target_type), ...]
|
||||
"""
|
||||
patterns = []
|
||||
for rel_name, pattern in self.all_patterns.items():
|
||||
patterns.append((
|
||||
pattern.pattern,
|
||||
pattern.relation,
|
||||
pattern.source_type,
|
||||
pattern.target_type
|
||||
))
|
||||
return patterns
|
||||
|
||||
def get_pattern_stats(self) -> Dict:
|
||||
"""获取模式统计"""
|
||||
pc_count = len(self.pc_patterns)
|
||||
learned_count = len(self.learned_patterns)
|
||||
|
||||
# 按分类统计
|
||||
category_stats = {}
|
||||
for p in self.pc_patterns.values():
|
||||
category_stats[p.category] = category_stats.get(p.category, 0) + 1
|
||||
|
||||
return {
|
||||
'pc_patterns': pc_count,
|
||||
'learned_patterns': learned_count,
|
||||
'total_patterns': len(self.all_patterns),
|
||||
'category_distribution': category_stats,
|
||||
'learned_relations': list(self.learned_patterns.keys())
|
||||
}
|
||||
|
||||
# ========== 自动学习接口 ==========
|
||||
|
||||
def learn_from_feedback(self,
|
||||
feedback_type: str,
|
||||
original_text: str,
|
||||
missed_relation: Optional[str] = None,
|
||||
context: str = "",
|
||||
feedback_id: Optional[str] = None) -> Optional[RelationPattern]:
|
||||
"""
|
||||
从用户反馈学习新模式
|
||||
|
||||
Args:
|
||||
feedback_type: 'missed_relation', 'wrong_pattern', etc.
|
||||
original_text: 原始文本
|
||||
missed_relation: 遗漏的关系名
|
||||
context: 上下文
|
||||
feedback_id: 反馈ID
|
||||
|
||||
Returns:
|
||||
学习到的模式,如果学习失败返回None
|
||||
"""
|
||||
if feedback_type != 'missed_relation' or not missed_relation:
|
||||
return None
|
||||
|
||||
relation_name = missed_relation.lower().strip().replace(' ', '_').replace('-', '_')
|
||||
|
||||
# 检查是否已存在
|
||||
if relation_name in self.all_patterns:
|
||||
print(f" 模式已存在: {relation_name}")
|
||||
return self.all_patterns[relation_name]
|
||||
|
||||
# 生成正则模式
|
||||
pattern = self._generate_pattern_from_text(missed_relation, original_text)
|
||||
|
||||
if not pattern:
|
||||
print(f" 无法生成模式: {missed_relation}")
|
||||
return None
|
||||
|
||||
# 创建新模式
|
||||
new_pattern = RelationPattern(
|
||||
relation=relation_name,
|
||||
pattern=pattern,
|
||||
source_type='gene',
|
||||
target_type='target',
|
||||
category='LEARNED',
|
||||
confidence=0.8, # 学习模式的初始置信度
|
||||
source='learned',
|
||||
learned_from=feedback_id,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
usage_count=0
|
||||
)
|
||||
|
||||
# 保存
|
||||
self.learned_patterns[relation_name] = new_pattern
|
||||
self.all_patterns[relation_name] = new_pattern
|
||||
self._save_learned_patterns()
|
||||
|
||||
print(f"✅ 学习到新模式: {relation_name}")
|
||||
print(f" 模式: {pattern}")
|
||||
|
||||
return new_pattern
|
||||
|
||||
def _generate_pattern_from_text(self, relation: str, text: str) -> Optional[str]:
|
||||
"""从文本生成关系正则模式"""
|
||||
try:
|
||||
# 基因名模式
|
||||
gene_pattern = r'[A-Z][a-z]{1,2}[A-Z]{1,4}\d{1,4}[a-z]?'
|
||||
|
||||
# 转义关系词
|
||||
escaped_relation = re.escape(relation.lower())
|
||||
|
||||
# 生成模式: (基因) + 关系 + (目标)
|
||||
# 支持变体: modulates, modulate, modulated
|
||||
pattern = rf'(\b{gene_pattern}\b).*?{escaped_relation}[sd]?\b.*?(.+?)(?:\.|,|;|$)'
|
||||
|
||||
# 验证模式是否有效
|
||||
re.compile(pattern)
|
||||
|
||||
return pattern
|
||||
except Exception as e:
|
||||
print(f" 模式生成失败: {e}")
|
||||
return None
|
||||
|
||||
def update_pattern_confidence(self, relation_name: str, success: bool):
|
||||
"""更新模式置信度"""
|
||||
if relation_name not in self.learned_patterns:
|
||||
return
|
||||
|
||||
pattern = self.learned_patterns[relation_name]
|
||||
pattern.usage_count += 1
|
||||
|
||||
# 简单的置信度更新算法
|
||||
if success:
|
||||
pattern.confidence = min(1.0, pattern.confidence + 0.05)
|
||||
else:
|
||||
pattern.confidence = max(0.5, pattern.confidence - 0.1)
|
||||
|
||||
self._save_learned_patterns()
|
||||
|
||||
def get_pattern_by_category(self, category: str) -> List[RelationPattern]:
|
||||
"""按分类获取模式"""
|
||||
return [p for p in self.all_patterns.values() if p.category == category]
|
||||
|
||||
def export_patterns_report(self, output_path: Path):
|
||||
"""导出模式报告"""
|
||||
stats = self.get_pattern_stats()
|
||||
|
||||
report = f"""
|
||||
# V14 关系模式报告
|
||||
|
||||
生成时间: {datetime.now().isoformat()}
|
||||
|
||||
## 统计概览
|
||||
|
||||
- PC模式: {stats['pc_patterns']} 种
|
||||
- 学习模式: {stats['learned_patterns']} 种
|
||||
- 总模式: {stats['total_patterns']} 种
|
||||
|
||||
## PC模式分类分布
|
||||
|
||||
"""
|
||||
for cat, count in sorted(stats['category_distribution'].items(), key=lambda x: -x[1]):
|
||||
report += f"- {cat}: {count} 种\n"
|
||||
|
||||
if stats['learned_relations']:
|
||||
report += f"\n## 学习到的模式\n\n"
|
||||
for rel in stats['learned_relations']:
|
||||
p = self.learned_patterns[rel]
|
||||
report += f"- **{rel}** (置信度: {p.confidence:.2f}, 使用: {p.usage_count}次)\n"
|
||||
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
f.write(report)
|
||||
|
||||
print(f"✅ 模式报告已导出: {output_path}")
|
||||
|
||||
|
||||
# 全局单例
|
||||
_pc_patterns_loader: Optional[PCPatternsLoader] = None
|
||||
|
||||
def get_pc_patterns_loader() -> PCPatternsLoader:
|
||||
"""获取PC模式加载器单例"""
|
||||
global _pc_patterns_loader
|
||||
if _pc_patterns_loader is None:
|
||||
_pc_patterns_loader = PCPatternsLoader()
|
||||
return _pc_patterns_loader
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 测试
|
||||
print("="*80)
|
||||
print("PC模式加载器测试")
|
||||
print("="*80)
|
||||
|
||||
loader = PCPatternsLoader()
|
||||
|
||||
# 显示统计
|
||||
stats = loader.get_pattern_stats()
|
||||
print(f"\n模式统计:")
|
||||
print(f" PC模式: {stats['pc_patterns']}")
|
||||
print(f" 学习模式: {stats['learned_patterns']}")
|
||||
print(f" 总模式: {stats['total_patterns']}")
|
||||
|
||||
# 显示分类
|
||||
print(f"\n前5个分类:")
|
||||
for cat, count in list(stats['category_distribution'].items())[:5]:
|
||||
print(f" {cat}: {count}")
|
||||
|
||||
# 测试学习
|
||||
print(f"\n测试学习新模式:")
|
||||
new_pattern = loader.learn_from_feedback(
|
||||
feedback_type='missed_relation',
|
||||
original_text='SUC5 modulates sugar uptake in leaves.',
|
||||
missed_relation='modulates',
|
||||
feedback_id='test_001'
|
||||
)
|
||||
|
||||
if new_pattern:
|
||||
print(f" 学习成功: {new_pattern.relation}")
|
||||
|
||||
# 导出报告
|
||||
loader.export_patterns_report(Path('./patterns_report.md'))
|
||||
@@ -0,0 +1,501 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import json
|
||||
import re
|
||||
from html.parser import HTMLParser
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Optional
|
||||
|
||||
from .sectionizer_v16_1_1 import classify_heading, is_major_heading, is_subheading
|
||||
|
||||
|
||||
REFERENCE_HEADING = re.compile(
|
||||
r"^\s*(?:references?|bibliography|literature cited)\s*$", re.I
|
||||
)
|
||||
IMAGE_PLACEHOLDER = re.compile(
|
||||
r"^\s*(?:!\[[^\]]*\]\([^)]*\)|\[image\]|<image[^>]*>|image\s+placeholder)\s*$",
|
||||
re.I,
|
||||
)
|
||||
FORMULA = re.compile(r"(?:\$[^$]+\$|\\(?:mathrm|mathbf|mathit|text)\s*\{)")
|
||||
FIGURE_CAPTION = re.compile(
|
||||
r"^\s*(?:supplementary\s+)?fig(?:ure)?\.?\s*[A-Za-z]?\d+\b", re.I
|
||||
)
|
||||
HYPHENATED_BREAK = re.compile(r"(?<=[A-Za-z])-\s*\n\s*(?=[a-z])")
|
||||
SOFT_BREAK = re.compile(r"(?<![.!?:;])\n(?=\S)")
|
||||
|
||||
|
||||
def is_figure_panel_fragment(text: str) -> bool:
|
||||
"""Return true for isolated panel labels emitted as fake captions by MinerU."""
|
||||
value = clean_block_text(text)
|
||||
if not value or len(value) > 30 or re.search(r"\bfig(?:ure)?\b", value, re.I):
|
||||
return False
|
||||
if re.search(r"[a-z]{3,}", value):
|
||||
return False
|
||||
tokens = re.findall(r"[A-Za-z]+|\d+|[\u4e00-\u9fff]+", value)
|
||||
return bool(tokens) and all(
|
||||
len(token) <= 2
|
||||
or token.isdigit()
|
||||
or re.fullmatch(r"Cas", token, re.I)
|
||||
for token in tokens
|
||||
)
|
||||
|
||||
|
||||
class _TableTextParser(HTMLParser):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.rows: List[List[str]] = []
|
||||
self._row: List[str] = []
|
||||
self._cell: List[str] = []
|
||||
self._in_cell = False
|
||||
|
||||
def handle_starttag(self, tag: str, attrs) -> None:
|
||||
tag = tag.lower()
|
||||
if tag == "tr":
|
||||
self._row = []
|
||||
elif tag in {"td", "th"}:
|
||||
self._cell = []
|
||||
self._in_cell = True
|
||||
|
||||
def handle_data(self, data: str) -> None:
|
||||
if self._in_cell:
|
||||
self._cell.append(data)
|
||||
|
||||
def handle_endtag(self, tag: str) -> None:
|
||||
tag = tag.lower()
|
||||
if tag in {"td", "th"}:
|
||||
value = re.sub(r"\s+", " ", "".join(self._cell)).strip()
|
||||
self._row.append(value)
|
||||
self._in_cell = False
|
||||
elif tag == "tr" and self._row:
|
||||
self.rows.append(self._row)
|
||||
|
||||
|
||||
def _as_text(value: Any) -> str:
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, list):
|
||||
return " ".join(_as_text(item) for item in value if _as_text(item))
|
||||
if value is None:
|
||||
return ""
|
||||
return str(value)
|
||||
|
||||
|
||||
def clean_block_text(value: Any) -> str:
|
||||
text = _as_text(value).replace("\r\n", "\n").replace("\r", "\n")
|
||||
text = HYPHENATED_BREAK.sub("", text)
|
||||
text = SOFT_BREAK.sub(" ", text)
|
||||
return re.sub(r"\s+", " ", text).strip()
|
||||
|
||||
|
||||
def render_table(table_body: str, caption: str = "", footnote: str = "") -> str:
|
||||
parser = _TableTextParser()
|
||||
try:
|
||||
parser.feed(table_body or "")
|
||||
except Exception:
|
||||
parser.rows = []
|
||||
lines: List[str] = []
|
||||
if caption:
|
||||
lines.append(f"Table caption: {clean_block_text(caption)}")
|
||||
if parser.rows:
|
||||
header = " | ".join(parser.rows[0])
|
||||
lines.append(f"Table header: {header}")
|
||||
for row in parser.rows[1:]:
|
||||
lines.append("Table row: " + " | ".join(row))
|
||||
elif table_body:
|
||||
lines.append("Table content: " + clean_block_text(re.sub(r"<[^>]+>", " ", table_body)))
|
||||
if footnote:
|
||||
lines.append(f"Table footnote: {clean_block_text(footnote)}")
|
||||
return "\n".join(line for line in lines if line.strip())
|
||||
|
||||
|
||||
def looks_like_reference_list(items: List[Any]) -> bool:
|
||||
values = [clean_block_text(item) for item in items if clean_block_text(item)]
|
||||
if not values:
|
||||
return False
|
||||
citation_like = sum(
|
||||
bool(
|
||||
re.search(r"\b(?:19|20)\d{2}\b", value)
|
||||
and (
|
||||
re.search(r"\bdoi\b|https?://doi\.org", value, re.I)
|
||||
or re.search(r"\b(?:vol|pp|journal|plant|science|biol|genet|fish|aquac|mar|zool)\b", value, re.I)
|
||||
)
|
||||
)
|
||||
for value in values
|
||||
)
|
||||
return citation_like >= max(2, len(values) // 2)
|
||||
|
||||
|
||||
class MinerUBlockBuilderV16_1_2:
|
||||
"""Build ordered semantic blocks from MinerU content_list with MD fallback."""
|
||||
|
||||
def __init__(self, output_dir: Path):
|
||||
self.output_dir = Path(output_dir)
|
||||
|
||||
def build(
|
||||
self,
|
||||
mineru_dir: Optional[Path],
|
||||
paper_id: str,
|
||||
fallback_text: str = "",
|
||||
) -> List[Dict[str, Any]]:
|
||||
content_path = self._find_content_list(mineru_dir, paper_id)
|
||||
blocks: List[Dict[str, Any]] = []
|
||||
source = "content_list"
|
||||
if content_path:
|
||||
try:
|
||||
payload = json.loads(content_path.read_text(encoding="utf-8"))
|
||||
if isinstance(payload, list):
|
||||
blocks = self._from_content_list(payload, paper_id)
|
||||
except Exception as exc:
|
||||
source = f"content_list_failed:{exc}"
|
||||
if not blocks:
|
||||
source = "markdown_or_parser_text"
|
||||
text = self._load_markdown(mineru_dir, paper_id) or fallback_text
|
||||
blocks = self._from_markdown(text, paper_id)
|
||||
self._write_outputs(blocks, paper_id, source, content_path)
|
||||
return blocks
|
||||
|
||||
def reconstruct_text(self, blocks: Iterable[Dict[str, Any]]) -> str:
|
||||
parts = []
|
||||
for block in blocks:
|
||||
if block.get("is_reference") or block.get("is_image_placeholder"):
|
||||
continue
|
||||
text = block.get("clean_text") or ""
|
||||
if text:
|
||||
parts.append(text)
|
||||
return "\n\n".join(parts)
|
||||
|
||||
def _find_content_list(self, mineru_dir: Optional[Path], paper_id: str) -> Optional[Path]:
|
||||
if not mineru_dir:
|
||||
return None
|
||||
root = Path(mineru_dir)
|
||||
direct = root / f"{paper_id}_content_list.json"
|
||||
if direct.exists():
|
||||
return direct
|
||||
matches = sorted(root.glob("*_content_list.json"))
|
||||
if matches:
|
||||
return matches[0]
|
||||
matches = sorted(root.rglob("*_content_list.json"))
|
||||
return matches[0] if matches else None
|
||||
|
||||
def _load_markdown(self, mineru_dir: Optional[Path], paper_id: str) -> str:
|
||||
if not mineru_dir:
|
||||
return ""
|
||||
root = Path(mineru_dir)
|
||||
direct = root / f"{paper_id}.md"
|
||||
if direct.exists():
|
||||
return direct.read_text(encoding="utf-8", errors="ignore")
|
||||
matches = sorted(root.glob("*.md")) or sorted(root.rglob("*.md"))
|
||||
return matches[0].read_text(encoding="utf-8", errors="ignore") if matches else ""
|
||||
|
||||
def _base_block(
|
||||
self,
|
||||
paper_id: str,
|
||||
order: int,
|
||||
page: Optional[int],
|
||||
mineru_type: str,
|
||||
raw_text: str,
|
||||
block_type: str,
|
||||
**extra: Any,
|
||||
) -> Dict[str, Any]:
|
||||
clean = clean_block_text(raw_text)
|
||||
heading_type = classify_heading(clean)
|
||||
is_major = block_type == "heading" and is_major_heading(clean)
|
||||
is_sub = block_type == "subheading" or (
|
||||
block_type == "heading" and not is_major and is_subheading(clean)
|
||||
)
|
||||
is_reference = bool(extra.get("is_reference") or heading_type == "references")
|
||||
is_panel_fragment = (
|
||||
block_type == "figure_caption" and is_figure_panel_fragment(clean)
|
||||
)
|
||||
is_administrative = bool(extra.get("is_administrative"))
|
||||
is_image = block_type == "image_placeholder"
|
||||
keep = (
|
||||
bool(clean)
|
||||
and not is_reference
|
||||
and not is_administrative
|
||||
and not is_panel_fragment
|
||||
and not is_image
|
||||
and block_type not in {"heading", "subheading"}
|
||||
)
|
||||
reason = (
|
||||
"references_after_heading"
|
||||
if is_reference
|
||||
else "administrative_section"
|
||||
if is_administrative
|
||||
else "figure_panel_label_only"
|
||||
if is_panel_fragment
|
||||
else "image_placeholder"
|
||||
if is_image
|
||||
else "empty_after_cleaning"
|
||||
if not clean
|
||||
else "heading_metadata_only"
|
||||
if block_type in {"heading", "subheading"}
|
||||
else ""
|
||||
)
|
||||
return {
|
||||
"block_id": f"{paper_id}_b{order:06d}",
|
||||
"paper_id": paper_id,
|
||||
"page": page,
|
||||
"order": order,
|
||||
"mineru_type": mineru_type,
|
||||
"v16_block_type": block_type,
|
||||
"raw_text": raw_text,
|
||||
"clean_text": clean,
|
||||
"char_count": len(clean),
|
||||
"is_major_heading": is_major,
|
||||
"is_subheading": is_sub,
|
||||
"section_type": heading_type if is_major else "",
|
||||
"is_reference": is_reference,
|
||||
"is_figure_caption": block_type == "figure_caption",
|
||||
"is_table_block": block_type == "table_block",
|
||||
"is_image_placeholder": is_image,
|
||||
"is_formula_heavy": len(FORMULA.findall(clean)) >= 3,
|
||||
"is_administrative": is_administrative,
|
||||
"is_figure_panel_fragment": is_panel_fragment,
|
||||
"keep_for_llm": keep,
|
||||
"drop_reason": reason,
|
||||
**{k: v for k, v in extra.items() if k != "is_reference"},
|
||||
}
|
||||
|
||||
def _from_content_list(
|
||||
self, payload: List[Dict[str, Any]], paper_id: str
|
||||
) -> List[Dict[str, Any]]:
|
||||
blocks: List[Dict[str, Any]] = []
|
||||
in_references = False
|
||||
in_administrative = False
|
||||
order = 0
|
||||
|
||||
def append(
|
||||
item: Dict[str, Any],
|
||||
raw: str,
|
||||
block_type: str,
|
||||
mineru_type: Optional[str] = None,
|
||||
is_reference: bool = False,
|
||||
) -> None:
|
||||
nonlocal order, in_references, in_administrative
|
||||
order += 1
|
||||
clean = clean_block_text(raw)
|
||||
heading_class = classify_heading(clean)
|
||||
if heading_class == "references":
|
||||
in_references = True
|
||||
elif heading_class == "administrative":
|
||||
in_administrative = True
|
||||
elif heading_class and heading_class not in {
|
||||
"subheading_inside_parent",
|
||||
"data_availability",
|
||||
}:
|
||||
in_administrative = False
|
||||
elif heading_class == "data_availability":
|
||||
in_administrative = False
|
||||
ref = in_references or is_reference
|
||||
blocks.append(
|
||||
self._base_block(
|
||||
paper_id,
|
||||
order,
|
||||
item.get("page_idx"),
|
||||
mineru_type or str(item.get("type") or "unknown"),
|
||||
raw,
|
||||
block_type,
|
||||
is_reference=ref,
|
||||
is_administrative=in_administrative,
|
||||
bbox=item.get("bbox"),
|
||||
text_level=item.get("text_level"),
|
||||
)
|
||||
)
|
||||
|
||||
for item in payload:
|
||||
typ = str(item.get("type") or "").lower()
|
||||
if typ in {"header", "footer", "page_number", "page_footnote"}:
|
||||
append(item, _as_text(item.get("text")), "image_placeholder", typ)
|
||||
continue
|
||||
if typ == "text":
|
||||
raw = _as_text(item.get("text"))
|
||||
level = item.get("text_level")
|
||||
classification = classify_heading(raw)
|
||||
clean = clean_block_text(raw)
|
||||
heading_like = bool(
|
||||
level is not None
|
||||
and len(clean) <= 220
|
||||
and len(clean.split()) <= 24
|
||||
and not re.search(r"[.!?]\s*$", clean)
|
||||
)
|
||||
block_type = (
|
||||
"figure_caption"
|
||||
if FIGURE_CAPTION.match(clean)
|
||||
else
|
||||
"heading"
|
||||
if classification not in {None, "subheading_inside_parent"}
|
||||
else "subheading"
|
||||
if heading_like or classification == "subheading_inside_parent"
|
||||
else "paragraph"
|
||||
)
|
||||
append(item, raw, block_type, typ)
|
||||
continue
|
||||
if typ == "list":
|
||||
list_items = item.get("list_items") or []
|
||||
raw = "\n".join(_as_text(v) for v in list_items)
|
||||
append(
|
||||
item,
|
||||
raw,
|
||||
"paragraph",
|
||||
typ,
|
||||
is_reference=(
|
||||
str(item.get("sub_type") or "").lower() == "ref_text"
|
||||
and looks_like_reference_list(list_items)
|
||||
),
|
||||
)
|
||||
continue
|
||||
if typ in {"image", "chart"}:
|
||||
captions = item.get("image_caption") or item.get("chart_caption") or []
|
||||
caption = "\n".join(_as_text(v) for v in captions if _as_text(v))
|
||||
if caption:
|
||||
append(item, caption, "figure_caption", typ)
|
||||
else:
|
||||
append(item, "[image]", "image_placeholder", typ)
|
||||
continue
|
||||
if typ == "table":
|
||||
table_text = render_table(
|
||||
_as_text(item.get("table_body")),
|
||||
_as_text(item.get("table_caption")),
|
||||
_as_text(item.get("table_footnote")),
|
||||
)
|
||||
append(item, table_text, "table_block", typ)
|
||||
continue
|
||||
raw = (
|
||||
_as_text(item.get("text"))
|
||||
or _as_text(item.get("content"))
|
||||
or _as_text(item.get("latex"))
|
||||
)
|
||||
if raw:
|
||||
append(item, raw, "paragraph", typ)
|
||||
return blocks
|
||||
|
||||
def _from_markdown(self, text: str, paper_id: str) -> List[Dict[str, Any]]:
|
||||
blocks: List[Dict[str, Any]] = []
|
||||
paragraphs: List[str] = []
|
||||
order = 0
|
||||
in_references = False
|
||||
in_administrative = False
|
||||
|
||||
def emit(raw: str, forced_type: str = "") -> None:
|
||||
nonlocal order, in_references, in_administrative
|
||||
if not raw.strip():
|
||||
return
|
||||
order += 1
|
||||
clean = clean_block_text(raw)
|
||||
heading_match = re.match(r"^\s{0,3}#{1,6}\s+(.+?)\s*$", raw.strip())
|
||||
heading_text = heading_match.group(1) if heading_match else clean
|
||||
classification = classify_heading(heading_text)
|
||||
if classification == "references":
|
||||
in_references = True
|
||||
elif classification == "administrative":
|
||||
in_administrative = True
|
||||
elif classification and classification not in {
|
||||
"subheading_inside_parent",
|
||||
"data_availability",
|
||||
}:
|
||||
in_administrative = False
|
||||
elif classification == "data_availability":
|
||||
in_administrative = False
|
||||
if forced_type:
|
||||
block_type = forced_type
|
||||
elif heading_match:
|
||||
block_type = (
|
||||
"heading"
|
||||
if classification not in {None, "subheading_inside_parent"}
|
||||
else "subheading"
|
||||
)
|
||||
elif re.match(r"^(?:supplementary\s+)?fig(?:ure)?\.?\s*\w+", clean, re.I):
|
||||
block_type = "figure_caption"
|
||||
elif "<table" in clean.lower() or re.match(r"^\s*\|.+\|\s*$", raw):
|
||||
block_type = "table_block"
|
||||
elif IMAGE_PLACEHOLDER.match(clean):
|
||||
block_type = "image_placeholder"
|
||||
else:
|
||||
block_type = "paragraph"
|
||||
blocks.append(
|
||||
self._base_block(
|
||||
paper_id,
|
||||
order,
|
||||
None,
|
||||
"markdown",
|
||||
heading_text if heading_match else raw,
|
||||
block_type,
|
||||
is_reference=in_references,
|
||||
is_administrative=in_administrative,
|
||||
)
|
||||
)
|
||||
|
||||
for line in (text or "").splitlines():
|
||||
special = (
|
||||
not line.strip()
|
||||
or re.match(r"^\s{0,3}#{1,6}\s+", line)
|
||||
or IMAGE_PLACEHOLDER.match(line.strip())
|
||||
or "<table" in line.lower()
|
||||
or re.match(r"^\s*\|.+\|\s*$", line)
|
||||
)
|
||||
if special and paragraphs:
|
||||
emit("\n".join(paragraphs))
|
||||
paragraphs = []
|
||||
if not line.strip():
|
||||
continue
|
||||
if special:
|
||||
emit(line)
|
||||
else:
|
||||
paragraphs.append(line)
|
||||
if paragraphs:
|
||||
emit("\n".join(paragraphs))
|
||||
return blocks
|
||||
|
||||
def _write_outputs(
|
||||
self,
|
||||
blocks: List[Dict[str, Any]],
|
||||
paper_id: str,
|
||||
source: str,
|
||||
content_path: Optional[Path],
|
||||
) -> None:
|
||||
self.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
raw_path = self.output_dir / "01_mineru_blocks_raw.jsonl"
|
||||
clean_path = self.output_dir / "02_mineru_blocks_clean.jsonl"
|
||||
with raw_path.open("w", encoding="utf-8") as raw_handle, clean_path.open(
|
||||
"w", encoding="utf-8"
|
||||
) as clean_handle:
|
||||
for block in blocks:
|
||||
raw_handle.write(json.dumps(block, ensure_ascii=False) + "\n")
|
||||
clean_handle.write(json.dumps(block, ensure_ascii=False) + "\n")
|
||||
with (self.output_dir / "dropped_block_audit.csv").open(
|
||||
"w", newline="", encoding="utf-8"
|
||||
) as handle:
|
||||
fields = [
|
||||
"block_id",
|
||||
"page",
|
||||
"order",
|
||||
"mineru_type",
|
||||
"v16_block_type",
|
||||
"drop_reason",
|
||||
"raw_text",
|
||||
]
|
||||
writer = csv.DictWriter(handle, fieldnames=fields)
|
||||
writer.writeheader()
|
||||
for block in blocks:
|
||||
if not block.get("keep_for_llm"):
|
||||
writer.writerow({key: block.get(key, "") for key in fields})
|
||||
summary = {
|
||||
"paper_id": paper_id,
|
||||
"block_source": source,
|
||||
"content_list_path": str(content_path or ""),
|
||||
"block_count": len(blocks),
|
||||
"kept_block_count": sum(bool(b.get("keep_for_llm")) for b in blocks),
|
||||
"reference_block_count": sum(bool(b.get("is_reference")) for b in blocks),
|
||||
"figure_caption_block_count": sum(
|
||||
bool(b.get("is_figure_caption")) for b in blocks
|
||||
),
|
||||
"table_block_count": sum(bool(b.get("is_table_block")) for b in blocks),
|
||||
"image_placeholder_count": sum(
|
||||
bool(b.get("is_image_placeholder")) for b in blocks
|
||||
),
|
||||
}
|
||||
(self.output_dir / "mineru_block_summary.json").write_text(
|
||||
json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
|
||||
)
|
||||
+293
@@ -0,0 +1,293 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Phase 1: 文档解析
|
||||
整合v12的marker-pdf章节识别 + PyMuPDF章节补充
|
||||
"""
|
||||
|
||||
import re
|
||||
import os
|
||||
import shlex
|
||||
import subprocess
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Tuple, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
import fitz # PyMuPDF
|
||||
|
||||
try:
|
||||
from ...project_config import conda_profile
|
||||
except Exception: # pragma: no cover - supports direct script execution
|
||||
def conda_profile() -> Path:
|
||||
profile = os.environ.get("PLANTAI_CONDA_PROFILE", "").strip()
|
||||
if profile:
|
||||
return Path(profile)
|
||||
root = os.environ.get("PLANTAI_CONDA_ROOT", "").strip()
|
||||
return (Path(root) if root else Path.home() / "miniconda3") / "etc/profile.d/conda.sh"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Section:
|
||||
"""章节信息 - V14.2: 添加position和end字段,与phase1_document_parser_mineru兼容"""
|
||||
title: str
|
||||
level: int
|
||||
parent: Optional[str]
|
||||
start_page: int
|
||||
end_page: int
|
||||
section_type: str
|
||||
position: int = 0 # 在文本中的起始位置
|
||||
end: int = 0 # 在文本中的结束位置
|
||||
|
||||
|
||||
class DocumentParser:
|
||||
"""文档解析器 - 整合v12能力"""
|
||||
|
||||
def __init__(self, marker_env: str = "markerpdf"):
|
||||
self.marker_env = marker_env
|
||||
self.marker_available = self._check_marker()
|
||||
|
||||
def _check_marker(self) -> bool:
|
||||
"""检查marker-pdf是否可用(在markerpdf环境中)"""
|
||||
try:
|
||||
profile = shlex.quote(str(conda_profile()))
|
||||
env_name = shlex.quote(self.marker_env)
|
||||
# 检查conda环境和markerpdf环境是否存在
|
||||
cmd = [
|
||||
"bash", "-c",
|
||||
f"source {profile} && conda env list | grep {env_name}"
|
||||
]
|
||||
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||
return result.returncode == 0 and self.marker_env in result.stdout
|
||||
except:
|
||||
return False
|
||||
|
||||
def parse(self, pdf_path: Path) -> Tuple[List[Section], str]:
|
||||
"""
|
||||
解析PDF文档
|
||||
|
||||
Returns:
|
||||
(sections, full_text)
|
||||
"""
|
||||
print(" Step 1.1: 提取PDF文本...")
|
||||
full_text = self._extract_text_pymupdf(pdf_path)
|
||||
print(f" ✓ 提取 {len(full_text)} 字符")
|
||||
|
||||
print(" Step 1.2: 识别章节...")
|
||||
sections = None
|
||||
|
||||
# 尝试使用marker-pdf
|
||||
if self.marker_available:
|
||||
try:
|
||||
sections = self._extract_sections_marker(pdf_path)
|
||||
if sections:
|
||||
print(f" ✓ marker-pdf识别到 {len(sections)} 个章节")
|
||||
except Exception as e:
|
||||
print(f" ⚠️ marker-pdf失败: {e}")
|
||||
|
||||
# 回退到PyMuPDF
|
||||
if sections is None:
|
||||
sections = self._extract_sections_pymupdf(pdf_path)
|
||||
print(f" ⚠️ PyMuPDF回退识别到 {len(sections)} 个章节")
|
||||
|
||||
return sections, full_text
|
||||
|
||||
def _extract_text_pymupdf(self, pdf_path: Path) -> str:
|
||||
"""使用PyMuPDF提取全文"""
|
||||
doc = fitz.open(str(pdf_path))
|
||||
text = ""
|
||||
for page in doc:
|
||||
text += page.get_text() + "\n"
|
||||
doc.close()
|
||||
return text
|
||||
|
||||
def _extract_sections_marker(self, pdf_path: Path) -> Optional[List[Section]]:
|
||||
"""使用marker-pdf提取章节(v12方法,conda activate markerpdf)"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
output_dir = Path(tmpdir) / "marker_output"
|
||||
output_dir.mkdir()
|
||||
|
||||
# 使用conda activate markerpdf环境(v12方法)
|
||||
profile = shlex.quote(str(conda_profile()))
|
||||
env_name = shlex.quote(self.marker_env)
|
||||
quoted_pdf = shlex.quote(str(pdf_path))
|
||||
quoted_output = shlex.quote(str(output_dir))
|
||||
cmd = [
|
||||
"bash", "-c",
|
||||
f"source {profile} && "
|
||||
f"conda activate {env_name} && "
|
||||
f"marker_single {quoted_pdf} --output_dir {quoted_output}"
|
||||
]
|
||||
|
||||
print(f" 调用marker-pdf (conda activate markerpdf)...")
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=300
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
print(f" ⚠️ marker-pdf失败: {result.stderr[:200]}")
|
||||
return None
|
||||
|
||||
# marker-pdf输出到子目录
|
||||
pdf_name = pdf_path.stem
|
||||
subdir = output_dir / pdf_name
|
||||
|
||||
if subdir.exists():
|
||||
md_files = list(subdir.glob("*.md"))
|
||||
else:
|
||||
md_files = list(output_dir.glob("*.md"))
|
||||
|
||||
if not md_files:
|
||||
print(" ⚠️ marker-pdf未生成markdown文件")
|
||||
return None
|
||||
|
||||
# 解析markdown
|
||||
md_content = md_files[0].read_text(encoding='utf-8')
|
||||
sections = self._parse_marker_markdown(md_content)
|
||||
|
||||
return sections if sections else None
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
print(" ⚠️ marker-pdf超时")
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f" ⚠️ marker-pdf错误: {e}")
|
||||
return None
|
||||
|
||||
def _parse_marker_markdown(self, md_content: str) -> List[Section]:
|
||||
"""解析marker的markdown输出(v12方法)"""
|
||||
sections = []
|
||||
lines = md_content.split('\n')
|
||||
heading_pattern = r'^(#{1,3})\s+(.+)$'
|
||||
|
||||
current_parent = None
|
||||
|
||||
for i, line in enumerate(lines):
|
||||
match = re.match(heading_pattern, line)
|
||||
if match:
|
||||
level = len(match.group(1))
|
||||
title = match.group(2).strip()
|
||||
|
||||
# 更新父章节
|
||||
if level == 1:
|
||||
current_parent = title
|
||||
parent = None
|
||||
else:
|
||||
parent = current_parent
|
||||
|
||||
sections.append(Section(
|
||||
title=title,
|
||||
level=level,
|
||||
parent=parent,
|
||||
start_page=0,
|
||||
end_page=0,
|
||||
section_type=title.lower().replace(' ', '_').replace('-', '_')
|
||||
))
|
||||
|
||||
return sections
|
||||
|
||||
def _extract_sections_pymupdf(self, pdf_path: Path) -> List[Section]:
|
||||
"""使用PyMuPDF提取章节(回退方法)"""
|
||||
sections = []
|
||||
|
||||
# 常见章节标题模式
|
||||
section_patterns = [
|
||||
(r'\n\s*Abstract\s*\n', 'Abstract', 1),
|
||||
(r'\n\s*Introduction\s*\n', 'Introduction', 1),
|
||||
(r'\n\s*Materials?\s+and\s+Methods?\s*\n', 'Methods', 1),
|
||||
(r'\n\s*Results?\s*\n', 'Results', 1),
|
||||
(r'\n\s*Discussion\s*\n', 'Discussion', 1),
|
||||
(r'\n\s*Conclusion\s*\n', 'Conclusion', 1),
|
||||
(r'\n\s*References?\s*\n', 'References', 1),
|
||||
]
|
||||
|
||||
text = self._extract_text_pymupdf(pdf_path)
|
||||
|
||||
for pattern, title, level in section_patterns:
|
||||
if re.search(pattern, text, re.IGNORECASE):
|
||||
sections.append(Section(
|
||||
title=title,
|
||||
level=level,
|
||||
parent=None,
|
||||
start_page=0,
|
||||
end_page=0,
|
||||
section_type=title.lower()
|
||||
))
|
||||
|
||||
return sections
|
||||
|
||||
def extract_section_text(self, full_text: str, section: Section) -> str:
|
||||
"""
|
||||
从全文中提取章节文本
|
||||
|
||||
修复:处理marker-pdf返回的markdown格式标题(如**RESULTS**)
|
||||
"""
|
||||
import re
|
||||
|
||||
# 清理章节标题(去除markdown格式)
|
||||
clean_title = re.sub(r'\*+', '', section.title).strip()
|
||||
|
||||
# 常见章节标题列表(用于定位下一个章节)
|
||||
common_sections = [
|
||||
'abstract', 'introduction', 'materials', 'methods', 'results',
|
||||
'discussion', 'conclusion', 'references', 'acknowledgements',
|
||||
'supplemental', 'additional information'
|
||||
]
|
||||
|
||||
# 构建正则:匹配章节标题(支持markdown格式)
|
||||
# 模式1: 纯文本标题
|
||||
pattern1 = rf'\n\s*{re.escape(clean_title)}\s*\n'
|
||||
# 模式2: markdown加粗标题
|
||||
pattern2 = rf'\n\s*\*+\s*{re.escape(clean_title)}\s*\*+\s*\n'
|
||||
|
||||
# 尝试匹配章节标题
|
||||
match = None
|
||||
for pattern in [pattern2, pattern1]:
|
||||
match = re.search(pattern, full_text, re.IGNORECASE)
|
||||
if match:
|
||||
break
|
||||
|
||||
if not match:
|
||||
return ""
|
||||
|
||||
# 找到章节开始位置
|
||||
start_pos = match.end()
|
||||
|
||||
# 查找下一个章节的开始位置
|
||||
end_pos = len(full_text)
|
||||
for next_section in common_sections:
|
||||
if next_section.lower() == clean_title.lower():
|
||||
continue
|
||||
# 查找下一个章节标题
|
||||
next_pattern = rf'\n\s*(?:\*+\s*)?{re.escape(next_section)}\s*(?:\*+)?\s*\n'
|
||||
next_match = re.search(next_pattern, full_text[start_pos:], re.IGNORECASE)
|
||||
if next_match:
|
||||
potential_end = start_pos + next_match.start()
|
||||
if potential_end < end_pos:
|
||||
end_pos = potential_end
|
||||
|
||||
# 提取章节文本
|
||||
section_text = full_text[start_pos:end_pos].strip()
|
||||
|
||||
return section_text
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 测试
|
||||
parser = DocumentParser()
|
||||
test_pdf_value = os.environ.get("PLANTAI_TEST_PDF", "").strip()
|
||||
test_pdf = Path(test_pdf_value) if test_pdf_value else None
|
||||
|
||||
if test_pdf and test_pdf.exists():
|
||||
sections, text = parser.parse(test_pdf)
|
||||
print(f"\n测试完成!")
|
||||
print(f"章节数: {len(sections)}")
|
||||
print(f"文本长度: {len(text)}")
|
||||
for s in sections:
|
||||
print(f" - {s.title} (level {s.level})")
|
||||
else:
|
||||
print("设置 PLANTAI_TEST_PDF 后可运行本地解析测试")
|
||||
+566
@@ -0,0 +1,566 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Phase 1: 文档解析 - MinerU 版本
|
||||
使用 MinerU 作为主要解析器,替代 marker-pdf
|
||||
"""
|
||||
|
||||
import re
|
||||
import os
|
||||
import shlex
|
||||
import subprocess
|
||||
import tempfile
|
||||
import json
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Tuple, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
import fitz # PyMuPDF 作为回退
|
||||
|
||||
try:
|
||||
from ...project_config import conda_env_path, conda_profile
|
||||
except Exception: # pragma: no cover - supports direct script execution
|
||||
def _conda_root() -> Path:
|
||||
value = os.environ.get("PLANTAI_CONDA_ROOT", "").strip()
|
||||
if value:
|
||||
return Path(value)
|
||||
return Path.home() / "miniconda3"
|
||||
|
||||
def conda_env_path(env_name: str) -> Path:
|
||||
return _conda_root() / "envs" / env_name
|
||||
|
||||
def conda_profile() -> Path:
|
||||
value = os.environ.get("PLANTAI_CONDA_PROFILE", "").strip()
|
||||
if value:
|
||||
return Path(value)
|
||||
return _conda_root() / "etc/profile.d/conda.sh"
|
||||
|
||||
# Import text cleaning function from pdf_extractor
|
||||
try:
|
||||
from ...pdf_extractor import clean_text
|
||||
except ImportError:
|
||||
# Fallback: define clean_text locally if import fails
|
||||
def clean_text(text: str) -> str:
|
||||
"""Clean text to handle surrogate characters, invalid Unicode, and PDF line breaks"""
|
||||
if not text:
|
||||
return text
|
||||
|
||||
# Step 1: Replace invalid surrogate characters
|
||||
cleaned = []
|
||||
for char in text:
|
||||
code = ord(char)
|
||||
if 0xD800 <= code <= 0xDFFF:
|
||||
continue
|
||||
if code < 32 and code not in (9, 10, 13):
|
||||
continue
|
||||
cleaned.append(char)
|
||||
|
||||
text = ''.join(cleaned)
|
||||
|
||||
# Step 2: Fix hyphenated line breaks
|
||||
text = re.sub(r'(\w)-\n(\w)', r'\1\2', text)
|
||||
text = re.sub(r'(\w)-\r?\n(\w)', r'\1\2', text)
|
||||
|
||||
# Step 3: Fix soft line breaks within sentences (enhanced version)
|
||||
text = re.sub(r'(?<=[a-zA-Z.,;:!?])\s*\n\s*(?=[a-z])', ' ', text)
|
||||
text = re.sub(r'(?<=[a-zA-Z.,;:!?])\s*\n\s*(?=[A-Z][a-z])', ' ', text)
|
||||
text = re.sub(r'([A-Z]\.)\s*\n\s*([a-z])', r'\1 \2', text)
|
||||
|
||||
# Step 4: Normalize multiple newlines
|
||||
text = re.sub(r'\n{3,}', '\n\n', text)
|
||||
|
||||
return text
|
||||
|
||||
|
||||
@dataclass
|
||||
class Section:
|
||||
"""章节信息 - V13.1 添加文本位置字段"""
|
||||
title: str
|
||||
level: int
|
||||
parent: Optional[str]
|
||||
start_page: int
|
||||
end_page: int
|
||||
section_type: str
|
||||
position: int = 0 # 在文本中的起始位置
|
||||
end: int = 0 # 在文本中的结束位置
|
||||
|
||||
|
||||
class MinerUDocumentParser:
|
||||
"""MinerU 文档解析器"""
|
||||
|
||||
def __init__(self, conda_env: str = "minerupy312", api_url: str = "http://127.0.0.1:50215", use_api: bool = False):
|
||||
self.conda_env = conda_env
|
||||
self.api_url = api_url
|
||||
self.use_api = use_api
|
||||
self.mineru_available = self._check_mineru()
|
||||
self.api_available = self._check_api() if use_api else False
|
||||
|
||||
def _check_mineru(self) -> bool:
|
||||
"""检查 MinerU 是否可用 - V14.2优化: 使用快速检查"""
|
||||
try:
|
||||
# V15.3.4 Spark/portable mode: allow callers to point directly to
|
||||
# the MinerU executable instead of relying on a hard-coded conda
|
||||
# environment path from the original 239 host.
|
||||
env_cmd = (
|
||||
os.environ.get("V15_MINERU_CMD", "").strip()
|
||||
or os.environ.get("MINERU_CMD", "").strip()
|
||||
)
|
||||
if env_cmd and Path(env_cmd).exists():
|
||||
return True
|
||||
if shutil.which("mineru"):
|
||||
return True
|
||||
|
||||
# V14.2优化: 直接检查conda环境目录是否存在,避免conda命令开销
|
||||
env_path = conda_env_path(self.conda_env)
|
||||
if env_path.exists():
|
||||
return True
|
||||
|
||||
# 回退到conda命令检查
|
||||
profile = shlex.quote(str(conda_profile()))
|
||||
env_name = shlex.quote(self.conda_env)
|
||||
cmd = [
|
||||
"bash", "-c",
|
||||
f"source {profile} && conda env list | grep {env_name}"
|
||||
]
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
|
||||
return result.returncode == 0 and self.conda_env in result.stdout
|
||||
except subprocess.TimeoutExpired:
|
||||
print(f" ⚠️ 检查 MinerU 环境超时")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f" ⚠️ 检查 MinerU 环境失败: {e}")
|
||||
return False
|
||||
|
||||
def _check_api(self) -> bool:
|
||||
"""检查 MinerU API 服务是否可用"""
|
||||
try:
|
||||
import requests
|
||||
response = requests.get(f"{self.api_url}/docs", timeout=5)
|
||||
if response.status_code == 200:
|
||||
print(f" ✓ MinerU API 服务可用: {self.api_url}")
|
||||
return True
|
||||
else:
|
||||
print(f" ⚠️ MinerU API 服务返回状态码: {response.status_code}")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f" ⚠️ MinerU API 服务不可用: {e}")
|
||||
return False
|
||||
|
||||
def _check_api(self) -> bool:
|
||||
"""检查 MinerU API 是否可用"""
|
||||
try:
|
||||
import urllib.request
|
||||
urllib.request.urlopen(self.api_url, timeout=2)
|
||||
return True
|
||||
except:
|
||||
return False
|
||||
|
||||
def parse(self, pdf_path: Path, output_base_dir: Optional[Path] = None) -> Tuple[List[Section], str, Path]:
|
||||
"""
|
||||
解析 PDF 文档
|
||||
|
||||
Returns:
|
||||
(sections, full_text, mineru_output_dir)
|
||||
"""
|
||||
print(" Step 1.1: 使用 MinerU 提取 PDF...")
|
||||
|
||||
# 为每篇论文创建独立文件夹
|
||||
if output_base_dir is None:
|
||||
output_base_dir = pdf_path.parent / 'mineru'
|
||||
|
||||
paper_name = pdf_path.stem
|
||||
mineru_output_dir = output_base_dir / 'test' / paper_name / 'hybrid_auto'
|
||||
|
||||
# V14.2优化: 检查多个可能的输出路径
|
||||
# MinerU 输出可能在子目录中 (paper_name/hybrid_auto/ 或 paper_name/auto/)
|
||||
possible_content_files = [
|
||||
mineru_output_dir / f'{paper_name}_content_list.json',
|
||||
mineru_output_dir / paper_name / 'hybrid_auto' / f'{paper_name}_content_list.json',
|
||||
mineru_output_dir / paper_name / 'auto' / f'{paper_name}_content_list.json',
|
||||
]
|
||||
|
||||
content_list_file = None
|
||||
for f in possible_content_files:
|
||||
if f.exists():
|
||||
content_list_file = f
|
||||
break
|
||||
|
||||
if mineru_output_dir.exists() and content_list_file:
|
||||
print(f" ✓ 找到已有 MinerU 结果: {mineru_output_dir}")
|
||||
sections, full_text, actual_output_dir = self._parse_mineru_output(mineru_output_dir)
|
||||
min_chars = int(os.environ.get("V15_MINERU_MIN_TEXT_CHARS", "5000") or "5000")
|
||||
force_bad_cache_rerun = os.environ.get("V15_MINERU_RERUN_BAD_CACHE", "1").strip().lower() in {"1", "true", "yes", "on"}
|
||||
short_notice = re.search(r"\b(Retraction|Correction|Erratum)\s*:", full_text or "", re.I) is not None
|
||||
cache_looks_bad = (
|
||||
force_bad_cache_rerun
|
||||
and len(full_text or "") < min_chars
|
||||
and len(sections or []) == 0
|
||||
and not short_notice
|
||||
)
|
||||
if not cache_looks_bad:
|
||||
return sections, full_text, actual_output_dir
|
||||
print(
|
||||
f" ⚠️ 已有 MinerU 结果疑似坏缓存: text_chars={len(full_text or '')}, "
|
||||
f"sections={len(sections or [])}, min_chars={min_chars}; 将重新解析"
|
||||
)
|
||||
try:
|
||||
import shutil
|
||||
shutil.rmtree(mineru_output_dir)
|
||||
except Exception as cache_exc:
|
||||
print(f" ⚠️ 删除坏 MinerU 缓存失败: {cache_exc}")
|
||||
|
||||
# 调用 MinerU (使用 pipeline 模式,避免 vLLM)
|
||||
if self.mineru_available:
|
||||
try:
|
||||
success = self._run_mineru_pipeline(pdf_path, mineru_output_dir)
|
||||
if success:
|
||||
return self._parse_mineru_output(mineru_output_dir)
|
||||
except Exception as e:
|
||||
print(f" ⚠️ MinerU pipeline 失败: {e}")
|
||||
|
||||
# 禁止回退到 PyMuPDF,直接返回空结果
|
||||
print(" ❌ MinerU 失败,禁止回退到 PyMuPDF")
|
||||
return [], "", Path()
|
||||
|
||||
def _run_mineru(self, pdf_path: Path, output_dir: Path) -> bool:
|
||||
"""运行 MinerU (使用 pipeline 模式,避免 vLLM)"""
|
||||
return self._run_mineru_pipeline(pdf_path, output_dir)
|
||||
|
||||
def _run_mineru_modelscope(self, pdf_path: Path, output_dir: Path) -> bool:
|
||||
"""使用 modelscope 备选方案运行 MinerU"""
|
||||
profile = shlex.quote(str(conda_profile()))
|
||||
env_name = shlex.quote(self.conda_env)
|
||||
quoted_pdf = shlex.quote(str(pdf_path))
|
||||
quoted_output = shlex.quote(str(output_dir))
|
||||
cmd = [
|
||||
"bash", "-c",
|
||||
f"source {profile} && "
|
||||
f"conda activate {env_name} && "
|
||||
f"MINERU_MODEL_SOURCE=modelscope mineru -p {quoted_pdf} -o {quoted_output}"
|
||||
]
|
||||
|
||||
print(f" 执行: MINERU_MODEL_SOURCE=modelscope mineru -p {pdf_path} -o {output_dir}")
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd, capture_output=True, text=True, timeout=900 # V15修订: 从300秒增加到900秒(15分钟)
|
||||
)
|
||||
return result.returncode == 0
|
||||
except Exception as e:
|
||||
print(f" ⚠️ Modelscope 备选方案失败: {e}")
|
||||
return False
|
||||
|
||||
def _run_mineru_pipeline(self, pdf_path: Path, output_dir: Path) -> bool:
|
||||
"""使用 MinerU pipeline 模式运行(不使用 vLLM)"""
|
||||
import subprocess
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
backend = os.environ.get("V15_MINERU_BACKEND", "pipeline").strip() or "pipeline"
|
||||
method = os.environ.get("V15_MINERU_METHOD", "").strip()
|
||||
lang = os.environ.get("V15_MINERU_LANG", "").strip()
|
||||
use_api_env = os.environ.get("V15_MINERU_USE_API", "").strip().lower()
|
||||
use_api = self.use_api or use_api_env in {"1", "true", "yes", "on"}
|
||||
api_url = os.environ.get("V15_MINERU_API_URL", self.api_url).strip()
|
||||
extra_args = []
|
||||
if method:
|
||||
extra_args.extend(["-m", method])
|
||||
if lang:
|
||||
extra_args.extend(["-l", lang])
|
||||
if use_api and api_url:
|
||||
extra_args.extend(["--api-url", api_url])
|
||||
extra_arg_text = " ".join(extra_args)
|
||||
|
||||
env_cmd = (
|
||||
os.environ.get("V15_MINERU_CMD", "").strip()
|
||||
or os.environ.get("MINERU_CMD", "").strip()
|
||||
)
|
||||
mineru_cmd = env_cmd if env_cmd else (shutil.which("mineru") or "")
|
||||
if mineru_cmd:
|
||||
cmd = [
|
||||
mineru_cmd,
|
||||
"-p", str(pdf_path),
|
||||
"-o", str(output_dir),
|
||||
"-b", backend,
|
||||
*extra_args,
|
||||
]
|
||||
run_env = os.environ.copy()
|
||||
run_env["MINERU_MODEL_SOURCE"] = run_env.get("MINERU_MODEL_SOURCE", "modelscope")
|
||||
cmd_display = " ".join(cmd)
|
||||
else:
|
||||
profile = shlex.quote(str(conda_profile()))
|
||||
env_name = shlex.quote(self.conda_env)
|
||||
quoted_pdf = shlex.quote(str(pdf_path))
|
||||
quoted_output = shlex.quote(str(output_dir))
|
||||
quoted_backend = shlex.quote(backend)
|
||||
cmd = [
|
||||
"bash", "-c",
|
||||
f"source {profile} && "
|
||||
f"conda activate {env_name} && "
|
||||
f"export MINERU_MODEL_SOURCE=modelscope && "
|
||||
f"mineru -p {quoted_pdf} -o {quoted_output} -b {quoted_backend} {extra_arg_text}"
|
||||
]
|
||||
run_env = None
|
||||
cmd_display = f"mineru -p {pdf_path} -o {output_dir} -b {backend} {extra_arg_text}".strip()
|
||||
|
||||
print(f" 执行: {cmd_display}")
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd, capture_output=True, text=True, timeout=900, env=run_env # V15修订: 从300秒增加到900秒(15分钟)
|
||||
)
|
||||
try:
|
||||
(output_dir / "mineru_stdout.log").write_text(result.stdout or "", encoding="utf-8", errors="ignore")
|
||||
(output_dir / "mineru_stderr.log").write_text(result.stderr or "", encoding="utf-8", errors="ignore")
|
||||
except Exception as log_exc:
|
||||
print(f" ⚠️ MinerU 日志写入失败: {log_exc}")
|
||||
|
||||
if result.returncode == 0:
|
||||
print(f" ✓ MinerU pipeline 提取成功")
|
||||
return True
|
||||
else:
|
||||
print(f" ⚠️ MinerU pipeline 退出码: {result.returncode}")
|
||||
if result.stderr:
|
||||
print(f" stderr: {result.stderr[:1000]}")
|
||||
print(f" 完整 stderr: {output_dir / 'mineru_stderr.log'}")
|
||||
return False
|
||||
except subprocess.TimeoutExpired:
|
||||
print(f" ⚠️ MinerU pipeline 超时 (已设置900秒)")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f" ⚠️ MinerU pipeline 异常: {e}")
|
||||
return False
|
||||
|
||||
def _parse_mineru_output(self, mineru_dir: Path) -> Tuple[List[Section], str, Path]:
|
||||
"""解析 MinerU 输出 - V13.1 改进版"""
|
||||
sections = []
|
||||
full_text = ""
|
||||
|
||||
# V14.2: MinerU 输出在子目录中 (paper_name/hybrid_auto/ 或 paper_name/auto/)
|
||||
paper_name = mineru_dir.parent.name
|
||||
|
||||
# 尝试不同的输出目录结构
|
||||
possible_dirs = [
|
||||
mineru_dir / paper_name / 'hybrid_auto',
|
||||
mineru_dir / paper_name / 'auto',
|
||||
mineru_dir,
|
||||
]
|
||||
|
||||
actual_output_dir = None
|
||||
for d in possible_dirs:
|
||||
if d.exists() and (d / f'{paper_name}.md').exists():
|
||||
actual_output_dir = d
|
||||
break
|
||||
|
||||
if actual_output_dir:
|
||||
mineru_dir = actual_output_dir
|
||||
print(f" ✓ 找到 MinerU 输出: {mineru_dir}")
|
||||
else:
|
||||
print(f" ⚠️ 未找到 MinerU 输出文件,尝试默认路径: {mineru_dir}")
|
||||
|
||||
# 先读取完整文本(用于计算章节位置)
|
||||
text_file = mineru_dir / f'{paper_name}.md'
|
||||
if text_file.exists():
|
||||
try:
|
||||
with open(text_file, 'r') as f:
|
||||
full_text = f.read()
|
||||
# Apply text cleaning to fix PDF line breaks
|
||||
full_text = clean_text(full_text)
|
||||
print(f" ✓ 提取 {len(full_text)} 字符 (已清理PDF换行)")
|
||||
except Exception as e:
|
||||
print(f" ⚠️ 读取文本失败: {e}")
|
||||
else:
|
||||
print(f" ⚠️ 文本文件不存在: {text_file}")
|
||||
|
||||
# V13.1: 从 content_list.json 读取章节信息(替代 layout.json)
|
||||
content_list_file = mineru_dir / f'{paper_name}_content_list.json'
|
||||
if content_list_file.exists() and full_text:
|
||||
try:
|
||||
with open(content_list_file, 'r') as f:
|
||||
content_list = json.load(f)
|
||||
|
||||
# V13.1: 正则匹配章节名称(不区分大小写)
|
||||
section_patterns = [
|
||||
(r'(?i)^\s*(?:\d+(?:\.\d+)*\.?\s*)?(abstract|summary|overview)\s*$', 'abstract'),
|
||||
(r'(?i)^\s*(?:\d+(?:\.\d+)*\.?\s*)?(introduction|background|context)\s*$', 'introduction'),
|
||||
(r'(?i)^\s*(?:\d+(?:\.\d+)*\.?\s*)?(materials?\s*(&|and)?\s*methods?|methodology|experimental\s*procedures?)\s*$', 'methods'),
|
||||
(r'(?i)^\s*(?:\d+(?:\.\d+)*\.?\s*)?(results?|findings?|outcomes?)\s*$', 'results'),
|
||||
(r'(?i)^\s*(?:\d+(?:\.\d+)*\.?\s*)?(discussions?|interpretation|analysis)\s*$', 'discussion'),
|
||||
(r'(?i)^\s*(?:\d+(?:\.\d+)*\.?\s*)?(conclusions?|final\s*remarks)\s*$', 'conclusion'),
|
||||
]
|
||||
|
||||
# 排除的章节
|
||||
exclude_patterns = [
|
||||
r'(?i)references?',
|
||||
r'(?i)acknowledgments?',
|
||||
r'(?i)author\s*contributions?',
|
||||
r'(?i)supplementary\s*materials?',
|
||||
]
|
||||
|
||||
# 提取章节标题(短文本 + 匹配正则)
|
||||
for item in content_list:
|
||||
if item.get('type') == 'text':
|
||||
text = item.get('text', '').strip()
|
||||
# 过滤:短文本(<100字符)
|
||||
if len(text) < 100:
|
||||
# 检查是否匹配排除规则
|
||||
if any(re.match(p, text) for p in exclude_patterns):
|
||||
continue
|
||||
|
||||
# 匹配章节模式
|
||||
matched = False
|
||||
for pattern, normalized in section_patterns:
|
||||
if re.match(pattern, text):
|
||||
matched = True
|
||||
break
|
||||
|
||||
if matched:
|
||||
# 在文本中查找章节位置
|
||||
title_lower = text.lower()
|
||||
position = full_text.lower().find(title_lower)
|
||||
if position == -1:
|
||||
position = 0
|
||||
|
||||
# 使用标准化名称
|
||||
section = Section(
|
||||
title=normalized,
|
||||
level=1,
|
||||
parent=None,
|
||||
start_page=item.get('page_idx', 0) + 1,
|
||||
end_page=item.get('page_idx', 0) + 1,
|
||||
section_type='heading',
|
||||
position=position,
|
||||
end=len(full_text) # 临时设置为文本末尾
|
||||
)
|
||||
sections.append(section)
|
||||
|
||||
# 计算每个章节的结束位置(下一个章节的起始位置)
|
||||
sections.sort(key=lambda s: s.position)
|
||||
for i in range(len(sections) - 1):
|
||||
sections[i].end = sections[i + 1].position
|
||||
|
||||
sections = self._split_large_sections(sections, full_text)
|
||||
|
||||
print(f" ✓ 从 MinerU content_list 识别 {len(sections)} 个章节(含子章节)")
|
||||
except Exception as e:
|
||||
print(f" ⚠️ 解析 content_list.json 失败: {e}")
|
||||
|
||||
# 如果没有章节,使用 PyMuPDF 补充
|
||||
if not sections:
|
||||
print(" ⚠️ MinerU 未识别章节,使用 PyMuPDF 补充...")
|
||||
pdf_path = mineru_dir.parent.with_suffix('.pdf')
|
||||
if pdf_path.exists():
|
||||
sections = self._extract_sections_pymupdf(pdf_path)
|
||||
|
||||
return sections, full_text, mineru_dir
|
||||
|
||||
def _fallback_pymupdf(self, pdf_path: Path) -> Tuple[List[Section], str, Path]:
|
||||
"""PyMuPDF 回退"""
|
||||
print(" 使用 PyMuPDF 提取...")
|
||||
|
||||
doc = fitz.open(str(pdf_path))
|
||||
text = ""
|
||||
for page in doc:
|
||||
text += page.get_text() + "\n"
|
||||
doc.close()
|
||||
|
||||
sections = self._extract_sections_pymupdf(pdf_path)
|
||||
|
||||
print(f" ✓ PyMuPDF 提取 {len(text)} 字符, {len(sections)} 个章节")
|
||||
|
||||
return sections, text, Path()
|
||||
|
||||
def _split_large_sections(self, sections: List[Section], full_text: str) -> List[Section]:
|
||||
"""
|
||||
细粒度章节划分 - 匹配 PC 风格
|
||||
将大章节(如 results, discussion)分割为子章节(results1, results2)
|
||||
"""
|
||||
# 需要分割的章节类型(使用完整命名)
|
||||
splittable_types = ['results', 'discussion', 'introduction']
|
||||
# 子章节大小:约 5000 字符
|
||||
subsection_size = 5000
|
||||
|
||||
new_sections = []
|
||||
subsection_counter = {}
|
||||
|
||||
for sec in sections:
|
||||
sec_length = sec.end - sec.position
|
||||
|
||||
# 检查是否需要分割
|
||||
if sec.title in splittable_types and sec_length > subsection_size * 2:
|
||||
# 需要分割为大章节 + 子章节
|
||||
num_subsections = max(2, sec_length // subsection_size)
|
||||
subsection_length = sec_length // num_subsections
|
||||
|
||||
# 添加大章节(第一个子章节)
|
||||
base_title = sec.title
|
||||
subsection_counter[base_title] = 1
|
||||
|
||||
for i in range(num_subsections):
|
||||
sub_position = sec.position + i * subsection_length
|
||||
sub_end = min(sec.position + (i + 1) * subsection_length, sec.end)
|
||||
|
||||
if i == 0:
|
||||
# 第一个使用基础名称(无编号)
|
||||
sub_title = base_title
|
||||
else:
|
||||
# 后续使用数字后缀,从1开始(results1, results2...)
|
||||
sub_title = f"{base_title}{i}"
|
||||
|
||||
new_sec = Section(
|
||||
title=sub_title,
|
||||
level=sec.level,
|
||||
parent=sec.parent,
|
||||
start_page=sec.start_page,
|
||||
end_page=sec.end_page,
|
||||
section_type=sec.section_type,
|
||||
position=sub_position,
|
||||
end=sub_end
|
||||
)
|
||||
new_sections.append(new_sec)
|
||||
else:
|
||||
# 不需要分割,直接添加
|
||||
new_sections.append(sec)
|
||||
|
||||
return new_sections
|
||||
|
||||
def extract_section_text(self, full_text: str, section: Section) -> str:
|
||||
"""根据章节位置从全文提取章节文本"""
|
||||
if section.position is not None and section.end is not None:
|
||||
return full_text[section.position:section.end]
|
||||
return full_text
|
||||
|
||||
def _extract_sections_pymupdf(self, pdf_path: Path) -> List[Section]:
|
||||
"""使用 PyMuPDF 提取章节"""
|
||||
sections = []
|
||||
|
||||
# 简单的章节识别规则
|
||||
heading_patterns = [
|
||||
r'^\d+\.\s+([A-Z][^\n]+)', # 1. Introduction
|
||||
r'^([A-Z][A-Z\s]+)$', # INTRODUCTION
|
||||
r'^(Abstract|Introduction|Methods|Results|Discussion|Conclusion)s?$',
|
||||
]
|
||||
|
||||
doc = fitz.open(str(pdf_path))
|
||||
for page_num, page in enumerate(doc, 1):
|
||||
text = page.get_text()
|
||||
for pattern in heading_patterns:
|
||||
for match in re.finditer(pattern, text, re.MULTILINE):
|
||||
section = Section(
|
||||
title=match.group(1).strip(),
|
||||
level=1,
|
||||
parent=None,
|
||||
start_page=page_num,
|
||||
end_page=page_num,
|
||||
section_type='heading'
|
||||
)
|
||||
sections.append(section)
|
||||
doc.close()
|
||||
|
||||
return sections
|
||||
|
||||
|
||||
# 便捷函数
|
||||
def parse_document(pdf_path: Path, output_base_dir: Optional[Path] = None) -> Tuple[List[Section], str, Path]:
|
||||
"""解析文档"""
|
||||
parser = MinerUDocumentParser()
|
||||
return parser.parse(pdf_path, output_base_dir)
|
||||
File diff suppressed because it is too large
Load Diff
+496
@@ -0,0 +1,496 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Phase 3: 实体识别
|
||||
整合v12的实体过滤 + fusion的实体识别
|
||||
"""
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Optional
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
# 导入PC数据加载器
|
||||
try:
|
||||
from .pc_data_loader import get_pc_data_loader
|
||||
PC_DATA_AVAILABLE = True
|
||||
except ImportError:
|
||||
PC_DATA_AVAILABLE = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class Entity:
|
||||
"""实体定义"""
|
||||
id: str
|
||||
name: str
|
||||
type: str
|
||||
synonyms: List[str] = field(default_factory=list)
|
||||
properties: Dict = field(default_factory=dict)
|
||||
confidence: float = 1.0
|
||||
source: str = ""
|
||||
# PC风格字段
|
||||
name_resolved: str = "" # 标准化后的名称
|
||||
aliases: List[str] = field(default_factory=list) # 别名列表
|
||||
|
||||
|
||||
class EntityRecognizer:
|
||||
"""实体识别器 - 整合v12+fusion+PC+V14.2能力"""
|
||||
|
||||
# V14.2优化: 类级别的单例,避免重复加载
|
||||
_pc_loader = None
|
||||
_entity_linker = None
|
||||
|
||||
def __init__(self, use_pc_data: bool = True, use_entity_linker: bool = True):
|
||||
self.entities: List[Entity] = []
|
||||
self.entity_map: Dict[str, Entity] = {}
|
||||
|
||||
# PC数据加载器
|
||||
self.pc_loader = None
|
||||
self.entity_linker = None
|
||||
|
||||
if use_pc_data and PC_DATA_AVAILABLE:
|
||||
try:
|
||||
# V14.2优化: 使用类级别的单例
|
||||
if EntityRecognizer._pc_loader is None:
|
||||
EntityRecognizer._pc_loader = get_pc_data_loader()
|
||||
self.pc_loader = EntityRecognizer._pc_loader
|
||||
print(" ✓ PC数据加载器已启用")
|
||||
|
||||
# V14.2: 实体链接器
|
||||
if use_entity_linker:
|
||||
try:
|
||||
from .entity_linker import get_entity_linker
|
||||
# V14.2优化: 使用类级别的单例
|
||||
if EntityRecognizer._entity_linker is None:
|
||||
EntityRecognizer._entity_linker = get_entity_linker(self.pc_loader)
|
||||
self.entity_linker = EntityRecognizer._entity_linker
|
||||
print(" ✓ V14.2实体链接器已启用")
|
||||
except Exception as e:
|
||||
print(f" ⚠️ 实体链接器加载失败: {e}")
|
||||
|
||||
except Exception as e:
|
||||
print(f" ⚠️ PC数据加载失败: {e}")
|
||||
|
||||
def recognize(self, text: str) -> List[Entity]:
|
||||
"""
|
||||
识别实体
|
||||
|
||||
识别8种实体类型:
|
||||
- gene (基因)
|
||||
- tissue (组织)
|
||||
- metabolite (代谢物)
|
||||
- process (过程)
|
||||
- compartment (亚细胞定位)
|
||||
- species (物种)
|
||||
- phylogeny_group (系统发育分组)
|
||||
- chromosome (染色体位置,作为属性)
|
||||
"""
|
||||
print(" Step 3.1: 识别基因实体...")
|
||||
self._extract_genes(text)
|
||||
|
||||
print(" Step 3.2: 识别组织...")
|
||||
self._extract_tissues(text)
|
||||
|
||||
print(" Step 3.3: 识别代谢物...")
|
||||
self._extract_metabolites(text)
|
||||
|
||||
print(" Step 3.4: 识别过程...")
|
||||
self._extract_processes(text)
|
||||
|
||||
print(" Step 3.5: 识别表型...")
|
||||
self._extract_phenotypes(text)
|
||||
|
||||
print(" Step 3.6: 识别亚细胞定位...")
|
||||
self._extract_compartments(text)
|
||||
|
||||
print(" Step 3.7: 识别物种...")
|
||||
self._extract_species(text)
|
||||
|
||||
print(" Step 3.8: 识别系统发育分组...")
|
||||
self._extract_phylogeny_groups(text)
|
||||
|
||||
# 去重
|
||||
self._deduplicate_entities()
|
||||
|
||||
return self.entities
|
||||
|
||||
def _extract_genes(self, text: str):
|
||||
"""识别基因(通用模式 - 支持任何物种)"""
|
||||
|
||||
# 模式1: 通用基因名模式(最优先)
|
||||
# 格式: 前缀(1-2字母) + 家族名(1-4大写字母) + 数字(1-4位) + 可选亚型(a-z)
|
||||
# 示例: AtSPL1, OsSPL16a, CiSPL12b, BpSPL9, MdMYB1
|
||||
pattern1 = r'\b([A-Z][a-z]?[A-Z]{1,4}\d{1,4}[a-z]?)\b'
|
||||
for match in re.finditer(pattern1, text):
|
||||
gene_id = match.group(1)
|
||||
# 验证格式
|
||||
if re.match(r'^[A-Z][a-z]?[A-Z]{1,4}\d{1,4}[a-z]?$', gene_id):
|
||||
self._add_entity(gene_id, "gene", match.group(0))
|
||||
# 如果有亚型字母,同时添加基础基因名
|
||||
if gene_id[-1].islower():
|
||||
base_gene = gene_id[:-1]
|
||||
self._add_entity(base_gene, "gene", match.group(0))
|
||||
|
||||
# 模式2: 带连字符的基因名 (如 CiSPL12-a)
|
||||
pattern2 = r'\b([A-Z][a-z]?[A-Z]{1,4}\d{1,4})-([a-z])\b'
|
||||
for match in re.finditer(pattern2, text):
|
||||
gene_id = match.group(1) + match.group(2)
|
||||
self._add_entity(gene_id, "gene", match.group(0))
|
||||
|
||||
# 模式3: 复数基因形式 (如 CiSPLs, SPL genes)
|
||||
pattern3 = r'\b([A-Z][a-z]?[A-Z]{1,4}\d{1,4}[a-z]?)s\b'
|
||||
for match in re.finditer(pattern3, text):
|
||||
gene_id = match.group(1)
|
||||
if re.match(r'^[A-Z][a-z]?[A-Z]{1,4}\d{1,4}[a-z]?$', gene_id):
|
||||
self._add_entity(gene_id, "gene", match.group(0))
|
||||
|
||||
# 模式4: 苦荞特殊格式 FtPinG00001.1(保留兼容性)
|
||||
pattern4 = r'\b(FtPinG\d+\.\d+)\b'
|
||||
for match in re.finditer(pattern4, text):
|
||||
self._add_entity(match.group(1), "gene", match.group(0))
|
||||
|
||||
# 模式5-8: PC 复合实体识别
|
||||
self._extract_composite_entities(text)
|
||||
|
||||
def _extract_composite_entities(self, text: str):
|
||||
"""提取 PC 风格的复合实体"""
|
||||
|
||||
# 格式1: "32 SPL gene(s) (CiSPLs)" - 数量+基因家族+别名
|
||||
pattern1 = r'(\d+)\s+([A-Z][a-z]{1,2}[A-Z]{1,4})\s+genes?\s*\(\s*([A-Z][a-z]{1,2}[A-Z]{1,4})s?\s*\)'
|
||||
for match in re.finditer(pattern1, text, re.IGNORECASE):
|
||||
count = match.group(1)
|
||||
family = match.group(2)
|
||||
alias_raw = match.group(3)
|
||||
# 修正 alias:去掉末尾的 s(如果有)
|
||||
alias = alias_raw[:-1] if alias_raw.endswith('s') else alias_raw
|
||||
entity_name = f"{count} {family} gene(s) ({alias}s)"
|
||||
|
||||
# 查找该群体的成员基因
|
||||
member_genes = self._find_member_genes(text, alias)
|
||||
|
||||
entity = self._add_entity(entity_name, "gene_group", match.group(0))
|
||||
if entity:
|
||||
entity.properties['count'] = count
|
||||
entity.properties['family'] = family
|
||||
entity.properties['alias'] = alias
|
||||
entity.properties['member_genes'] = member_genes
|
||||
entity.synonyms = [f"{alias}s", f"{family} gene(s)", f"{count} {alias}s"]
|
||||
|
||||
# 格式2: "32 CiSPLs" - 数量+别名
|
||||
pattern2 = r'(\d+)\s+([A-Z][a-z]{1,2}[A-Z]{1,4})s\b'
|
||||
for match in re.finditer(pattern2, text):
|
||||
count = match.group(1)
|
||||
family = match.group(2)
|
||||
entity_name = f"{count} {family}s"
|
||||
|
||||
member_genes = self._find_member_genes(text, family)
|
||||
|
||||
entity = self._add_entity(entity_name, "gene_group", match.group(0))
|
||||
if entity:
|
||||
entity.properties['count'] = count
|
||||
entity.properties['family'] = family
|
||||
entity.properties['member_genes'] = member_genes
|
||||
|
||||
# 格式3: "SPL gene(s) family" - 基因家族+family
|
||||
pattern3 = r'([A-Z][a-z]{1,2}[A-Z]{1,4})\s+genes?\s+family'
|
||||
for match in re.finditer(pattern3, text, re.IGNORECASE):
|
||||
family = match.group(1)
|
||||
entity_name = f"{family} gene(s) family"
|
||||
|
||||
entity = self._add_entity(entity_name, "gene_family", match.group(0))
|
||||
if entity:
|
||||
entity.properties['family'] = family
|
||||
entity.synonyms = [f"{family} genes", f"{family}s"]
|
||||
|
||||
# 格式4: "SPL gene(s)" - 基因家族+gene(s)
|
||||
pattern4 = r'([A-Z][a-z]{1,2}[A-Z]{1,4})\s+genes?'
|
||||
for match in re.finditer(pattern4, text, re.IGNORECASE):
|
||||
family = match.group(1)
|
||||
# 避免与格式3重复
|
||||
if not re.search(rf'{family}\s+gene\(s\)\s+family', text[match.start():match.end()+20], re.IGNORECASE):
|
||||
entity_name = f"{family} gene(s)"
|
||||
|
||||
entity = self._add_entity(entity_name, "gene_group", match.group(0))
|
||||
if entity:
|
||||
entity.properties['family'] = family
|
||||
entity.synonyms = [f"{family}s"]
|
||||
|
||||
# 格式5: "nine segmentally duplicated gene(s) pairs"
|
||||
pattern5 = r'(\d+|nine|eight|seven|six|five|four|three|two|one)\s+segmentally\s+duplicated\s+genes?\s*pairs?'
|
||||
for match in re.finditer(pattern5, text, re.IGNORECASE):
|
||||
count = match.group(1)
|
||||
entity_name = f"{count} segmentally duplicated gene pairs"
|
||||
|
||||
entity = self._add_entity(entity_name, "gene_group", match.group(0))
|
||||
if entity:
|
||||
entity.properties['count'] = count
|
||||
entity.properties['duplication_type'] = 'segmental'
|
||||
|
||||
# 格式6: "16 collinear gene(s) pairs"
|
||||
pattern6 = r'(\d+)\s+collinear\s+genes?\s*pairs?'
|
||||
for match in re.finditer(pattern6, text, re.IGNORECASE):
|
||||
count = match.group(1)
|
||||
entity_name = f"{count} collinear gene pairs"
|
||||
|
||||
entity = self._add_entity(entity_name, "gene_group", match.group(0))
|
||||
if entity:
|
||||
entity.properties['count'] = count
|
||||
entity.properties['relationship'] = 'collinear'
|
||||
|
||||
# 格式7: "SPL transcription factors"
|
||||
pattern7 = r'([A-Z][a-z]{1,2}[A-Z]{1,4})\s+transcription\s+factors?'
|
||||
for match in re.finditer(pattern7, text, re.IGNORECASE):
|
||||
family = match.group(1)
|
||||
entity_name = f"{family} transcription factors"
|
||||
|
||||
entity = self._add_entity(entity_name, "gene_group", match.group(0))
|
||||
if entity:
|
||||
entity.properties['family'] = family
|
||||
entity.properties['type'] = 'transcription_factor'
|
||||
|
||||
def _find_member_genes(self, text: str, family_prefix: str) -> list:
|
||||
"""查找复合实体中的成员基因 - V14.2优化: 限制搜索范围"""
|
||||
# V14.2优化: 只在文本的前10000个字符中查找成员基因
|
||||
# 这样可以避免处理超大文本时的性能问题
|
||||
search_text = text[:10000] if len(text) > 10000 else text
|
||||
pattern = rf'\b{family_prefix}\d+[a-z]?\b'
|
||||
matches = re.findall(pattern, search_text)
|
||||
return sorted(set(matches))
|
||||
|
||||
def _extract_tissues(self, text: str):
|
||||
"""识别组织(V16.aqua 水产化词表)"""
|
||||
tissues = [
|
||||
'gill', 'gills', 'liver', 'hepatopancreas', 'muscle', 'intestine', 'gut',
|
||||
'skin', 'spleen', 'kidney', 'head kidney', 'brain', 'eye', 'heart',
|
||||
'gonad', 'ovary', 'testis', 'fin', 'fins', 'scales', 'mucus',
|
||||
'blood', 'serum', 'plasma', 'hemolymph', 'haemolymph',
|
||||
'mantle', 'shell', 'adductor muscle', 'foot', 'siphon',
|
||||
'thallus', 'frond', 'blade', 'stipe'
|
||||
]
|
||||
|
||||
for tissue in tissues:
|
||||
pattern = rf'\b{tissue}\b'
|
||||
for match in re.finditer(pattern, text, re.IGNORECASE):
|
||||
self._add_entity(tissue.lower(), "tissue", match.group(0))
|
||||
|
||||
def _extract_metabolites(self, text: str):
|
||||
"""识别代谢物(V16.aqua 水产化词表)"""
|
||||
metabolites = [
|
||||
'glucose', 'glycogen', 'starch', 'fatty acid', 'lipid', 'protein',
|
||||
'amino acid', 'lysine', 'methionine', 'arginine',
|
||||
'cholesterol', 'triglyceride', 'astaxanthin', 'carotenoid',
|
||||
'cortisol', 'testosterone', 'estradiol', 'thyroid hormone',
|
||||
'insulin-like growth factor', 'melatonin'
|
||||
]
|
||||
|
||||
for met in metabolites:
|
||||
pattern = rf'\b{met}\b'
|
||||
for match in re.finditer(pattern, text, re.IGNORECASE):
|
||||
self._add_entity(met.lower(), "metabolite", match.group(0))
|
||||
|
||||
def _extract_processes(self, text: str):
|
||||
"""识别过程(V16.aqua 水产化词表)"""
|
||||
# 基础过程词汇
|
||||
processes = [
|
||||
'transport', 'accumulation', 'biosynthesis', 'metabolism',
|
||||
'signaling', 'regulation', 'development', 'growth',
|
||||
'stress response', 'immune response', 'immune defense',
|
||||
'osmoregulation', 'ion regulation', 'digestion', 'absorption',
|
||||
'molting', 'moulting', 'ecdysis', 'spawning', 'hatching', 'reproduction',
|
||||
'maturation', 'smoltification', 'metamorphosis',
|
||||
'respiration', 'nitrogen excretion', 'ammonia excretion',
|
||||
'differentiation', 'proliferation', 'activation',
|
||||
'repression', 'induction', 'inhibition', 'apoptosis',
|
||||
'phagocytosis', 'vaccination', 'immunization'
|
||||
]
|
||||
|
||||
for proc in processes:
|
||||
pattern = rf'\b{proc}\b'
|
||||
for match in re.finditer(pattern, text, re.IGNORECASE):
|
||||
self._add_entity(proc.lower(), "process", match.group(0))
|
||||
|
||||
# V16.aqua: 水产高频过程短语
|
||||
process_phrases = [
|
||||
r'(?:larval|juvenile|embryonic)\s+development',
|
||||
r'(?:cell|tissue)\s+proliferation',
|
||||
r'(?:specific|relative)\s+growth',
|
||||
r'(?:gonadal|ovarian|testicular)\s+development',
|
||||
r'(?:thermal|salinity|hypoxia|ammonia|nitrite)\s+stress',
|
||||
r'(?:immune|innate immune|antiviral)\s+response',
|
||||
r'(?:feed|protein)\s+(?:utilization|efficiency)',
|
||||
r'(?:disease|viral|bacterial)\s+resistance',
|
||||
r'(?:water|environmental)\s+quality',
|
||||
]
|
||||
|
||||
for pattern in process_phrases:
|
||||
for match in re.finditer(pattern, text, re.IGNORECASE):
|
||||
self._add_entity(match.group(0).lower(), "process", match.group(0))
|
||||
|
||||
def _extract_phenotypes(self, text: str):
|
||||
"""识别表型(V16.aqua 水产化词表)"""
|
||||
phenotypes = [
|
||||
'growth performance', 'specific growth rate', 'weight gain',
|
||||
'feed conversion ratio', 'feed conversion efficiency',
|
||||
'survival rate', 'mortality', 'cumulative mortality',
|
||||
'disease resistance', 'immune response', 'immunity',
|
||||
'final body weight', 'body weight', 'condition factor',
|
||||
'hepatosomatic index', 'viscerosomatic index',
|
||||
'gonadosomatic index', 'fecundity', 'hatching rate',
|
||||
'fertilization rate', 'egg quality', 'sperm motility',
|
||||
'flesh quality', 'fillet yield', 'protein content',
|
||||
'lipid content', 'moisture content', 'ash content',
|
||||
'carotenoid content', 'pigmentation'
|
||||
]
|
||||
|
||||
for pheno in phenotypes:
|
||||
pattern = rf'\b{pheno}\b'
|
||||
for match in re.finditer(pattern, text, re.IGNORECASE):
|
||||
self._add_entity(match.group(0).lower(), "phenotype", match.group(0))
|
||||
|
||||
# 表型短语模式
|
||||
pheno_patterns = [
|
||||
r'(?:enhanced|improved|increased|reduced|decreased)\s+(?:growth|survival|resistance|tolerance|immunity|fcr|feed conversion)\b[\w\s]*',
|
||||
r'(?:higher|lower|better|worse)\s+(?:growth|survival|mortality|fcr|feed conversion)\b[\w\s]*',
|
||||
r'(?:early|late|delayed)\s+(?:maturation|spawning|hatching|molting)',
|
||||
]
|
||||
|
||||
for pattern in pheno_patterns:
|
||||
for match in re.finditer(pattern, text, re.IGNORECASE):
|
||||
# 截断过长的匹配
|
||||
pheno_text = match.group(0)
|
||||
if len(pheno_text) > 60:
|
||||
pheno_text = pheno_text[:60] + "..."
|
||||
self._add_entity(pheno_text.lower(), "phenotype", match.group(0))
|
||||
|
||||
def _extract_compartments(self, text: str):
|
||||
"""识别亚细胞定位(fusion方法,通用细胞器水产适用)"""
|
||||
compartments = [
|
||||
'plasma membrane', 'cell membrane', 'cytoplasm', 'nucleus',
|
||||
'vacuole', 'mitochondria', 'endoplasmic reticulum',
|
||||
'golgi apparatus', 'peroxisome', 'lysosome'
|
||||
]
|
||||
|
||||
for comp in compartments:
|
||||
pattern = rf'\b{comp}\b'
|
||||
for match in re.finditer(pattern, text, re.IGNORECASE):
|
||||
self._add_entity(comp.lower(), "compartment", match.group(0))
|
||||
|
||||
def _extract_species(self, text: str):
|
||||
"""识别物种(V16.aqua 水产化词表:主要水产养殖物种)"""
|
||||
species = [
|
||||
# 虾蟹
|
||||
'Litopenaeus vannamei', 'Penaeus monodon', 'Macrobrachium rosenbergii',
|
||||
'Eriocheir sinensis', 'Scylla paramamosain', 'Procambarus clarkii',
|
||||
'vannamei', 'shrimp', 'prawn', 'crab', 'crayfish',
|
||||
# 鱼类
|
||||
'Oreochromis niloticus', 'tilapia', 'Oncorhynchus mykiss',
|
||||
'rainbow trout', 'Salmo salar', 'Atlantic salmon', 'salmon',
|
||||
'Ctenopharyngodon idella', 'grass carp', 'Cyprinus carpio',
|
||||
'common carp', 'carp', 'Lateolabrax maculatus', 'sea bass',
|
||||
'Larimichthys crocea', 'yellow croaker', 'Epinephelus',
|
||||
'grouper', 'Danio rerio', 'zebrafish', 'Anguilla japonica',
|
||||
'eel', 'Ictalurus punctatus', 'channel catfish', 'catfish',
|
||||
'Dicentrarchus labrax', 'Sparus aurata', 'gilthead seabream',
|
||||
'Takifugu rubripes', 'pufferfish', 'Gadus morhua', 'cod',
|
||||
'Misgurnus anguillicaudatus', 'loach',
|
||||
# 贝类
|
||||
'Crassostrea gigas', 'oyster', 'Chlamys farreri', 'scallop',
|
||||
'Mytilus', 'mussel', 'Ruditapes philippinarum', 'clam',
|
||||
'Haliotis', 'abalone', 'Pinctada', 'pearl oyster',
|
||||
# 藻类
|
||||
'Saccharina japonica', 'kelp', 'Undaria pinnatifida',
|
||||
'Laminaria', 'Gracilaria', 'Pyropia', 'Porphyra', 'nori',
|
||||
'Ulva', 'Chlorella', 'Spirulina', 'Haematococcus pluvialis', 'Artemia',
|
||||
'Nannochloropsis', 'seaweed', 'macroalgae', 'microalgae',
|
||||
# 其他
|
||||
'Apostichopus japonicus', 'sea cucumber', 'Strongylocentrotus',
|
||||
'sea urchin', 'Rana catesbeiana', 'bullfrog', 'Cherax quadricarinatus',
|
||||
'Penaeus japonicus', 'Marsupenaeus japonicus'
|
||||
]
|
||||
|
||||
for sp in species:
|
||||
pattern = rf'\b{re.escape(sp)}\b'
|
||||
for match in re.finditer(pattern, text, re.IGNORECASE):
|
||||
self._add_entity(sp.lower(), "species", match.group(0))
|
||||
|
||||
def _extract_phylogeny_groups(self, text: str):
|
||||
"""识别系统发育分组(fusion方法)"""
|
||||
# 模式: Group I, Group II, Group III, Group IV
|
||||
pattern = r'(Group\s+[IV\d]+)'
|
||||
for match in re.finditer(pattern, text, re.IGNORECASE):
|
||||
self._add_entity(match.group(1), "phylogeny_group", match.group(0))
|
||||
|
||||
def _add_entity(self, entity_id: str, entity_type: str, evidence: str) -> Entity:
|
||||
"""添加实体(V14.2: 支持实体链接器)"""
|
||||
# V14.2优化: 先检查原始名称是否已存在,避免重复链接
|
||||
original_key = f"{entity_id.lower()}:{entity_type}"
|
||||
if original_key in self.entity_map:
|
||||
return self.entity_map[original_key]
|
||||
|
||||
# V14.2: 使用实体链接器进行标准化
|
||||
name_resolved = entity_id
|
||||
aliases = []
|
||||
link_method = "none"
|
||||
link_score = 0.0
|
||||
|
||||
# V14.2优化: 只对基因类型使用实体链接器,其他类型跳过
|
||||
if entity_type == 'gene' and self.entity_linker:
|
||||
# 使用V14.2实体链接器(模糊匹配)
|
||||
link_result = self.entity_linker.link(entity_id, entity_type)
|
||||
name_resolved = link_result['resolved']
|
||||
aliases = link_result['aliases']
|
||||
link_method = link_result['method']
|
||||
link_score = link_result['score']
|
||||
elif self.pc_loader:
|
||||
# 回退到V14.1方式
|
||||
if entity_type == 'gene':
|
||||
name_resolved = self.pc_loader.normalize_entity(entity_id)
|
||||
aliases = self.pc_loader.expand_gene_aliases(name_resolved)
|
||||
|
||||
# 使用标准化后的名称作为key
|
||||
entity_key = f"{name_resolved.lower()}:{entity_type}"
|
||||
|
||||
if entity_key not in self.entity_map:
|
||||
entity = Entity(
|
||||
id=name_resolved,
|
||||
name=entity_id,
|
||||
type=entity_type,
|
||||
source=evidence,
|
||||
name_resolved=name_resolved,
|
||||
aliases=aliases
|
||||
)
|
||||
# V14.2: 添加链接元数据
|
||||
entity.properties['link_method'] = link_method
|
||||
entity.properties['link_score'] = link_score
|
||||
|
||||
self.entities.append(entity)
|
||||
self.entity_map[original_key] = entity # 同时用原始名称索引
|
||||
self.entity_map[entity_key] = entity # 用标准化名称索引
|
||||
return entity
|
||||
return self.entity_map[entity_key]
|
||||
|
||||
def _deduplicate_entities(self):
|
||||
"""去重(简单实现)"""
|
||||
seen = set()
|
||||
unique = []
|
||||
for e in self.entities:
|
||||
key = (e.name.lower(), e.type)
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
unique.append(e)
|
||||
self.entities = unique
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 测试
|
||||
recognizer = EntityRecognizer()
|
||||
test_text = """
|
||||
SlSut1 was highly expressed in leaves and fruits.
|
||||
It transports sucrose and localizes to the plasma membrane.
|
||||
Grouped with Group I in Arabidopsis.
|
||||
"""
|
||||
|
||||
entities = recognizer.recognize(test_text)
|
||||
print(f"\n测试完成!")
|
||||
print(f"实体数: {len(entities)}")
|
||||
for e in entities:
|
||||
print(f" - {e.name} ({e.type})")
|
||||
@@ -0,0 +1,61 @@
|
||||
from __future__ import annotations
|
||||
import csv
|
||||
from pathlib import Path
|
||||
from typing import Iterable,List
|
||||
from .entity_validator_v16_1_1 import (
|
||||
hard_reject_reasons,
|
||||
infer_layer,
|
||||
is_administrative_relation,
|
||||
is_weak_relation,
|
||||
qa_weight,
|
||||
)
|
||||
class SoftGateV16_1_1:
|
||||
def __init__(self, output_dir: Path|None=None): self.output_dir=Path(output_dir) if output_dir else None; self.audit_rows=[]
|
||||
def apply(self, relations: Iterable, source_label: str="") -> List:
|
||||
out=[]
|
||||
for r in relations or []:
|
||||
reasons=hard_reject_reasons(r.source,r.source_type,r.relation,r.target,r.target_type,r.evidence); action="keep"
|
||||
if is_administrative_relation(r.relation, r.evidence):
|
||||
reasons.append("administrative_metadata")
|
||||
if reasons:
|
||||
r.output_layer="rejected"; r.candidate_status="rejected"; r.triple_quality="rejected"; r.triple_quality_reason=";".join(reasons); action="hard_reject"
|
||||
else:
|
||||
endpoint_problem=bool(getattr(r,"validation_reason",""))
|
||||
if is_weak_relation(r.relation):
|
||||
endpoint_problem = True
|
||||
weak_reason = "weak_relation_predicate"
|
||||
existing = getattr(r, "validation_reason", "")
|
||||
r.validation_reason = ";".join(
|
||||
value for value in [existing, weak_reason] if value
|
||||
)
|
||||
if getattr(r, "output_layer", "") not in {
|
||||
"method_relations",
|
||||
"dataset_metadata",
|
||||
"caption_relations",
|
||||
}:
|
||||
r.output_layer = "candidate_review"
|
||||
r.candidate_status = "candidate"
|
||||
r.triple_quality = "candidate"
|
||||
if not getattr(r,"output_layer",""):
|
||||
r.output_layer=infer_layer(r.source_type,getattr(r,"relation_category",""),getattr(r,"layer_hint",""),getattr(r,"source_label",source_label),getattr(r,"claim_scope",""),float(r.confidence or 0),endpoint_problem)
|
||||
if not r.candidate_status: r.candidate_status="candidate" if r.output_layer in {"candidate_review","caption_relations"} else "final"
|
||||
if not r.triple_quality: r.triple_quality="candidate" if r.candidate_status=="candidate" else "supported"
|
||||
r.qa_weight=qa_weight(r.output_layer)
|
||||
audit_reason = ";".join(
|
||||
dict.fromkeys(
|
||||
value
|
||||
for value in [
|
||||
getattr(r, "triple_quality_reason", ""),
|
||||
getattr(r, "validation_reason", ""),
|
||||
]
|
||||
if value
|
||||
)
|
||||
)
|
||||
self.audit_rows.append({"source":r.source,"relation":r.relation,"target":r.target,"output_layer":r.output_layer,"candidate_status":r.candidate_status,"action":action,"reason":audit_reason,"evidence":r.evidence})
|
||||
out.append(r)
|
||||
self.save_audit(); return out
|
||||
def save_audit(self):
|
||||
if not self.output_dir: return
|
||||
fields=["source","relation","target","output_layer","candidate_status","action","reason","evidence"]
|
||||
with open(self.output_dir/"phase45_soft_gate_audit.csv","w",newline="",encoding="utf-8") as f:
|
||||
w=csv.DictWriter(f,fieldnames=fields); w.writeheader(); w.writerows(self.audit_rows)
|
||||
+1022
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,26 @@
|
||||
from __future__ import annotations
|
||||
import csv
|
||||
from pathlib import Path
|
||||
from typing import Iterable,List
|
||||
from .entity_validator_v16_1_1 import infer_layer,qa_weight
|
||||
class LayerAssignerV16_1_1:
|
||||
def __init__(self, output_dir: Path|None=None): self.output_dir=Path(output_dir) if output_dir else None; self.audit_rows=[]
|
||||
def apply(self, relations: Iterable, source_label: str="") -> List:
|
||||
out=[]
|
||||
for r in relations or []:
|
||||
old=getattr(r,"output_layer","")
|
||||
if old in {"rejected","hard_reject"}: r.qa_weight=0.0
|
||||
else:
|
||||
endpoint_problem=bool(getattr(r,"validation_reason","")); layer=infer_layer(r.source_type,getattr(r,"relation_category",""),getattr(r,"layer_hint",""),getattr(r,"source_label",source_label),getattr(r,"claim_scope",""),float(r.confidence or 0),endpoint_problem)
|
||||
if old in {"caption_relations","candidate_review","method_relations","dataset_metadata","background_relations"}: layer=old
|
||||
r.output_layer=layer; r.qa_weight=qa_weight(layer)
|
||||
if not r.candidate_status: r.candidate_status="candidate" if layer in {"candidate_review","caption_relations"} else "final"
|
||||
if not r.triple_quality: r.triple_quality="candidate" if r.candidate_status=="candidate" else "supported"
|
||||
self.audit_rows.append({"source":r.source,"relation":r.relation,"target":r.target,"old_output_layer":old,"new_output_layer":r.output_layer,"qa_weight":getattr(r,"qa_weight",0),"relation_category":getattr(r,"relation_category",""),"reason":getattr(r,"validation_reason","")})
|
||||
out.append(r)
|
||||
self.save_audit(); return out
|
||||
def save_audit(self):
|
||||
if not self.output_dir: return
|
||||
fields=["source","relation","target","old_output_layer","new_output_layer","qa_weight","relation_category","reason"]
|
||||
with open(self.output_dir/"phase46_layer_assignment_audit.csv","w",newline="",encoding="utf-8") as f:
|
||||
w=csv.DictWriter(f,fieldnames=fields); w.writeheader(); w.writerows(self.audit_rows)
|
||||
+2574
File diff suppressed because it is too large
Load Diff
+2010
File diff suppressed because it is too large
Load Diff
Executable
+2
@@ -0,0 +1,2 @@
|
||||
from .phase4_llm_srt_extraction_v16_1_1 import PlantScienceLLMExtractorV16_1_1
|
||||
LLMRelationExtractorV15_2_1 = PlantScienceLLMExtractorV16_1_1
|
||||
+153
@@ -0,0 +1,153 @@
|
||||
from __future__ import annotations
|
||||
import csv,json,os,re,time
|
||||
from pathlib import Path
|
||||
from typing import Dict,List,Optional
|
||||
from .phase4_llm_relation_extraction import LLMRelationExtractor
|
||||
from .phase4_relation_extraction import Relation
|
||||
from .entity_validator_v16_1_1 import ALLOWED_ENTITY_TYPES,ALLOWED_RELATION_CATEGORIES,clean_text,entity_issues,hard_reject_reasons,normalize_relation,infer_layer,qa_weight
|
||||
|
||||
V16_SYSTEM_PROMPT = """You extract explicit aquaculture-science source-relation-target facts from scientific paper text.
|
||||
Domain: aquaculture, aquatic biology, fisheries, fish/mollusk/crustacean/algae physiology and nutrition.
|
||||
Input is scientific text mixed Chinese and English; extract entity names and evidence in the source language.
|
||||
Use only the provided numbered sentences.
|
||||
Favor high recall, but every relation must be directly supported by the evidence sentence ids.
|
||||
Do not use outside knowledge.
|
||||
Do not invent aliases, accessions, species names, or gene functions.
|
||||
Return JSON only."""
|
||||
|
||||
class PlantScienceLLMExtractorV16_1_1(LLMRelationExtractor):
|
||||
PAGE_CHAR_LIMIT=int(os.environ.get("V16_LLM1_PAGE_CHARS") or 2200)
|
||||
def __init__(self,*args,**kwargs):
|
||||
super().__init__(*args,**kwargs)
|
||||
self.v16_entity_rows=[]; self.v16_relation_rows=[]; self.v16_rejected_rows=[]
|
||||
self._cost_log_path=None
|
||||
def set_cost_log(self,path:Optional[Path])->None:
|
||||
"""V16.aqua cost instrumentation: JSONL append {paper_id, model, prompt_tokens, completion_tokens, timestamp}."""
|
||||
self._cost_log_path=Path(path) if path else None
|
||||
def _record_cost(self,paper_id,resp)->None:
|
||||
if not self._cost_log_path: return
|
||||
try:
|
||||
usage=getattr(resp,"usage",None)
|
||||
row={"paper_id":str(paper_id or ""),"model":str(getattr(self,"model","") or ""),
|
||||
"prompt_tokens":int(getattr(usage,"prompt_tokens",0) or 0),
|
||||
"completion_tokens":int(getattr(usage,"completion_tokens",0) or 0),
|
||||
"timestamp":time.strftime("%Y-%m-%dT%H:%M:%S%z")}
|
||||
with open(self._cost_log_path,"a",encoding="utf-8") as f: f.write(json.dumps(row,ensure_ascii=False)+"\n")
|
||||
except Exception as e:
|
||||
print(" V16 cost log failed: %s"%e,flush=True)
|
||||
def _split_sentences_v16(self,text):
|
||||
text=re.sub(r"\s+"," ",str(text or "")).strip()
|
||||
if not text: return []
|
||||
for a,b in [("et al.","et al<PRD>"),("e.g.","e<PRD>g<PRD>"),("i.e.","i<PRD>e<PRD>"),("Fig.","Fig<PRD>"),("Figs.","Figs<PRD>")]: text=text.replace(a,b)
|
||||
return [p.replace("<PRD>",".").strip() for p in re.split(r"(?<=[.!?])\s+(?=[A-Z0-9(\[])",text) if p.strip()]
|
||||
def _build_evidence_map(self,sents): return {"E%03d"%(i+1):s for i,s in enumerate(sents)}
|
||||
def _split_pages(self,sents,limit):
|
||||
pages=[]; cur=[]; n=0
|
||||
for s in sents:
|
||||
if cur and n+len(s)+1>limit: pages.append(" ".join(cur)); cur=[cur[-1]]; n=len(cur[0])+1
|
||||
cur.append(s); n+=len(s)+1
|
||||
if cur: pages.append(" ".join(cur))
|
||||
return pages
|
||||
def _extract_candidates_llm1(self,window_id,window_text,section_title,pmid,source_label):
|
||||
sents=self._split_sentences_v16(window_text); evmap=self._build_evidence_map(sents); pages=self._split_pages(sents,self.PAGE_CHAR_LIMIT); out=[]
|
||||
for pi,page in enumerate(pages):
|
||||
data=self._call_page(window_id,pi,page,evmap,section_title,pmid,source_label)
|
||||
if not data: continue
|
||||
for ent in data.get("entities",[]) or []:
|
||||
ent=dict(ent); ent.update({"window_id":window_id,"paper_id":pmid,"section":section_title,"source_label":source_label}); self.v16_entity_rows.append(ent)
|
||||
for group in data.get("evidence_groups",[]) or []:
|
||||
eids=group.get("evidence_sentence_ids",[]) or []
|
||||
ev=" ".join(evmap.get(e,"") for e in eids).strip()
|
||||
for r in group.get("relations",[]) or []:
|
||||
c=dict(r); c["evidence_ids"]=eids; c["evidence"]=ev; c.setdefault("confidence",0.70); c.setdefault("relation_unit_id",window_id); c.setdefault("sentence_ids",eids); out.append(c); self.v16_relation_rows.append({**c,"window_id":window_id,"paper_id":pmid,"section":section_title,"source_label":source_label})
|
||||
# Backward-compatible top-level relations, if model returns old schema.
|
||||
for r in data.get("relations",[]) or []:
|
||||
c=dict(r); eids=c.get("evidence_ids",[]) or []; c["evidence"]=" ".join(evmap.get(e,"") for e in eids).strip() or c.get("evidence",""); c.setdefault("relation_unit_id",window_id); c.setdefault("sentence_ids",eids); out.append(c); self.v16_relation_rows.append({**c,"window_id":window_id,"paper_id":pmid,"section":section_title,"source_label":source_label})
|
||||
return out
|
||||
def _call_page(self,window_id,page_idx,page_text,evidence_map,section_title,pmid,source_label):
|
||||
ev_idx="\n".join("[%s] %s"%(eid,s) for eid,s in evidence_map.items())
|
||||
prompt="\n".join([
|
||||
"Extract explicit aquaculture-science entities and SRT relations from the provided numbered evidence sentences.",
|
||||
"Extract all explicit aquaculture facts; do not only pick examples.",
|
||||
"Allowed entity types: "+", ".join(sorted(ALLOWED_ENTITY_TYPES)),
|
||||
"Allowed relation categories: "+", ".join(sorted(ALLOWED_RELATION_CATEGORIES)),
|
||||
"Rules:",
|
||||
"1. Entity names must be meaningful biological, aquacultural, environmental, methodological, or dataset entities, not bare numbers.",
|
||||
"2. Preserve negation, direction, treatment, tissue, time point, species, pathogen, and environmental context when visible.",
|
||||
"3. Figure caption evidence is allowed, but set layer_hint to caption_relations or candidate_review.",
|
||||
"4. Table facts may be extracted only from explicit row/column text.",
|
||||
"5. Do not treat promoter as promote.",
|
||||
"6. Each relation must express one atomic claim. Split clauses joined by but, whereas, while, or however into separate triples.",
|
||||
"7. Keep relation labels concise predicates. Do not place target context, contrast clauses, evidence prose, or complete sentences in relation.",
|
||||
"8. Source and target must be literal entities visible in the cited evidence sentence, except for direct expansion of an explicit range or enumeration.",
|
||||
"9. Return empty arrays only if no relation exists.",
|
||||
"10. Use canonical snake_case relation labels such as improves_growth_performance, reduces_fcr, increases_survival_rate, increases_resistance_to, enhances_immune_response, improves_water_quality, upregulates, downregulates, responds_to, affects, participates_in, measured_with, has_accession; keep each label at most 4 words.",
|
||||
"11. Map relation categories to the category field: growth metrics -> growth_performance/feed_conversion_efficiency/survival_rate; immune -> immune_response; water -> water_quality_effect; expression -> expression_tissue/expression_condition; metabolites -> metabolic_change; traits/SNPs -> variant_trait_association; methodology -> method_tool; otherwise phenotype_effect/pathway_participation/physiological_measurement/other_supported as applicable.",
|
||||
"", "paper_id: "+str(pmid), "relation_unit_id: "+str(window_id), "section: "+str(section_title), "source_label: "+str(source_label), "page: "+str(page_idx+1),
|
||||
"", "EVIDENCE SENTENCES:", ev_idx,
|
||||
"", "PAGE TEXT:", page_text,
|
||||
"", "Return JSON only using this schema:",
|
||||
'{"entities":[{"name":"","type":"","surface":"","sentence_ids":["E001"]}],"evidence_groups":[{"evidence_sentence_ids":["E001"],"relations":[{"source":"","source_type":"","relation":"","relation_category":"","target":"","target_type":"","claim_scope":"paper_finding|paper_interpretation|background|method_context|unknown","layer_hint":"main_kg|candidate_review|caption_relations|background_relations|method_relations|dataset_metadata","negated":false,"condition":"","species_context":"","confidence":0.0}]}]}'
|
||||
])
|
||||
self.stats["llm_requests"]+=1
|
||||
try:
|
||||
kwargs={"model":self.model,"messages":[{"role":"system","content":V16_SYSTEM_PROMPT},{"role":"user","content":prompt}],"temperature":0.1,"max_tokens":self.llm_max_tokens,"timeout":self.request_timeout}
|
||||
if "deepseek" in str(self.model).lower() and os.environ.get("V15_RELATION_DISABLE_THINKING","1").strip().lower() not in {"0","false","no","off"}: kwargs["extra_body"]={"thinking":{"type":"disabled"}}
|
||||
with self._hard_timeout(float(self.request_timeout)+5): resp=self._client.chat.completions.create(**kwargs)
|
||||
self._record_cost(pmid,resp)
|
||||
content=(resp.choices[0].message.content or "").strip(); self.raw_response_rows.append({"window_id":f"{window_id}#page{page_idx+1}","pmid":pmid,"section":section_title,"source_label":source_label,"status":"ok","response":content,"error":""})
|
||||
if not content: return None
|
||||
t=re.sub(r"^```(?:json)?","",content,flags=re.I).strip(); t=re.sub(r"```$","",t).strip()
|
||||
try: return json.loads(self._fix_json(t))
|
||||
except Exception:
|
||||
try:
|
||||
import json_repair
|
||||
return json_repair.loads(t)
|
||||
except Exception as e:
|
||||
self.stats["llm_failures"]+=1; self.raw_response_rows[-1]["error"]="json_parse_failed: "+str(e); print(" V16 LLM-1 JSON parse fail: %s"%e,flush=True); return None
|
||||
except Exception as e:
|
||||
self.raw_response_rows.append({"window_id":f"{window_id}#page{page_idx+1}","pmid":pmid,"section":section_title,"source_label":source_label,"status":"error","response":"","error":str(e)})
|
||||
self.stats["llm_failures"]+=1; print(" V16 LLM-1 page fail: %s"%e,flush=True); return None
|
||||
def _candidates_to_relations(self,candidates,window_text,section_title,pmid,source_label,validation_status="v16_llm1_rule_filtered"):
|
||||
kept=[]; rejected=[]; wn=re.sub(r"\s+"," ",window_text or "").lower()
|
||||
for i,c in enumerate(candidates or [],1):
|
||||
src=clean_text(c.get("source")); tgt=clean_text(c.get("target")); raw_rel=clean_text(c.get("relation")); ev=clean_text(c.get("evidence")); st=clean_text(c.get("source_type")) or "unknown"; tt=clean_text(c.get("target_type")) or "unknown"; conf=float(c.get("confidence",0) or 0)
|
||||
cid=str(c.get("candidate_id") or f"v16_{i:05d}"); reasons=hard_reject_reasons(src,st,raw_rel,tgt,tt,ev)
|
||||
if ev and re.sub(r"\s+"," ",ev).lower() not in wn: reasons.append("evidence_no_match")
|
||||
if reasons:
|
||||
row={"candidate_id":cid,"reason":";".join(reasons),"source":src,"relation":raw_rel,"target":tgt,"evidence":ev}; rejected.append(row); self.v16_rejected_rows.append(row); continue
|
||||
norm_rel,norm_reason=normalize_relation(raw_rel,c.get("relation_category",""),tt); endpoint_issues=entity_issues(src,st,"source")+entity_issues(tgt,tt,"target"); validation_reason=";".join([x for x in endpoint_issues if "unsupported_type" in x or "generic_fragment" in x or "starts_with" in x])
|
||||
layer=infer_layer(st,c.get("relation_category",""),c.get("layer_hint",""),source_label,c.get("claim_scope",""),conf,bool(validation_reason))
|
||||
r=Relation(source=src,source_type=st,relation=norm_rel,target=tgt,target_type=tt,evidence=ev,section=section_title,confidence=conf,method="llm_srt_v16",pubmed_id=pmid or "",p_source=f"{pmid}:{section_title}" if pmid else section_title,source_label=source_label or "text")
|
||||
r.candidate_id=cid; r.raw_relation=raw_rel; r.normalized_relation=norm_rel; r.relation_category=clean_text(c.get("relation_category")); r.normalization_reason=norm_reason; r.claim_scope=clean_text(c.get("claim_scope")); r.layer_hint=clean_text(c.get("layer_hint")); r.negated=bool(c.get("negated",False)); r.condition=clean_text(c.get("condition")); r.species_context=clean_text(c.get("species_context")); r.sentence_ids=c.get("sentence_ids") or c.get("evidence_ids") or []; r.relation_unit_id=c.get("relation_unit_id") or ""; r.source_block_ids=[]; r.block_types=[]; r.validation_reason=validation_reason; r.output_layer=layer; r.qa_weight=qa_weight(layer); r.candidate_status="candidate" if layer in {"candidate_review","caption_relations"} else "final"; r.triple_quality="candidate" if r.candidate_status=="candidate" else "supported"; r.validation_status=validation_status
|
||||
kept.append(r)
|
||||
return kept,rejected
|
||||
def save_audit_files(self,output_dir:Optional[Path]=None)->Dict:
|
||||
paths=super().save_audit_files(output_dir); out=Path(output_dir or self.output_dir or "."); out.mkdir(parents=True,exist_ok=True)
|
||||
def jsonl(name,rows):
|
||||
p=out/name
|
||||
with open(p,"w",encoding="utf-8") as f:
|
||||
for row in rows: f.write(json.dumps(row,ensure_ascii=False)+"\n")
|
||||
return p
|
||||
paths["entity_extraction_raw"]=jsonl("entity_extraction_raw.jsonl",self.v16_entity_rows)
|
||||
paths["llm1_parsed_relations"]=jsonl("llm1_parsed_relations.jsonl",self.v16_relation_rows)
|
||||
fields=["candidate_id","reason","source","relation","target","evidence"]
|
||||
with open(out/"relation_rejected_audit.csv","w",newline="",encoding="utf-8") as f:
|
||||
w=csv.DictWriter(f,fieldnames=fields); w.writeheader(); w.writerows(self.v16_rejected_rows)
|
||||
# Compatibility audit names requested by V16.1.1 plan.
|
||||
for name,rows,cols in [("entity_validation_audit.csv",self.v16_entity_rows,["window_id","paper_id","section","source_label","name","type","surface","sentence_ids"]),("relation_support_audit.csv",self.v16_relation_rows,["window_id","paper_id","section","source_label","source","relation","target","evidence","evidence_ids","confidence"]),("relation_normalization_audit.csv",self.v16_relation_rows,["window_id","paper_id","section","source","relation","target","relation_category","confidence"] )]:
|
||||
with open(out/name,"w",newline="",encoding="utf-8") as f:
|
||||
w=csv.DictWriter(f,fieldnames=cols); w.writeheader();
|
||||
for row in rows: w.writerow({c:row.get(c,"") for c in cols})
|
||||
def _json_safe(value):
|
||||
from pathlib import Path as _Path
|
||||
if isinstance(value, _Path):
|
||||
return str(value)
|
||||
if isinstance(value, dict):
|
||||
return {k: _json_safe(v) for k, v in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [_json_safe(v) for v in value]
|
||||
return value
|
||||
return _json_safe(paths)
|
||||
|
||||
LLMRelationExtractorV15_2_1 = PlantScienceLLMExtractorV16_1_1
|
||||
+392
@@ -0,0 +1,392 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from .entity_validator_v16_1_1 import (
|
||||
ALLOWED_ENTITY_TYPES,
|
||||
clean_text,
|
||||
entity_issues,
|
||||
normalize_entity_type,
|
||||
)
|
||||
from .phase4_llm_srt_extraction_v16_1_1 import (
|
||||
PlantScienceLLMExtractorV16_1_1,
|
||||
)
|
||||
|
||||
|
||||
class AquaScienceLLMExtractorV16_aqua_1_0(PlantScienceLLMExtractorV16_1_1):
|
||||
"""V16.aqua.1.0 evidence-local paging and entity validation (aquaculture SRT lane)."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.llm1_input_rows: List[Dict] = []
|
||||
self.validated_entity_rows: List[Dict] = []
|
||||
self.rejected_entity_rows: List[Dict] = []
|
||||
self.support_audit_rows: List[Dict] = []
|
||||
self.normalization_audit_rows: List[Dict] = []
|
||||
|
||||
def _local_bridge_relations(self, *args, **kwargs):
|
||||
"""Disable V15 gene-family bridge rules in the full-plant V16 lane."""
|
||||
return []
|
||||
|
||||
def _split_pages_with_ids(
|
||||
self, sentences: List[str], limit: int
|
||||
) -> List[List[Tuple[str, str]]]:
|
||||
pages: List[List[Tuple[str, str]]] = []
|
||||
current: List[Tuple[str, str]] = []
|
||||
current_chars = 0
|
||||
for index, sentence in enumerate(sentences, 1):
|
||||
item = (f"E{index:03d}", sentence)
|
||||
if current and current_chars + len(sentence) + 1 > limit:
|
||||
pages.append(current)
|
||||
current = current[-1:]
|
||||
current_chars = sum(len(value) + 1 for _, value in current)
|
||||
current.append(item)
|
||||
current_chars += len(sentence) + 1
|
||||
if current:
|
||||
pages.append(current)
|
||||
return pages
|
||||
|
||||
def _extract_candidates_llm1(
|
||||
self,
|
||||
window_id,
|
||||
window_text,
|
||||
section_title,
|
||||
pmid,
|
||||
source_label,
|
||||
):
|
||||
sentences = self._split_sentences_v16(window_text)
|
||||
pages = self._split_pages_with_ids(sentences, self.PAGE_CHAR_LIMIT)
|
||||
output = []
|
||||
for page_index, page_items in enumerate(pages):
|
||||
evidence_map = dict(page_items)
|
||||
page_text = " ".join(text for _, text in page_items)
|
||||
data = self._call_page(
|
||||
window_id,
|
||||
page_index,
|
||||
page_text,
|
||||
evidence_map,
|
||||
section_title,
|
||||
pmid,
|
||||
source_label,
|
||||
)
|
||||
if not data:
|
||||
continue
|
||||
for entity in data.get("entities", []) or []:
|
||||
row = self._validate_entity(
|
||||
entity,
|
||||
window_id,
|
||||
pmid,
|
||||
section_title,
|
||||
source_label,
|
||||
evidence_map,
|
||||
)
|
||||
self.v16_entity_rows.append(row)
|
||||
for group in data.get("evidence_groups", []) or []:
|
||||
evidence_ids = [
|
||||
value for value in group.get("evidence_sentence_ids", []) or []
|
||||
if value in evidence_map
|
||||
]
|
||||
evidence = " ".join(evidence_map[value] for value in evidence_ids).strip()
|
||||
for relation in group.get("relations", []) or []:
|
||||
candidate = dict(relation)
|
||||
candidate["evidence_ids"] = evidence_ids
|
||||
candidate["evidence"] = evidence
|
||||
candidate.setdefault("confidence", 0.70)
|
||||
candidate.setdefault("relation_unit_id", window_id)
|
||||
candidate.setdefault("sentence_ids", evidence_ids)
|
||||
output.append(candidate)
|
||||
self.v16_relation_rows.append(
|
||||
{
|
||||
**candidate,
|
||||
"window_id": window_id,
|
||||
"paper_id": pmid,
|
||||
"section": section_title,
|
||||
"source_label": source_label,
|
||||
}
|
||||
)
|
||||
for relation in data.get("relations", []) or []:
|
||||
candidate = dict(relation)
|
||||
evidence_ids = [
|
||||
value for value in candidate.get("evidence_ids", []) or []
|
||||
if value in evidence_map
|
||||
]
|
||||
candidate["evidence"] = (
|
||||
" ".join(evidence_map[value] for value in evidence_ids).strip()
|
||||
or candidate.get("evidence", "")
|
||||
)
|
||||
candidate.setdefault("relation_unit_id", window_id)
|
||||
candidate.setdefault("sentence_ids", evidence_ids)
|
||||
output.append(candidate)
|
||||
self.v16_relation_rows.append(
|
||||
{
|
||||
**candidate,
|
||||
"window_id": window_id,
|
||||
"paper_id": pmid,
|
||||
"section": section_title,
|
||||
"source_label": source_label,
|
||||
}
|
||||
)
|
||||
return output
|
||||
|
||||
def _call_page(
|
||||
self,
|
||||
window_id,
|
||||
page_idx,
|
||||
page_text,
|
||||
evidence_map,
|
||||
section_title,
|
||||
pmid,
|
||||
source_label,
|
||||
):
|
||||
self.llm1_input_rows.append(
|
||||
{
|
||||
"window_id": window_id,
|
||||
"page": page_idx + 1,
|
||||
"paper_id": pmid,
|
||||
"section": section_title,
|
||||
"source_label": source_label,
|
||||
"evidence_sentence_ids": list(evidence_map),
|
||||
"evidence_sentences": evidence_map,
|
||||
"char_count": len(page_text),
|
||||
}
|
||||
)
|
||||
return super()._call_page(
|
||||
window_id,
|
||||
page_idx,
|
||||
page_text,
|
||||
evidence_map,
|
||||
section_title,
|
||||
pmid,
|
||||
source_label,
|
||||
)
|
||||
|
||||
def _validate_entity(
|
||||
self,
|
||||
entity: Dict,
|
||||
window_id: str,
|
||||
paper_id: str,
|
||||
section: str,
|
||||
source_label: str,
|
||||
evidence_map: Dict[str, str],
|
||||
) -> Dict:
|
||||
name = clean_text(entity.get("name") or entity.get("surface"))
|
||||
surface = clean_text(entity.get("surface")) or name
|
||||
entity_type = normalize_entity_type(entity.get("type"))
|
||||
sentence_ids = [
|
||||
value for value in entity.get("sentence_ids", []) or []
|
||||
if value in evidence_map
|
||||
]
|
||||
issues = entity_issues(name, entity_type, "entity")
|
||||
visible = (
|
||||
any(surface.lower() in evidence_map[value].lower() for value in sentence_ids)
|
||||
if surface and sentence_ids
|
||||
else False
|
||||
)
|
||||
if surface and not visible:
|
||||
visible = any(surface.lower() in text.lower() for text in evidence_map.values())
|
||||
if not visible:
|
||||
issues.append("entity_surface_not_in_local_evidence")
|
||||
row = {
|
||||
"window_id": window_id,
|
||||
"paper_id": paper_id,
|
||||
"section": section,
|
||||
"source_label": source_label,
|
||||
"name": name,
|
||||
"type": entity_type,
|
||||
"surface": surface,
|
||||
"sentence_ids": sentence_ids,
|
||||
"action": "reject" if issues else "accept",
|
||||
"reason": ";".join(issues),
|
||||
}
|
||||
if issues:
|
||||
self.rejected_entity_rows.append(row)
|
||||
else:
|
||||
self.validated_entity_rows.append(row)
|
||||
return row
|
||||
|
||||
def _candidates_to_relations(
|
||||
self,
|
||||
candidates,
|
||||
window_text,
|
||||
section_title,
|
||||
pmid,
|
||||
source_label,
|
||||
validation_status="v16_aqua_llm1_rule_filtered",
|
||||
):
|
||||
normalized = []
|
||||
for candidate in candidates or []:
|
||||
value = dict(candidate)
|
||||
value["source_type"] = normalize_entity_type(value.get("source_type"))
|
||||
value["target_type"] = normalize_entity_type(value.get("target_type"))
|
||||
normalized.append(value)
|
||||
kept, rejected = super()._candidates_to_relations(
|
||||
normalized,
|
||||
window_text,
|
||||
section_title,
|
||||
pmid,
|
||||
source_label,
|
||||
validation_status,
|
||||
)
|
||||
for relation in kept:
|
||||
quality_issues = [
|
||||
value
|
||||
for value in str(getattr(relation, "validation_reason", "") or "").split(";")
|
||||
if value
|
||||
]
|
||||
evidence_lower = re.sub(r"\s+", " ", relation.evidence or "").lower()
|
||||
for role, endpoint in [
|
||||
("source", relation.source),
|
||||
("target", relation.target),
|
||||
]:
|
||||
endpoint_lower = re.sub(r"\s+", " ", endpoint or "").lower()
|
||||
if endpoint_lower and endpoint_lower not in evidence_lower:
|
||||
quality_issues.append(f"{role}_not_literal_in_evidence")
|
||||
raw_relation = getattr(relation, "raw_relation", relation.relation) or ""
|
||||
if len(raw_relation.split()) > 8 or re.search(
|
||||
r"\b(?:but|whereas|while|however)\b", raw_relation, re.I
|
||||
):
|
||||
quality_issues.append("non_atomic_relation")
|
||||
quality_issues = list(dict.fromkeys(quality_issues))
|
||||
relation.validation_reason = ";".join(quality_issues)
|
||||
if quality_issues and relation.output_layer not in {
|
||||
"caption_relations",
|
||||
"method_relations",
|
||||
"dataset_metadata",
|
||||
"background_relations",
|
||||
}:
|
||||
relation.output_layer = "candidate_review"
|
||||
relation.candidate_status = "candidate"
|
||||
relation.triple_quality = "candidate"
|
||||
self.support_audit_rows.append(
|
||||
{
|
||||
"candidate_id": getattr(relation, "candidate_id", ""),
|
||||
"source": relation.source,
|
||||
"source_type": relation.source_type,
|
||||
"raw_relation": getattr(relation, "raw_relation", relation.relation),
|
||||
"normalized_relation": relation.relation,
|
||||
"target": relation.target,
|
||||
"target_type": relation.target_type,
|
||||
"evidence": relation.evidence,
|
||||
"action": "supported",
|
||||
"reason": getattr(relation, "validation_reason", ""),
|
||||
"output_layer": getattr(relation, "output_layer", ""),
|
||||
}
|
||||
)
|
||||
self.normalization_audit_rows.append(
|
||||
{
|
||||
"candidate_id": getattr(relation, "candidate_id", ""),
|
||||
"source": relation.source,
|
||||
"raw_relation": getattr(relation, "raw_relation", relation.relation),
|
||||
"normalized_relation": relation.relation,
|
||||
"target": relation.target,
|
||||
"relation_category": getattr(relation, "relation_category", ""),
|
||||
"normalization_reason": getattr(
|
||||
relation, "normalization_reason", ""
|
||||
),
|
||||
}
|
||||
)
|
||||
for row in rejected:
|
||||
self.support_audit_rows.append(
|
||||
{
|
||||
"candidate_id": row.get("candidate_id", ""),
|
||||
"source": row.get("source", ""),
|
||||
"source_type": "",
|
||||
"raw_relation": row.get("relation", ""),
|
||||
"normalized_relation": "",
|
||||
"target": row.get("target", ""),
|
||||
"target_type": "",
|
||||
"evidence": row.get("evidence", ""),
|
||||
"action": "rejected",
|
||||
"reason": row.get("reason", ""),
|
||||
"output_layer": "rejected",
|
||||
}
|
||||
)
|
||||
return kept, rejected
|
||||
|
||||
def save_audit_files(self, output_dir: Optional[Path] = None) -> Dict:
|
||||
paths = super().save_audit_files(output_dir)
|
||||
out = Path(output_dir or self.output_dir or ".")
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
with (out / "llm1_input.jsonl").open("w", encoding="utf-8") as handle:
|
||||
for row in self.llm1_input_rows:
|
||||
handle.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
with (out / "llm1_raw_response.jsonl").open("w", encoding="utf-8") as handle:
|
||||
for row in self.raw_response_rows:
|
||||
handle.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
fields = [
|
||||
"window_id",
|
||||
"paper_id",
|
||||
"section",
|
||||
"source_label",
|
||||
"name",
|
||||
"type",
|
||||
"surface",
|
||||
"sentence_ids",
|
||||
"action",
|
||||
"reason",
|
||||
]
|
||||
for filename, rows in [
|
||||
("entity_validation_audit.csv", self.validated_entity_rows + self.rejected_entity_rows),
|
||||
("entity_rejected_audit.csv", self.rejected_entity_rows),
|
||||
]:
|
||||
with (out / filename).open("w", newline="", encoding="utf-8") as handle:
|
||||
writer = csv.DictWriter(handle, fieldnames=fields)
|
||||
writer.writeheader()
|
||||
for row in rows:
|
||||
writer.writerow({field: row.get(field, "") for field in fields})
|
||||
support_fields = [
|
||||
"candidate_id",
|
||||
"source",
|
||||
"source_type",
|
||||
"raw_relation",
|
||||
"normalized_relation",
|
||||
"target",
|
||||
"target_type",
|
||||
"evidence",
|
||||
"action",
|
||||
"reason",
|
||||
"output_layer",
|
||||
]
|
||||
with (out / "relation_support_audit.csv").open(
|
||||
"w", newline="", encoding="utf-8"
|
||||
) as handle:
|
||||
writer = csv.DictWriter(handle, fieldnames=support_fields)
|
||||
writer.writeheader()
|
||||
writer.writerows(self.support_audit_rows)
|
||||
normalization_fields = [
|
||||
"candidate_id",
|
||||
"source",
|
||||
"raw_relation",
|
||||
"normalized_relation",
|
||||
"target",
|
||||
"relation_category",
|
||||
"normalization_reason",
|
||||
]
|
||||
with (out / "relation_normalization_audit.csv").open(
|
||||
"w", newline="", encoding="utf-8"
|
||||
) as handle:
|
||||
writer = csv.DictWriter(handle, fieldnames=normalization_fields)
|
||||
writer.writeheader()
|
||||
writer.writerows(self.normalization_audit_rows)
|
||||
paths.setdefault("audit_files", {}).update(
|
||||
{
|
||||
"llm1_input": str(out / "llm1_input.jsonl"),
|
||||
"llm1_raw_response": str(out / "llm1_raw_response.jsonl"),
|
||||
"entity_validation_audit": str(out / "entity_validation_audit.csv"),
|
||||
"entity_rejected_audit": str(out / "entity_rejected_audit.csv"),
|
||||
"relation_support_audit": str(out / "relation_support_audit.csv"),
|
||||
"relation_normalization_audit": str(
|
||||
out / "relation_normalization_audit.csv"
|
||||
),
|
||||
}
|
||||
)
|
||||
paths["validated_llm_entities"] = len(self.validated_entity_rows)
|
||||
paths["rejected_llm_entities"] = len(self.rejected_entity_rows)
|
||||
return paths
|
||||
|
||||
|
||||
LLMRelationExtractorV16_aqua_1_0 = AquaScienceLLMExtractorV16_aqua_1_0
|
||||
+1067
File diff suppressed because it is too large
Load Diff
+432
@@ -0,0 +1,432 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Phase 5: 后处理
|
||||
整合fusion的证据验证、去重 + v12的gene-centric chunks
|
||||
"""
|
||||
|
||||
import csv
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Set, Tuple
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
# V14.2: 使用phase4的完整Relation定义
|
||||
from .phase4_relation_extraction import Relation
|
||||
|
||||
|
||||
@dataclass
|
||||
class Chunk:
|
||||
"""Gene-centric chunk"""
|
||||
chunk_id: str
|
||||
gene_id: str
|
||||
text: str
|
||||
section: str
|
||||
relations: List[Relation] = field(default_factory=list)
|
||||
|
||||
|
||||
class PostProcessor:
|
||||
"""后处理器 - 整合fusion+v12能力"""
|
||||
|
||||
def __init__(self):
|
||||
self.duplicate_audit_rows: List[Dict] = []
|
||||
|
||||
def process(self, text: str, relations: List[Relation]) -> Tuple[List[Relation], List[Chunk]]:
|
||||
"""
|
||||
后处理流程
|
||||
|
||||
Returns:
|
||||
(validated_relations, gene_centric_chunks)
|
||||
"""
|
||||
print(" Step 5.1: 验证证据...")
|
||||
validated = self._validate_evidence(relations, text)
|
||||
print(f" ✓ 验证 {len(validated)} 个关系")
|
||||
|
||||
print(" Step 5.2: 去重...")
|
||||
unique = self._deduplicate(validated)
|
||||
print(f" ✓ 去重后 {len(unique)} 个关系")
|
||||
|
||||
print(" Step 5.3: 生成gene-centric chunks...")
|
||||
chunks = self._create_gene_centric_chunks(text, unique)
|
||||
print(f" ✓ 生成 {len(chunks)} 个chunks")
|
||||
|
||||
return unique, chunks
|
||||
|
||||
def _validate_evidence(self, relations: List[Relation], text: str) -> List[Relation]:
|
||||
"""
|
||||
验证证据(fusion方法)
|
||||
|
||||
确保证据文本确实存在于原文中
|
||||
"""
|
||||
validated = []
|
||||
text_lower = text.lower()
|
||||
|
||||
for rel in relations:
|
||||
evidence_lower = rel.evidence.lower()
|
||||
# 检查证据是否在原文中
|
||||
if evidence_lower in text_lower or len(evidence_lower) > 20:
|
||||
validated.append(rel)
|
||||
|
||||
return validated
|
||||
|
||||
def _deduplicate(self, relations: List[Relation]) -> List[Relation]:
|
||||
"""
|
||||
去重(fusion方法)
|
||||
|
||||
基于(source, relation, target)去重
|
||||
"""
|
||||
self.duplicate_audit_rows = []
|
||||
unique: List[Relation] = []
|
||||
|
||||
for rel in relations:
|
||||
duplicate_idx, duplicate_reason = self._find_semantic_duplicate(unique, rel)
|
||||
if duplicate_idx is None:
|
||||
unique.append(rel)
|
||||
continue
|
||||
|
||||
kept = unique[duplicate_idx]
|
||||
preferred = self._prefer_relation(kept, rel)
|
||||
removed = rel if preferred is kept else kept
|
||||
if preferred is not kept:
|
||||
unique[duplicate_idx] = rel
|
||||
self.duplicate_audit_rows.append({
|
||||
"action": "merged",
|
||||
"reason": duplicate_reason,
|
||||
"kept_source": preferred.source,
|
||||
"kept_relation": preferred.relation,
|
||||
"kept_target": preferred.target,
|
||||
"kept_method": getattr(preferred, "method", ""),
|
||||
"kept_source_label": getattr(preferred, "source_label", ""),
|
||||
"kept_claim_scope": getattr(preferred, "claim_scope", ""),
|
||||
"removed_source": removed.source,
|
||||
"removed_relation": removed.relation,
|
||||
"removed_target": removed.target,
|
||||
"removed_method": getattr(removed, "method", ""),
|
||||
"removed_source_label": getattr(removed, "source_label", ""),
|
||||
"removed_claim_scope": getattr(removed, "claim_scope", ""),
|
||||
"kept_evidence": preferred.evidence,
|
||||
"removed_evidence": removed.evidence,
|
||||
})
|
||||
|
||||
return unique
|
||||
|
||||
def save_duplicate_audit(self, path: Path) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
fields: List[str] = []
|
||||
for row in self.duplicate_audit_rows:
|
||||
for key in row:
|
||||
if key not in fields:
|
||||
fields.append(key)
|
||||
with open(path, "w", newline="", encoding="utf-8-sig") as handle:
|
||||
if not fields:
|
||||
handle.write("")
|
||||
return
|
||||
writer = csv.DictWriter(handle, fieldnames=fields, extrasaction="ignore")
|
||||
writer.writeheader()
|
||||
writer.writerows(self.duplicate_audit_rows)
|
||||
|
||||
def _find_semantic_duplicate(self, unique: List[Relation], rel: Relation) -> Tuple[object, str]:
|
||||
"""Find only strict endpoint duplicates.
|
||||
|
||||
Earlier versions merged relations with the same source and a
|
||||
semantically overlapping target when their relation labels belonged to a
|
||||
compatible family. That removed too much information during manual gold
|
||||
review. For the high-recall SRT pool, only merge when SOURCE and TARGET
|
||||
normalize to exactly the same text. The relation label is used only to
|
||||
avoid merging clearly opposite facts such as upregulated_in vs
|
||||
downregulated_in.
|
||||
"""
|
||||
rel_source = self._norm_endpoint(rel.source)
|
||||
rel_target = self._norm_endpoint(rel.target)
|
||||
rel_family = self._relation_family(rel.relation)
|
||||
for idx, existing in enumerate(unique):
|
||||
if self._norm_endpoint(existing.source) != rel_source:
|
||||
continue
|
||||
existing_target = self._norm_endpoint(existing.target)
|
||||
if existing_target == rel_target:
|
||||
existing_family = self._relation_family(existing.relation)
|
||||
if self._relation_polarity_conflict(existing_family, rel_family):
|
||||
continue
|
||||
if self._norm_endpoint(existing.relation) == self._norm_endpoint(rel.relation):
|
||||
return idx, "same_source_relation_target_exact"
|
||||
if self._relation_families_compatible(existing_family, rel_family):
|
||||
return idx, "same_source_target_expression_synonym"
|
||||
return None, ""
|
||||
|
||||
def _prefer_relation(self, first: Relation, second: Relation) -> Relation:
|
||||
return first if self._relation_quality_score(first) >= self._relation_quality_score(second) else second
|
||||
|
||||
def _relation_quality_score(self, rel: Relation) -> int:
|
||||
score = 0
|
||||
relation = self._relation_family(rel.relation)
|
||||
relation_norm = self._norm_endpoint(rel.relation)
|
||||
if relation in {"upregulated", "downregulated"}:
|
||||
score += 10
|
||||
elif relation == "expression_change":
|
||||
score += 7
|
||||
elif relation == "expression":
|
||||
score += 5
|
||||
if relation_norm in {"contains_cis_element", "lacks_binding_site_for", "localizes_to"}:
|
||||
score += 6
|
||||
target_norm = self._norm_endpoint(rel.target)
|
||||
if len(str(rel.target or "").split()) > 1:
|
||||
score += 2
|
||||
if target_norm not in {"tissues", "tissue", "conditions", "condition", "stress", "treatments", "levels", "expression"}:
|
||||
score += 2
|
||||
if len(str(rel.evidence or "")) <= 650:
|
||||
score += 1
|
||||
score += self._source_priority(rel)
|
||||
if getattr(rel, "method", "") in {"rule_fallback", "pc_aware_typed_fallback"}:
|
||||
score -= 3
|
||||
return score
|
||||
|
||||
def _source_priority(self, rel: Relation) -> int:
|
||||
label = " ".join([
|
||||
str(getattr(rel, "source_label", "") or ""),
|
||||
str(getattr(rel, "section", "") or ""),
|
||||
str(getattr(rel, "claim_scope", "") or ""),
|
||||
]).lower()
|
||||
for key, value in [
|
||||
("supplement", 10),
|
||||
("table", 9),
|
||||
("result", 8),
|
||||
("abstract", 7),
|
||||
("discussion", 5),
|
||||
("background", 3),
|
||||
("caption", 1),
|
||||
]:
|
||||
if key in label:
|
||||
return value
|
||||
return 4
|
||||
|
||||
def _target_is_semantic_overlap(self, left: str, right: str) -> bool:
|
||||
if not left or not right:
|
||||
return False
|
||||
if left == right:
|
||||
return True
|
||||
generic = {
|
||||
"tissues", "tissue", "organs", "conditions", "condition", "stress",
|
||||
"treatments", "proteins", "genes", "expression", "levels", "response",
|
||||
"responses",
|
||||
}
|
||||
if left in generic or right in generic:
|
||||
return True
|
||||
return left in right or right in left
|
||||
|
||||
def _relation_family(self, relation: str) -> str:
|
||||
text = self._norm_endpoint(relation)
|
||||
if "positive" in text or text in {"promotes", "enhances", "induces"}:
|
||||
return "positive_regulation"
|
||||
if "negative" in text or text in {"inhibits", "suppresses", "reduces"}:
|
||||
return "negative_regulation"
|
||||
if (
|
||||
"downregulat" in text
|
||||
or "decreas" in text
|
||||
or "reduc" in text
|
||||
or "repress" in text
|
||||
):
|
||||
return "downregulated"
|
||||
if (
|
||||
"upregulat" in text
|
||||
or "increas" in text
|
||||
or "induc" in text
|
||||
or "enhanc" in text
|
||||
):
|
||||
return "upregulated"
|
||||
if "differential" in text or "expression changed" in text or "changed in" in text:
|
||||
return "expression_change"
|
||||
if "express" in text or text in {"regulated", "regulates"}:
|
||||
return "expression"
|
||||
if "lack" in text or "not targeted" in text or "without" in text:
|
||||
return "absence"
|
||||
if "contain" in text or "harbor" in text or "binding_site" in text:
|
||||
return "containment"
|
||||
if "localiz" in text:
|
||||
return "localization"
|
||||
return text
|
||||
|
||||
def _relation_families_compatible(self, left: str, right: str) -> bool:
|
||||
if left == right:
|
||||
return True
|
||||
expression_families = {"expression", "expression_change", "upregulated", "downregulated"}
|
||||
return left in expression_families and right in expression_families
|
||||
|
||||
def _relation_polarity_conflict(self, left: str, right: str) -> bool:
|
||||
polarity_pairs = {
|
||||
("upregulated", "downregulated"),
|
||||
("positive_regulation", "negative_regulation"),
|
||||
("containment", "absence"),
|
||||
}
|
||||
return (left, right) in polarity_pairs or (right, left) in polarity_pairs
|
||||
|
||||
def _norm_endpoint(self, text: str) -> str:
|
||||
value = str(text or "").lower()
|
||||
value = re.sub(r"^(?:in|under|the|a|an)\s+", "", value)
|
||||
value = re.sub(r"[^a-z0-9]+", " ", value)
|
||||
return re.sub(r"\s+", " ", value).strip()
|
||||
|
||||
def _create_gene_centric_chunks(self, text: str, relations: List[Relation], sections: List = None) -> List[Chunk]:
|
||||
"""
|
||||
生成gene-centric chunks(v12方法)+ 全局章节chunks
|
||||
|
||||
为每个基因创建一个包含相关文本和关系的chunk
|
||||
同时保留关键章节(Abstract, Methods等)的完整内容
|
||||
"""
|
||||
chunks = []
|
||||
|
||||
# 按基因分组关系
|
||||
gene_relations: Dict[str, List[Relation]] = {}
|
||||
for rel in relations:
|
||||
if rel.source not in gene_relations:
|
||||
gene_relations[rel.source] = []
|
||||
gene_relations[rel.source].append(rel)
|
||||
|
||||
# 为每个基因创建chunk
|
||||
for gene_id, rels in gene_relations.items():
|
||||
# 提取包含该基因的句子
|
||||
relevant_text = self._extract_relevant_text(text, gene_id, rels)
|
||||
|
||||
# 获取主要章节
|
||||
rel_sections = [r.section for r in rels if r.section]
|
||||
main_section = rel_sections[0] if rel_sections else "unknown"
|
||||
|
||||
chunk = Chunk(
|
||||
chunk_id=f"chunk_{gene_id}",
|
||||
gene_id=gene_id,
|
||||
text=relevant_text,
|
||||
section=main_section,
|
||||
relations=rels
|
||||
)
|
||||
chunks.append(chunk)
|
||||
|
||||
# 添加全局章节chunks(Abstract, Methods等)
|
||||
section_chunks = self._create_section_chunks(text, sections)
|
||||
chunks.extend(section_chunks)
|
||||
|
||||
return chunks
|
||||
|
||||
def _create_section_chunks(self, text: str, sections: List = None) -> List[Chunk]:
|
||||
"""创建关键章节的完整内容chunks"""
|
||||
chunks = []
|
||||
|
||||
# 关键章节:这些章节包含全局信息,需要完整保留
|
||||
CRITICAL_SECTIONS = ['abstract', 'introduction', 'methods', 'materials', 'discussion']
|
||||
|
||||
# 如果没有提供章节信息,尝试从文本解析
|
||||
if not sections:
|
||||
sections = self._parse_sections_from_text(text)
|
||||
|
||||
for section in sections:
|
||||
section_name = section.get('title', '').lower() if isinstance(section, dict) else str(section).lower()
|
||||
section_text = section.get('text', '') if isinstance(section, dict) else ''
|
||||
|
||||
# 检查是否是关键章节
|
||||
if any(cs in section_name for cs in CRITICAL_SECTIONS):
|
||||
if section_text:
|
||||
chunks.append(Chunk(
|
||||
chunk_id=f"chunk_section_{section_name.replace(' ', '_')}",
|
||||
gene_id=f"SECTION_{section_name.upper().replace(' ', '_')}",
|
||||
text=section_text[:3000], # 限制长度避免过大
|
||||
section=section_name,
|
||||
relations=[]
|
||||
))
|
||||
|
||||
return chunks
|
||||
|
||||
def _parse_sections_from_text(self, text: str) -> List[dict]:
|
||||
"""从文本中解析章节(简单版本)"""
|
||||
sections = []
|
||||
|
||||
# 常见章节标题模式
|
||||
section_patterns = [
|
||||
(r'(?i)^\s*abstract\s*$', 'abstract'),
|
||||
(r'(?i)^\s*introduction\s*$', 'introduction'),
|
||||
(r'(?i)^\s*(materials?\s+and\s+methods?|methods?)\s*$', 'methods'),
|
||||
(r'(?i)^\s*results?\s*$', 'results'),
|
||||
(r'(?i)^\s*discussion\s*$', 'discussion'),
|
||||
(r'(?i)^\s*conclusion\s*$', 'conclusion'),
|
||||
]
|
||||
|
||||
lines = text.split('\n')
|
||||
current_section = None
|
||||
current_text = []
|
||||
|
||||
for line in lines:
|
||||
# 检查是否是章节标题
|
||||
for pattern, section_name in section_patterns:
|
||||
if re.match(pattern, line.strip()):
|
||||
# 保存上一个章节
|
||||
if current_section and current_text:
|
||||
sections.append({
|
||||
'title': current_section,
|
||||
'text': '\n'.join(current_text).strip()
|
||||
})
|
||||
# 开始新章节
|
||||
current_section = section_name
|
||||
current_text = []
|
||||
break
|
||||
else:
|
||||
# 不是标题,添加到当前章节
|
||||
if current_section:
|
||||
current_text.append(line)
|
||||
|
||||
# 保存最后一个章节
|
||||
if current_section and current_text:
|
||||
sections.append({
|
||||
'title': current_section,
|
||||
'text': '\n'.join(current_text).strip()
|
||||
})
|
||||
|
||||
return sections
|
||||
|
||||
def _extract_relevant_text(self, text: str, gene_id: str, relations: List[Relation]) -> str:
|
||||
"""提取与基因相关的文本"""
|
||||
# Improved sentence splitting: avoid splitting on abbreviations like "A. thaliana"
|
||||
# Pattern: split on .!? but not if preceded by single uppercase letter (abbreviation)
|
||||
sentences = []
|
||||
# First, protect abbreviations by temporarily replacing them
|
||||
protected_text = re.sub(r'([A-Z])\. ([a-z])', r'\1<DOT> \2', text)
|
||||
# Split sentences
|
||||
raw_sentences = [s.strip() for s in re.split(r"(?<=[.!?])\s+", protected_text) if len(s.strip()) > 20]
|
||||
# Restore abbreviations
|
||||
sentences = [s.replace('<DOT>', '.') for s in raw_sentences]
|
||||
relevant = []
|
||||
|
||||
gene_pattern = re.compile(rf'\b{re.escape(gene_id)}\b', re.IGNORECASE)
|
||||
|
||||
for sent in sentences:
|
||||
# 包含基因名的句子
|
||||
if gene_pattern.search(sent):
|
||||
relevant.append(sent)
|
||||
# 包含关系证据的句子
|
||||
else:
|
||||
for rel in relations:
|
||||
if rel.evidence and len(rel.evidence) > 10:
|
||||
evidence_words = rel.evidence.lower().split()[:5] # 前5个词
|
||||
if all(word in sent.lower() for word in evidence_words if len(word) > 3):
|
||||
relevant.append(sent)
|
||||
break
|
||||
|
||||
# 限制长度
|
||||
combined = " ".join(relevant[:10]) # 最多10个句子
|
||||
if len(combined) > 2000:
|
||||
combined = combined[:2000] + "..."
|
||||
|
||||
return combined
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 测试
|
||||
processor = PostProcessor()
|
||||
|
||||
test_text = "SlSut1 was highly expressed in leaves. SlSut2 was expressed in roots."
|
||||
test_relations = [
|
||||
Relation("SlSut1", "gene", "expressed_in", "leaves", "tissue", "SlSut1 was highly expressed in leaves", "Results"),
|
||||
Relation("SlSut2", "gene", "expressed_in", "roots", "tissue", "SlSut2 was expressed in roots", "Results"),
|
||||
]
|
||||
|
||||
relations, chunks = processor.process(test_text, test_relations)
|
||||
print(f"\n测试完成!")
|
||||
print(f"验证后关系: {len(relations)}")
|
||||
print(f"Chunks: {len(chunks)}")
|
||||
for c in chunks:
|
||||
print(f" - {c.gene_id}: {len(c.relations)} 个关系")
|
||||
@@ -0,0 +1,179 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Phase 5.5: 审核机制
|
||||
整合 v7 的审查功能到 V13
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import List, Dict, Tuple
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReviewResult:
|
||||
"""审核结果"""
|
||||
passed: bool
|
||||
issues: List[str]
|
||||
warnings: List[str]
|
||||
score: int
|
||||
severity: str # 'high', 'medium', 'low'
|
||||
|
||||
|
||||
class V15Reviewer:
|
||||
"""V13 审核器 - 基于 v7_reviewer 改进"""
|
||||
|
||||
@staticmethod
|
||||
def rule_review(entities: List[Dict], relations: List[Dict]) -> ReviewResult:
|
||||
"""规则审查"""
|
||||
issues = []
|
||||
warnings = []
|
||||
score = 60
|
||||
|
||||
# 检查实体
|
||||
gene_count = len([e for e in entities if e.get('type') == 'gene'])
|
||||
if gene_count == 0:
|
||||
issues.append("未提取到基因")
|
||||
elif gene_count > 1000:
|
||||
warnings.append(f"基因数{gene_count}过多,可能包含假阳性")
|
||||
score -= 10
|
||||
elif gene_count < 5:
|
||||
warnings.append(f"基因数{gene_count}过少,可能遗漏")
|
||||
score -= 5
|
||||
else:
|
||||
score += 10
|
||||
|
||||
# 检查关系
|
||||
rel_count = len(relations)
|
||||
if rel_count == 0:
|
||||
issues.append("未提取到任何关系")
|
||||
score -= 20
|
||||
elif rel_count < 10:
|
||||
warnings.append(f"关系数{rel_count}过少")
|
||||
score -= 5
|
||||
else:
|
||||
score += 10
|
||||
|
||||
# 检查关系类型多样性
|
||||
rel_types = set(r.get('relation') for r in relations)
|
||||
if len(rel_types) < 3:
|
||||
warnings.append(f"关系类型仅{len(rel_types)}种,多样性不足")
|
||||
score -= 5
|
||||
else:
|
||||
score += 10
|
||||
|
||||
# 检查高置信度关系比例
|
||||
high_conf = len([r for r in relations if r.get('confidence', 0) >= 0.9])
|
||||
if rel_count > 0 and high_conf / rel_count < 0.5:
|
||||
warnings.append(f"高置信度关系比例过低 ({high_conf}/{rel_count})")
|
||||
|
||||
severity = 'high' if issues else 'medium' if warnings else 'low'
|
||||
|
||||
return ReviewResult(
|
||||
passed=len(issues) == 0,
|
||||
issues=issues,
|
||||
warnings=warnings,
|
||||
score=max(0, min(100, score)),
|
||||
severity=severity
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def cross_validation(entities: List[Dict], relations: List[Dict], text: str) -> Dict:
|
||||
"""交叉验证 - 验证提取结果与原文一致性"""
|
||||
validations = []
|
||||
boost = 0
|
||||
|
||||
# 验证基因数
|
||||
gene_count = len([e for e in entities if e.get('type') == 'gene'])
|
||||
if gene_count > 0:
|
||||
patterns = [
|
||||
rf'{gene_count}\s+genes?',
|
||||
rf'{gene_count}\s+members?',
|
||||
rf'{gene_count}\s+SPL',
|
||||
]
|
||||
found = sum(1 for p in patterns if re.search(p, text, re.I))
|
||||
if found >= 1:
|
||||
validations.append(f"基因数{gene_count}在原文中得到验证")
|
||||
boost += 10
|
||||
|
||||
# 验证家族名
|
||||
family = next((e.get('name') for e in entities if e.get('type') == 'gene_family'), None)
|
||||
if family:
|
||||
count = text.lower().count(family.lower())
|
||||
if count >= 3:
|
||||
validations.append(f"家族名'{family}'出现{count}次")
|
||||
boost += 5
|
||||
|
||||
# 验证物种
|
||||
species = next((e.get('name') for e in entities if e.get('type') == 'species'), None)
|
||||
if species and species.lower() in text.lower():
|
||||
validations.append(f"物种'{species}'在原文中得到确认")
|
||||
boost += 5
|
||||
|
||||
# 验证关键关系
|
||||
key_rels = ['identified_in', 'expressed_in', 'collinear_with']
|
||||
for rel in key_rels:
|
||||
if any(r.get('relation') == rel for r in relations):
|
||||
validations.append(f"关键关系'{rel}'已提取")
|
||||
boost += 3
|
||||
|
||||
return {
|
||||
'validations': validations,
|
||||
'boost': boost,
|
||||
'score_boost': min(20, boost)
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def final_assessment(review: ReviewResult, validation: Dict) -> Dict:
|
||||
"""最终评估"""
|
||||
final_score = review.score + validation.get('score_boost', 0)
|
||||
final_score = max(0, min(100, final_score))
|
||||
|
||||
assessment = {
|
||||
'score': final_score,
|
||||
'grade': 'A' if final_score >= 90 else 'B' if final_score >= 75 else 'C' if final_score >= 60 else 'D',
|
||||
'passed': review.passed and final_score >= 60,
|
||||
'issues': review.issues,
|
||||
'warnings': review.warnings,
|
||||
'validations': validation.get('validations', []),
|
||||
'severity': review.severity,
|
||||
'recommendations': []
|
||||
}
|
||||
|
||||
# 生成建议
|
||||
if final_score < 60:
|
||||
assessment['recommendations'].append("建议人工审核提取结果")
|
||||
if review.issues:
|
||||
assessment['recommendations'].append("修复关键问题后再使用")
|
||||
if validation.get('boost', 0) < 10:
|
||||
assessment['recommendations'].append("增加交叉验证项")
|
||||
|
||||
return assessment
|
||||
|
||||
@staticmethod
|
||||
def review_all(entities: List[Dict], relations: List[Dict], text: str) -> Dict:
|
||||
"""完整审核流程"""
|
||||
print(" Step 5.5.1: 规则审查...")
|
||||
review = V15Reviewer.rule_review(entities, relations)
|
||||
print(f" 得分: {review.score}/100, 严重度: {review.severity}")
|
||||
if review.issues:
|
||||
print(f" ⚠️ 问题: {review.issues}")
|
||||
if review.warnings:
|
||||
print(f" 📝 警告: {review.warnings}")
|
||||
|
||||
print(" Step 5.5.2: 交叉验证...")
|
||||
validation = V15Reviewer.cross_validation(entities, relations, text)
|
||||
print(f" 验证项: {len(validation['validations'])}, 加分: +{validation['score_boost']}")
|
||||
|
||||
print(" Step 5.5.3: 最终评估...")
|
||||
assessment = V15Reviewer.final_assessment(review, validation)
|
||||
print(f" 最终得分: {assessment['score']}/100, 等级: {assessment['grade']}")
|
||||
print(f" 审核结果: {'✅ 通过' if assessment['passed'] else '❌ 未通过'}")
|
||||
|
||||
return assessment
|
||||
|
||||
|
||||
# 便捷函数
|
||||
def review_extraction_result(entities: List[Dict], relations: List[Dict], text: str) -> Dict:
|
||||
"""审核提取结果"""
|
||||
reviewer = V15Reviewer()
|
||||
return reviewer.review_all(entities, relations, text)
|
||||
+118
@@ -0,0 +1,118 @@
|
||||
"""
|
||||
Phase 6.5: 图表说明提取 - V13.1
|
||||
从 Markdown 提取 Figure/Table 说明并关联图片
|
||||
"""
|
||||
|
||||
import re
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class FigureCaption:
|
||||
"""图表说明"""
|
||||
figure: str # Figure 1 / Table 1
|
||||
caption: str # 完整说明文字
|
||||
page: int # 页码
|
||||
image_file: Optional[str] # 关联的图片文件名
|
||||
|
||||
|
||||
class FigureCaptionExtractor:
|
||||
"""图表说明提取器"""
|
||||
|
||||
def __init__(self):
|
||||
# V13.1: 改进正则,要求Figure/Table在行首或段落开头
|
||||
self.figure_pattern = re.compile(
|
||||
r'(?:^|\n)\s*(?:Figure|Fig\.?|Table)\s+(\d+[A-Za-z]?)\s*[:.]?\s*\n?\s*(.+?)(?=\n\n|\n(?:Figure|Fig\.?|Table)|\Z)',
|
||||
re.IGNORECASE | re.DOTALL
|
||||
)
|
||||
self.image_pattern = re.compile(r'!\[\]\(images/([^)]+)\)')
|
||||
|
||||
def extract_from_markdown(self, md_file: Path, content_list_file: Path) -> List[FigureCaption]:
|
||||
"""从 Markdown 提取图表说明 - V13.1 改进版"""
|
||||
captions = []
|
||||
|
||||
# 读取 Markdown
|
||||
md_text = md_file.read_text()
|
||||
|
||||
# 读取 content_list 获取图片页码和位置映射
|
||||
image_info = {} # image_file -> {page, position}
|
||||
if content_list_file.exists():
|
||||
with open(content_list_file) as f:
|
||||
content_list = json.load(f)
|
||||
for i, item in enumerate(content_list):
|
||||
if item.get('type') == 'image':
|
||||
# 从 content_list 获取页码
|
||||
page = item.get('page_idx', 0) + 1
|
||||
# 存储位置索引
|
||||
image_info[i] = {'page': page, 'index': i}
|
||||
|
||||
# 提取所有图片引用位置
|
||||
image_refs = list(self.image_pattern.finditer(md_text))
|
||||
image_positions = {m.group(1): m.start() for m in image_refs}
|
||||
|
||||
# 提取 Figure/Table 说明
|
||||
for match in self.figure_pattern.finditer(md_text):
|
||||
figure_num = match.group(1)
|
||||
caption_text = match.group(2).strip()
|
||||
figure_type = 'Figure' if 'fig' in match.group(0).lower() else 'Table'
|
||||
match_pos = match.start()
|
||||
|
||||
# 查找最近的图片引用(在说明前后200字符内)
|
||||
closest_image = None
|
||||
closest_dist = float('inf')
|
||||
for img_file, img_pos in image_positions.items():
|
||||
dist = abs(img_pos - match_pos)
|
||||
if dist < 200 and dist < closest_dist: # 200字符内
|
||||
closest_dist = dist
|
||||
closest_image = img_file
|
||||
|
||||
# 估算页码(从图片文件名哈希无法直接获取,使用文本位置估算)
|
||||
estimated_page = 1 + (match_pos // 3000) # 估算:每页约3000字符
|
||||
|
||||
captions.append(FigureCaption(
|
||||
figure=f"{figure_type} {figure_num}",
|
||||
caption=caption_text[:200] + "..." if len(caption_text) > 200 else caption_text,
|
||||
page=estimated_page,
|
||||
image_file=closest_image
|
||||
))
|
||||
|
||||
return captions
|
||||
|
||||
|
||||
def extract_captions(mineru_dir: Path) -> List[Dict]:
|
||||
"""便捷函数:从 MinerU 输出提取图表说明"""
|
||||
extractor = FigureCaptionExtractor()
|
||||
|
||||
md_file = mineru_dir / f'{mineru_dir.parent.name}.md'
|
||||
content_list_file = mineru_dir / f'{mineru_dir.parent.name}_content_list.json'
|
||||
|
||||
if not md_file.exists():
|
||||
return []
|
||||
|
||||
captions = extractor.extract_from_markdown(md_file, content_list_file)
|
||||
|
||||
return [{
|
||||
'figure': c.figure,
|
||||
'caption': c.caption,
|
||||
'page': c.page,
|
||||
'image_file': c.image_file
|
||||
} for c in captions]
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# 测试
|
||||
mineru_dir_value = os.environ.get("PLANTAI_TEST_MINERU_DIR", "").strip()
|
||||
if not mineru_dir_value:
|
||||
print("设置 PLANTAI_TEST_MINERU_DIR 后可运行本地图表说明测试")
|
||||
else:
|
||||
mineru_dir = Path(mineru_dir_value)
|
||||
captions = extract_captions(mineru_dir)
|
||||
print(f"提取到 {len(captions)} 个图表说明")
|
||||
for c in captions[:3]:
|
||||
print(f"\n{c['figure']}:")
|
||||
print(f" Caption: {c['caption'][:80]}...")
|
||||
print(f" Image: {c['image_file']}")
|
||||
+309
@@ -0,0 +1,309 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Phase 6: VLM 图像提取
|
||||
使用阿里云百炼 VLM 模型提取图表中的关系
|
||||
"""
|
||||
|
||||
import base64
|
||||
import requests
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
try:
|
||||
from ...project_config import env
|
||||
except ImportError:
|
||||
import os
|
||||
|
||||
def env(name: str, default: str = "") -> str:
|
||||
return os.environ.get(name, default).strip()
|
||||
|
||||
|
||||
@dataclass
|
||||
class VLMRelation:
|
||||
"""VLM 提取的关系"""
|
||||
source: str
|
||||
relation: str
|
||||
target: str
|
||||
evidence: str
|
||||
confidence: str # 'high', 'medium', 'low'
|
||||
image_source: str
|
||||
|
||||
|
||||
class VLMImageExtractor:
|
||||
"""VLM 图像提取器"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None):
|
||||
self.api_key = api_key or env('DASHSCOPE_API_KEY')
|
||||
self.api_base = "https://coding.dashscope.aliyuncs.com/v1"
|
||||
self.model = "qwen3.6-plus" # VLM 模型
|
||||
|
||||
# 启动子元件列表
|
||||
self.promoter_elements = [
|
||||
'O2-site', 'CAT-box', 'RY-element', 'circadian', 'GCN4 motif',
|
||||
'AACA motif', 'NON-box', 'MBSI', 'CGTCA-motif', 'TCA-element',
|
||||
'TGA-element', 'TGA-box', 'TGACG-motif', 'GARE-motif', 'ABRE',
|
||||
'P-box', 'TATC-box', 'AuxRR-core', 'ARE', 'LTR',
|
||||
'TC-rich repeats', 'MBS', 'GC-motif'
|
||||
]
|
||||
|
||||
def extract_from_image(self, image_path: Path) -> Dict:
|
||||
"""从单张图片提取文本"""
|
||||
if not self.api_key:
|
||||
return {'error': 'API key not set'}
|
||||
|
||||
# 读取图片
|
||||
with open(image_path, 'rb') as f:
|
||||
image_data = f.read()
|
||||
|
||||
image_base64 = base64.b64encode(image_data).decode('utf-8')
|
||||
|
||||
# 调用 VLM API
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
payload = {
|
||||
'model': self.model,
|
||||
'messages': [
|
||||
{
|
||||
'role': 'user',
|
||||
'content': [
|
||||
{
|
||||
'type': 'text',
|
||||
'text': '请识别这张图片中的所有文字内容。如果是图表,请特别关注:\n1. 表格数据(基因名、位置、ID等)\n2. 启动子元件(CAT-box, RY-element等)\n3. 表达数据(时间序列)\n4. 基因对信息(Ka/Ks值)\n\n直接输出识别到的文字,不要添加额外解释。'
|
||||
},
|
||||
{
|
||||
'type': 'image_url',
|
||||
'image_url': {
|
||||
'url': f'data:image/png;base64,{image_base64}'
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
'max_tokens': 4096
|
||||
}
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f'{self.api_base}/chat/completions',
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=120
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
text = result['choices'][0]['message']['content']
|
||||
return {
|
||||
'image': image_path.name,
|
||||
'text': text,
|
||||
'success': True
|
||||
}
|
||||
else:
|
||||
return {
|
||||
'error': f'API error: {response.status_code}',
|
||||
'success': False
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
'error': str(e),
|
||||
'success': False
|
||||
}
|
||||
|
||||
def extract_relations(self, vlm_text: str, image_name: str) -> List[VLMRelation]:
|
||||
"""从 VLM 文本中提取关系"""
|
||||
relations = []
|
||||
|
||||
# 1. 基因信息表格 -> located_on, has_gene_id
|
||||
if 'Gene Id' in vlm_text and 'Gene location' in vlm_text:
|
||||
relations.extend(self._extract_gene_info(vlm_text, image_name))
|
||||
|
||||
# 2. 启动子元件表格 -> had
|
||||
if any(elem in vlm_text for elem in self.promoter_elements[:5]):
|
||||
relations.extend(self._extract_promoter_elements(vlm_text, image_name))
|
||||
|
||||
# 3. Ka/Ks 表格 -> collinear_with
|
||||
if 'Gene pairs' in vlm_text and 'Ka/Ks' in vlm_text:
|
||||
relations.extend(self._extract_kaks_pairs(vlm_text, image_name))
|
||||
|
||||
# 4. 表达热图 -> expressed_in
|
||||
if '0h' in vlm_text and '6h' in vlm_text:
|
||||
relations.extend(self._extract_expression(vlm_text, image_name))
|
||||
|
||||
return relations
|
||||
|
||||
def _extract_gene_info(self, text: str, image_name: str) -> List[VLMRelation]:
|
||||
"""提取基因信息关系"""
|
||||
relations = []
|
||||
|
||||
for line in text.split('\n')[1:]: # 跳过表头
|
||||
if line.strip() and line.startswith('CiSPL'):
|
||||
parts = line.split()
|
||||
if len(parts) >= 3:
|
||||
gene_name = parts[0]
|
||||
gene_id = parts[1] if len(parts) > 1 else ''
|
||||
location = parts[2] if len(parts) > 2 else ''
|
||||
|
||||
# located_on
|
||||
if location and 'scaffold' in location:
|
||||
scaffold = location.split(':')[0]
|
||||
relations.append(VLMRelation(
|
||||
source=gene_name,
|
||||
relation='located_on',
|
||||
target=scaffold,
|
||||
evidence=f'Table: {line[:80]}',
|
||||
confidence='high',
|
||||
image_source=image_name
|
||||
))
|
||||
|
||||
# has_gene_id
|
||||
if gene_id and gene_id.startswith('CIL'):
|
||||
relations.append(VLMRelation(
|
||||
source=gene_name,
|
||||
relation='has_gene_id',
|
||||
target=gene_id,
|
||||
evidence=f'Table: {line[:80]}',
|
||||
confidence='high',
|
||||
image_source=image_name
|
||||
))
|
||||
|
||||
return relations
|
||||
|
||||
def _extract_promoter_elements(self, text: str, image_name: str) -> List[VLMRelation]:
|
||||
"""提取启动子元件关系"""
|
||||
relations = []
|
||||
|
||||
for line in text.split('\n'):
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
if line.startswith('CtSPL') or line.startswith('CiSPL'):
|
||||
gene_match = re.match(r'(CtSPL|CiSPL)(\d+[a-z]?)', line)
|
||||
if gene_match:
|
||||
gene_name = 'CiSPL' + gene_match.group(2)
|
||||
numbers = re.findall(r'\s+(\d+)', line)
|
||||
|
||||
for i, count in enumerate(numbers):
|
||||
if int(count) > 0 and i < len(self.promoter_elements):
|
||||
element = self.promoter_elements[i]
|
||||
relations.append(VLMRelation(
|
||||
source=gene_name,
|
||||
relation='had',
|
||||
target=element,
|
||||
evidence=f'Promoter table: {gene_name} has {count} {element}',
|
||||
confidence='high',
|
||||
image_source=image_name
|
||||
))
|
||||
|
||||
return relations
|
||||
|
||||
def _extract_kaks_pairs(self, text: str, image_name: str) -> List[VLMRelation]:
|
||||
"""提取 Ka/Ks 基因对关系"""
|
||||
relations = []
|
||||
|
||||
for line in text.split('\n')[1:]:
|
||||
if '/' in line and 'CiSPL' in line:
|
||||
match = re.match(r'(CiSPL\d+[a-z]?)/(CiSPL\d+[a-z]?)', line)
|
||||
if match:
|
||||
gene1, gene2 = match.groups()
|
||||
relations.append(VLMRelation(
|
||||
source=gene1,
|
||||
relation='collinear_with',
|
||||
target=gene2,
|
||||
evidence=f'Ka/Ks table: {line[:50]}',
|
||||
confidence='high',
|
||||
image_source=image_name
|
||||
))
|
||||
|
||||
return relations
|
||||
|
||||
def _extract_expression(self, text: str, image_name: str) -> List[VLMRelation]:
|
||||
"""提取表达关系"""
|
||||
relations = []
|
||||
genes = re.findall(r'(CiSPL\d+[a-z]?)', text)
|
||||
genes = list(set(genes))
|
||||
|
||||
for gene in genes:
|
||||
relations.append(VLMRelation(
|
||||
source=gene,
|
||||
relation='expressed_in',
|
||||
target='stress response time course',
|
||||
evidence='Expression heatmap with time points',
|
||||
confidence='medium',
|
||||
image_source=image_name
|
||||
))
|
||||
|
||||
return relations
|
||||
|
||||
def process_images(self, image_dir: Path, max_images: int = 10) -> List[VLMRelation]:
|
||||
"""处理目录中的所有图片"""
|
||||
all_relations = []
|
||||
|
||||
images = list(image_dir.glob('*.jpg')) + list(image_dir.glob('*.png'))
|
||||
images = images[:max_images]
|
||||
|
||||
print(f"Processing {len(images)} images with VLM...")
|
||||
|
||||
for i, img_path in enumerate(images):
|
||||
print(f" [{i+1}/{len(images)}] {img_path.name}")
|
||||
|
||||
# 提取文本
|
||||
result = self.extract_from_image(img_path)
|
||||
|
||||
if result.get('success'):
|
||||
# 提取关系
|
||||
relations = self.extract_relations(result['text'], img_path.name)
|
||||
all_relations.extend(relations)
|
||||
print(f" -> {len(relations)} relations")
|
||||
else:
|
||||
print(f" -> Failed: {result.get('error', 'Unknown')}")
|
||||
|
||||
return all_relations
|
||||
|
||||
|
||||
# 便捷函数
|
||||
def extract_from_mineru_images(mineru_image_dir: str, api_key: Optional[str] = None, max_images: int = 10) -> List[Dict]:
|
||||
"""从 MinerU 提取的图片中提取关系"""
|
||||
extractor = VLMImageExtractor(api_key)
|
||||
image_dir = Path(mineru_image_dir)
|
||||
|
||||
if not image_dir.exists():
|
||||
return []
|
||||
|
||||
relations = extractor.process_images(image_dir, max_images=max_images)
|
||||
|
||||
# 转换为字典列表
|
||||
return [{
|
||||
'source': r.source,
|
||||
'relation': r.relation,
|
||||
'target': r.target,
|
||||
'evidence': r.evidence,
|
||||
'confidence': r.confidence,
|
||||
'image_source': r.image_source
|
||||
} for r in relations]
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# 测试
|
||||
image_dir_value = env("PLANTAI_TEST_MINERU_IMAGE_DIR")
|
||||
if not image_dir_value:
|
||||
print("设置 PLANTAI_TEST_MINERU_IMAGE_DIR 后可运行本地 VLM 图像测试")
|
||||
else:
|
||||
image_dir = Path(image_dir_value)
|
||||
|
||||
extractor = VLMImageExtractor()
|
||||
relations = extractor.process_images(image_dir, max_images=5)
|
||||
|
||||
print(f"\nTotal relations: {len(relations)}")
|
||||
|
||||
# 统计
|
||||
from collections import Counter
|
||||
rel_types = Counter([r.relation for r in relations])
|
||||
print("\nRelation types:")
|
||||
for rel, count in rel_types.most_common():
|
||||
print(f" - {rel}: {count}")
|
||||
+637
@@ -0,0 +1,637 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Build cleaned relation-extraction blocks for V15 SRT extraction.
|
||||
|
||||
The block layer is intentionally scoped to V15 relation extraction. It does not
|
||||
change methodology extraction, which has its own Methods-only pipeline.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import json
|
||||
import re
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Callable, Iterable, List, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class RelationBlock:
|
||||
block_id: str
|
||||
section_title: str
|
||||
source_label: str
|
||||
block_type: str
|
||||
text: str
|
||||
char_count: int
|
||||
quality_flags: List[str] = field(default_factory=list)
|
||||
used_for_extraction: bool = True
|
||||
|
||||
|
||||
class RelationBlockBuilder:
|
||||
"""Normalize full-paper text into section-aware relation blocks."""
|
||||
|
||||
ABSTRACT_HEADING = re.compile(r"^\s*(abstract|summary)\b", re.I)
|
||||
REFERENCE_HEADING = re.compile(
|
||||
r"^\s*(references?|bibliography|literature cited|cited references)\s*$",
|
||||
re.I,
|
||||
)
|
||||
EXCLUDED_SECTION = re.compile(
|
||||
r"\b(references?|acknowledg(?:e)?ments?|author contributions?|funding|conflicts? of interest)\b",
|
||||
re.I,
|
||||
)
|
||||
CAPTION_START = re.compile(
|
||||
r"^\s*(fig(?:ure)?\.?\s*[s\d]+|table\s*[s\d]+|supplementary\s+(?:fig(?:ure)?|table))\b",
|
||||
re.I,
|
||||
)
|
||||
FULL_TEXT_SECTION_HEADING = re.compile(
|
||||
r"^\s*(?:#{1,6}\s*)?(?:\d+(?:\.\d+)*\.?\s*)?"
|
||||
r"(abstract|summary|introduction|background|results?|discussion|conclusions?|"
|
||||
r"materials and methods|methods?|supplementary(?: material| information)?|"
|
||||
r"references?|bibliography|literature cited)\s*[::]?\s*$",
|
||||
re.I,
|
||||
)
|
||||
INLINE_MAJOR_HEADING = re.compile(
|
||||
r"(?im)(?P<prefix>^|\n\s*|\s+#{1,6}\s*|\s+)"
|
||||
r"(?P<number>\d+(?:\.\d+)*\.?\s*)?"
|
||||
r"(?P<title>abstract|summary|introduction|background|results?|discussion|"
|
||||
r"conclusions?|materials and methods|methods?|supplementary(?: material| information)?|"
|
||||
r"references?|bibliography|literature cited)"
|
||||
r"\b(?P<punct>\s*[::]?)"
|
||||
)
|
||||
HIGH_RECALL_SECTION = re.compile(
|
||||
r"\b(abstract|summary|result|results|discussion|conclusion|background|introduction|"
|
||||
r"supplement(?:ary)?|table|figure|caption)\b",
|
||||
re.I,
|
||||
)
|
||||
URL_ONLY = re.compile(r"^\s*(?:\d{1,3}\s*)?(?:https?://|www\.)\S+\s*$", re.I)
|
||||
SUPPLEMENT_NOTICE = re.compile(
|
||||
r"\b(supplementary material|additional file|available at|data availability|supporting information)\b.*(?:https?://|doi:)",
|
||||
re.I,
|
||||
)
|
||||
METHOD_OR_RELATION_WORD = re.compile(
|
||||
r"\b("
|
||||
r"express(?:ed|ion|ing|es)?|regulated|regulates?|correlat(?:ed|ion|ing)?|associated|"
|
||||
r"interact(?:ed|ion|ing|s)?|bind(?:ing|s)?|target(?:ed|s)?|"
|
||||
r"member|subfamily|family|ortholog|homolog|paralog|cluster|synten|collinear|"
|
||||
r"duplicat|locali[sz]|involved|role|response|phenotype|domain|motif|"
|
||||
r"identified|contains?|harbors?|lacks?|C4 crop|nutrient"
|
||||
r")\b",
|
||||
re.I,
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
min_block_chars: int = 220,
|
||||
max_block_chars: int = 4800,
|
||||
include_caption_blocks: bool = False,
|
||||
include_table_caption_blocks: bool = True,
|
||||
include_figure_caption_blocks: bool = False,
|
||||
):
|
||||
self.min_block_chars = min_block_chars
|
||||
self.max_block_chars = max_block_chars
|
||||
self.include_caption_blocks = include_caption_blocks
|
||||
self.include_table_caption_blocks = include_table_caption_blocks
|
||||
self.include_figure_caption_blocks = include_figure_caption_blocks
|
||||
self.blocks: List[RelationBlock] = []
|
||||
self.removed_rows: List[dict] = []
|
||||
self.abstract_recovery_rows: List[dict] = []
|
||||
|
||||
def build(
|
||||
self,
|
||||
full_text: str,
|
||||
sections: Optional[Iterable] = None,
|
||||
section_text_getter: Optional[Callable[[str, object], str]] = None,
|
||||
pmid: str = "",
|
||||
) -> List[RelationBlock]:
|
||||
self.blocks = []
|
||||
self.removed_rows = []
|
||||
self.abstract_recovery_rows = []
|
||||
full_text = self._trim_reference_tail(full_text or "", "full_text")
|
||||
source_sections = list(sections or [])
|
||||
if not source_sections:
|
||||
for title, section_text in self._split_full_text_by_headings(full_text or ""):
|
||||
if self._is_reference_or_tail_section(title):
|
||||
self._record_removed(title, "section", "reference_or_tail_section", title)
|
||||
continue
|
||||
if self._section_kind(title) == "front_matter":
|
||||
self._record_removed(title, "section", "front_matter_metadata", section_text[:2000])
|
||||
continue
|
||||
self._build_for_section(title, section_text, pmid)
|
||||
else:
|
||||
self._recover_front_matter_abstract(full_text or "", source_sections, pmid)
|
||||
for section in source_sections:
|
||||
title = str(getattr(section, "title", "") or "section")
|
||||
if self._is_reference_or_tail_section(title):
|
||||
self._record_removed(title, "section", "reference_or_tail_section", title)
|
||||
continue
|
||||
text = ""
|
||||
if section_text_getter:
|
||||
try:
|
||||
text = section_text_getter(full_text, section)
|
||||
except Exception:
|
||||
text = ""
|
||||
if not text:
|
||||
start = int(getattr(section, "position", 0) or 0)
|
||||
end = int(getattr(section, "end", len(full_text)) or len(full_text))
|
||||
text = (full_text or "")[start:end]
|
||||
self._build_for_section(title, text, pmid)
|
||||
return list(self.blocks)
|
||||
|
||||
def save_audit_files(self, output_dir: Path | str) -> dict:
|
||||
out_dir = Path(output_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
blocks_path = out_dir / "relation_blocks.csv"
|
||||
removed_path = out_dir / "relation_removed_noise.csv"
|
||||
abstract_path = out_dir / "abstract_recovery_audit.csv"
|
||||
summary_path = out_dir / "relation_block_extraction_summary.json"
|
||||
|
||||
self._write_csv(
|
||||
blocks_path,
|
||||
[asdict(block) | {"quality_flags": "|".join(block.quality_flags)} for block in self.blocks],
|
||||
[
|
||||
"block_id",
|
||||
"section_title",
|
||||
"source_label",
|
||||
"block_type",
|
||||
"char_count",
|
||||
"quality_flags",
|
||||
"used_for_extraction",
|
||||
"text",
|
||||
],
|
||||
)
|
||||
self._write_csv(
|
||||
removed_path,
|
||||
self.removed_rows,
|
||||
["section_title", "block_type", "remove_reason", "text"],
|
||||
)
|
||||
self._write_csv(
|
||||
abstract_path,
|
||||
self.abstract_recovery_rows,
|
||||
["paper_id", "action", "reason", "section_title", "char_count", "text"],
|
||||
)
|
||||
summary = self.summary()
|
||||
summary["audit_files"] = {
|
||||
"relation_blocks": str(blocks_path),
|
||||
"relation_removed_noise": str(removed_path),
|
||||
"abstract_recovery_audit": str(abstract_path),
|
||||
"relation_block_extraction_summary": str(summary_path),
|
||||
}
|
||||
summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
return summary
|
||||
|
||||
def summary(self) -> dict:
|
||||
by_section: dict[str, int] = {}
|
||||
by_type: dict[str, int] = {}
|
||||
used = 0
|
||||
for block in self.blocks:
|
||||
by_section[block.section_title] = by_section.get(block.section_title, 0) + 1
|
||||
by_type[block.block_type] = by_type.get(block.block_type, 0) + 1
|
||||
if block.used_for_extraction:
|
||||
used += 1
|
||||
return {
|
||||
"relation_block_builder_enabled": True,
|
||||
"block_count": len(self.blocks),
|
||||
"used_block_count": used,
|
||||
"removed_noise_count": len(self.removed_rows),
|
||||
"front_matter_abstract_recovered": sum(
|
||||
1 for row in self.abstract_recovery_rows if row.get("action") == "selected"
|
||||
),
|
||||
"blocks_by_section": by_section,
|
||||
"blocks_by_type": by_type,
|
||||
}
|
||||
|
||||
def _recover_front_matter_abstract(self, full_text: str, source_sections: Iterable, pmid: str) -> None:
|
||||
"""Recover publisher-style abstracts that appear before Introduction without an Abstract heading."""
|
||||
if not full_text:
|
||||
return
|
||||
if any(self._section_kind(str(getattr(section, "title", "") or "")) == "abstract" for section in source_sections):
|
||||
self.abstract_recovery_rows.append(
|
||||
{
|
||||
"paper_id": pmid,
|
||||
"action": "skipped",
|
||||
"reason": "explicit_abstract_section_exists",
|
||||
"section_title": "abstract",
|
||||
"char_count": 0,
|
||||
"text": "",
|
||||
}
|
||||
)
|
||||
return
|
||||
|
||||
positions: List[int] = []
|
||||
for section in source_sections:
|
||||
try:
|
||||
pos = int(getattr(section, "position", 0) or 0)
|
||||
except Exception:
|
||||
pos = 0
|
||||
title = str(getattr(section, "title", "") or "")
|
||||
if pos > 0 and not self._is_reference_or_tail_section(title):
|
||||
positions.append(pos)
|
||||
|
||||
first_section_pos = min(positions) if positions else 0
|
||||
if first_section_pos <= 0:
|
||||
match = re.search(r"(?im)^\s*(?:\d+(?:\.\d+)*\.?\s*)?(introduction|background)\s*$", full_text)
|
||||
first_section_pos = match.start() if match else 0
|
||||
if first_section_pos <= 0:
|
||||
self._record_abstract_recovery(pmid, "skipped", "no_front_matter_boundary", "", "")
|
||||
return
|
||||
|
||||
front_text = full_text[:first_section_pos]
|
||||
candidate = self._select_front_matter_abstract_candidate(front_text)
|
||||
if not candidate:
|
||||
self._record_abstract_recovery(pmid, "skipped", "no_abstract_like_candidate", "", front_text[:2000])
|
||||
return
|
||||
|
||||
flags = self._quality_flags(candidate) + ["front_matter_recovered_abstract"]
|
||||
self._add_block("abstract", "abstract", "abstract", candidate, flags, True, pmid)
|
||||
self._record_abstract_recovery(pmid, "selected", "front_matter_before_first_section", "abstract", candidate)
|
||||
|
||||
def _select_front_matter_abstract_candidate(self, front_text: str) -> str:
|
||||
cleaned_lines: List[str] = []
|
||||
for raw_line in str(front_text or "").splitlines():
|
||||
line = self._normalize_line(raw_line)
|
||||
if not line:
|
||||
cleaned_lines.append("")
|
||||
continue
|
||||
if self._front_matter_line_reason(line):
|
||||
continue
|
||||
cleaned_lines.append(line)
|
||||
|
||||
cleaned = "\n".join(cleaned_lines)
|
||||
cleaned = re.sub(r"(?is)\bkeywords?\s*[::].*$", "", cleaned)
|
||||
cleaned = re.sub(r"(?is)^.*?\babstract\b\s*[::]?", "", cleaned, count=1) if re.search(r"(?i)\babstract\b", cleaned[:800]) else cleaned
|
||||
paragraphs = [self._compact(p) for p in re.split(r"\n\s*\n+", cleaned) if self._compact(p)]
|
||||
if not paragraphs:
|
||||
paragraphs = [self._compact(cleaned)]
|
||||
|
||||
abstract_signal = re.compile(
|
||||
r"\b(in this study|we identified|we analyzed|we analysed|we performed|we found|"
|
||||
r"results? showed|revealed|suggest(?:ed|s)?|genome-wide|gene family|"
|
||||
r"transcriptome|rna-seq|expression|phylogen|identified|characterized|"
|
||||
r"our current study|this study)\b",
|
||||
re.I,
|
||||
)
|
||||
candidates: List[str] = []
|
||||
for paragraph in paragraphs:
|
||||
if len(paragraph) < 220:
|
||||
continue
|
||||
if self._looks_like_front_matter_noise(paragraph):
|
||||
continue
|
||||
if abstract_signal.search(paragraph) or len(paragraph) >= 450:
|
||||
candidates.append(paragraph)
|
||||
if candidates:
|
||||
return self._compact(" ".join(candidates[:2]))[:3500]
|
||||
return ""
|
||||
|
||||
def _front_matter_line_reason(self, line: str) -> str:
|
||||
if line.startswith(" or line.startswith("# "):
|
||||
return "title_or_image"
|
||||
if re.search(r"\b(received|accepted|published|citation|cite this article|correspondence|edited by|reviewed by)\b", line, re.I):
|
||||
return "publisher_metadata"
|
||||
if re.search(r"\b(author contributions?|copyright|creative commons|frontiers in|keywords?)\b", line, re.I):
|
||||
return "publisher_metadata"
|
||||
if re.search(r"^[A-Za-z .,'-]+(?:\d+\s*,?\s*){2,}$", line) and len(line) < 220:
|
||||
return "author_affiliation_line"
|
||||
if re.fullmatch(r"(?:\d+\s*)+", line):
|
||||
return "numeric_marker"
|
||||
return ""
|
||||
|
||||
def _looks_like_front_matter_noise(self, text: str) -> bool:
|
||||
value = self._compact(text)
|
||||
if len(value) < 220:
|
||||
return True
|
||||
meta_hits = len(re.findall(
|
||||
r"\b(correspondence|received|accepted|published|citation|copyright|"
|
||||
r"affiliation|department|university|email|doi)\b",
|
||||
value,
|
||||
re.I,
|
||||
))
|
||||
if meta_hits >= 2 and not re.search(r"\b(in this study|results?|we identified|we analyzed|revealed)\b", value, re.I):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _record_abstract_recovery(self, pmid: str, action: str, reason: str, section_title: str, text: str) -> None:
|
||||
self.abstract_recovery_rows.append(
|
||||
{
|
||||
"paper_id": pmid,
|
||||
"action": action,
|
||||
"reason": reason,
|
||||
"section_title": section_title,
|
||||
"char_count": len(self._compact(text)),
|
||||
"text": self._compact(text)[:3000],
|
||||
}
|
||||
)
|
||||
|
||||
def _build_for_section(self, section_title: str, text: str, pmid: str) -> None:
|
||||
section_kind = self._section_kind(section_title)
|
||||
text = self._trim_reference_tail(str(text or ""), section_title)
|
||||
paragraphs: List[str] = []
|
||||
caption_rows: List[str] = []
|
||||
current: List[str] = []
|
||||
in_reference_tail = False
|
||||
|
||||
for raw_line in str(text or "").splitlines():
|
||||
line = self._normalize_line(raw_line)
|
||||
if not line:
|
||||
if current:
|
||||
paragraphs.append(" ".join(current).strip())
|
||||
current = []
|
||||
continue
|
||||
if self.REFERENCE_HEADING.match(line):
|
||||
in_reference_tail = True
|
||||
self._record_removed(section_title, "tail", "reference_heading", line)
|
||||
if current:
|
||||
paragraphs.append(" ".join(current).strip())
|
||||
current = []
|
||||
continue
|
||||
if in_reference_tail:
|
||||
self._record_removed(section_title, "tail", "reference_tail", line)
|
||||
continue
|
||||
reason = self._remove_reason(line)
|
||||
if reason:
|
||||
self._record_removed(section_title, "line", reason, line)
|
||||
continue
|
||||
if self.CAPTION_START.match(line):
|
||||
caption_rows.append(line)
|
||||
continue
|
||||
current.append(line)
|
||||
if current:
|
||||
paragraphs.append(" ".join(current).strip())
|
||||
|
||||
for block_text in self._pack_paragraphs(paragraphs):
|
||||
if not block_text:
|
||||
continue
|
||||
flags = self._quality_flags(block_text)
|
||||
if (
|
||||
"low_relation_signal" in flags
|
||||
and len(block_text) < self.min_block_chars
|
||||
and section_kind not in {"abstract", "supplementary", "high_recall"}
|
||||
):
|
||||
self._record_removed(section_title, "main_text", "low_relation_signal_short_block", block_text)
|
||||
continue
|
||||
source_label = "abstract" if section_kind == "abstract" else ("supplementary" if section_kind == "supplementary" else "text")
|
||||
self._add_block(section_title, source_label, "main_text", block_text, flags, True, pmid)
|
||||
|
||||
for caption in caption_rows:
|
||||
flags = self._quality_flags(caption) + ["caption"]
|
||||
is_figure_caption = caption.lower().startswith(("fig", "figure", "supplementary fig"))
|
||||
block_type = "figure_caption" if is_figure_caption else "table_caption"
|
||||
if self.include_caption_blocks:
|
||||
used_for_extraction = True
|
||||
elif is_figure_caption:
|
||||
used_for_extraction = self.include_figure_caption_blocks
|
||||
elif section_kind == "supplementary":
|
||||
used_for_extraction = True
|
||||
else:
|
||||
used_for_extraction = self.include_table_caption_blocks
|
||||
if not used_for_extraction:
|
||||
reason = "figure_caption_block_disabled" if is_figure_caption else "table_caption_block_disabled"
|
||||
self._record_removed(section_title, block_type, reason, caption)
|
||||
continue
|
||||
self._add_block(
|
||||
section_title,
|
||||
"caption",
|
||||
block_type,
|
||||
caption,
|
||||
flags,
|
||||
used_for_extraction,
|
||||
pmid,
|
||||
)
|
||||
|
||||
def _trim_reference_tail(self, text: str, section_title: str) -> str:
|
||||
"""Drop bibliography tails before relation-block construction.
|
||||
|
||||
References are useful metadata but extremely noisy for SRT extraction:
|
||||
they contain author names, years, gene symbols in cited titles, and
|
||||
unrelated biology. We keep this scoped to relation blocks so the
|
||||
raw MinerU files remain traceable.
|
||||
"""
|
||||
value = str(text or "")
|
||||
if not value:
|
||||
return value
|
||||
patterns = [
|
||||
r"(?im)^\s*(?:#+\s*)?(?:references|literature cited|bibliography)\s*$",
|
||||
r"(?im)(?:^|\n)\s*(?:#+\s*)?(?:references|literature cited|bibliography)\s*(?=\n|$)",
|
||||
r"(?im)\s+#{1,6}\s*(?:references|literature cited|bibliography)\b",
|
||||
]
|
||||
best = None
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, value)
|
||||
if match:
|
||||
best = match.start() if best is None else min(best, match.start())
|
||||
if best is None:
|
||||
return value
|
||||
removed = value[best:]
|
||||
self._record_removed(section_title, "tail", "reference_tail_truncated", removed[:2000])
|
||||
return value[:best]
|
||||
|
||||
def _pack_paragraphs(self, paragraphs: List[str]) -> List[str]:
|
||||
packed: List[str] = []
|
||||
current: List[str] = []
|
||||
current_len = 0
|
||||
for paragraph in paragraphs:
|
||||
paragraph = self._compact(paragraph)
|
||||
if not paragraph:
|
||||
continue
|
||||
if len(paragraph) > self.max_block_chars:
|
||||
if current:
|
||||
packed.append(" ".join(current).strip())
|
||||
current = []
|
||||
current_len = 0
|
||||
packed.extend(self._split_long_paragraph(paragraph))
|
||||
continue
|
||||
if current and current_len + len(paragraph) > self.max_block_chars:
|
||||
packed.append(" ".join(current).strip())
|
||||
current = []
|
||||
current_len = 0
|
||||
current.append(paragraph)
|
||||
current_len += len(paragraph) + 1
|
||||
if current:
|
||||
packed.append(" ".join(current).strip())
|
||||
return packed
|
||||
|
||||
def _split_long_paragraph(self, paragraph: str) -> List[str]:
|
||||
sentences = re.split(r"(?<=[.!?。!?])\s+", paragraph)
|
||||
parts: List[str] = []
|
||||
current: List[str] = []
|
||||
current_len = 0
|
||||
for sentence in sentences:
|
||||
sentence = sentence.strip()
|
||||
if not sentence:
|
||||
continue
|
||||
if current and current_len + len(sentence) > self.max_block_chars:
|
||||
parts.append(" ".join(current).strip())
|
||||
current = []
|
||||
current_len = 0
|
||||
current.append(sentence)
|
||||
current_len += len(sentence) + 1
|
||||
if current:
|
||||
parts.append(" ".join(current).strip())
|
||||
return parts
|
||||
|
||||
def _split_full_text_by_headings(self, full_text: str) -> List[tuple[str, str]]:
|
||||
"""Fallback section splitter when MinerU did not return section objects.
|
||||
|
||||
MinerU/PyMuPDF often flattens markdown-like headings into inline text,
|
||||
for example ``Abstract: ...`` or ``## 2. Results ## 2.1 ...``. A
|
||||
line-only heading splitter misses those cases and collapses the whole
|
||||
paper into ``full_text``. This scanner recognizes the major paper
|
||||
headings in-place and keeps the same-line content after the heading as
|
||||
the section body.
|
||||
"""
|
||||
text = str(full_text or "")
|
||||
if not text.strip():
|
||||
return []
|
||||
|
||||
matches = [m for m in self.INLINE_MAJOR_HEADING.finditer(text) if self._valid_heading_match(text, m)]
|
||||
if not matches:
|
||||
return [("full_text", full_text or "")]
|
||||
|
||||
sections: List[tuple[str, str]] = []
|
||||
first = matches[0]
|
||||
front = text[: first.start()].strip()
|
||||
if front:
|
||||
sections.append(("front_matter", front))
|
||||
|
||||
for index, match in enumerate(matches):
|
||||
title = self._canonical_heading_title(match.group("title") or "")
|
||||
body_start = match.end()
|
||||
body_end = matches[index + 1].start() if index + 1 < len(matches) else len(text)
|
||||
body = text[body_start:body_end].strip()
|
||||
if body:
|
||||
sections.append((title, body))
|
||||
|
||||
return [(title, body) for title, body in sections if body]
|
||||
|
||||
def _valid_heading_match(self, text: str, match: re.Match) -> bool:
|
||||
prefix = match.group("prefix") or ""
|
||||
title = (match.group("title") or "").lower()
|
||||
number = match.group("number") or ""
|
||||
punct = match.group("punct") or ""
|
||||
start = match.start()
|
||||
after = text[match.end() : match.end() + 40]
|
||||
if not ("#" in prefix or number or punct.strip()):
|
||||
if re.match(r"\s+[a-z]", after) and not re.match(r"\s*(?:\n|$)", after):
|
||||
return False
|
||||
if title in {"results", "result"} and re.search(r"\b(results?|showed|revealed)\s+(?:that|a|an|the)\b", text[max(0, start - 30): match.end() + 30], re.I):
|
||||
return bool("#" in prefix or prefix.startswith("\n") or start == 0)
|
||||
if "#" in prefix or prefix.startswith("\n") or start == 0:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _canonical_heading_title(self, title: str) -> str:
|
||||
clean = self._compact(title).lower()
|
||||
if clean == "result":
|
||||
clean = "results"
|
||||
if clean in {"method", "methods"}:
|
||||
return "Methods"
|
||||
if clean == "materials and methods":
|
||||
return "Materials and Methods"
|
||||
if clean == "literature cited":
|
||||
return "References"
|
||||
return clean[:1].upper() + clean[1:]
|
||||
|
||||
def _add_block(
|
||||
self,
|
||||
section_title: str,
|
||||
source_label: str,
|
||||
block_type: str,
|
||||
text: str,
|
||||
quality_flags: List[str],
|
||||
used_for_extraction: bool,
|
||||
pmid: str,
|
||||
) -> None:
|
||||
index = len(self.blocks) + 1
|
||||
section_slug = re.sub(r"[^A-Za-z0-9]+", "_", section_title.lower()).strip("_") or "section"
|
||||
block_id = f"{pmid or 'paper'}_{section_slug}_{index:04d}"
|
||||
clean_text = self._compact(text)
|
||||
self.blocks.append(
|
||||
RelationBlock(
|
||||
block_id=block_id,
|
||||
section_title=section_title,
|
||||
source_label=source_label,
|
||||
block_type=block_type,
|
||||
text=clean_text,
|
||||
char_count=len(clean_text),
|
||||
quality_flags=quality_flags,
|
||||
used_for_extraction=used_for_extraction,
|
||||
)
|
||||
)
|
||||
|
||||
def _remove_reason(self, line: str) -> str:
|
||||
if self.URL_ONLY.match(line):
|
||||
return "url_only"
|
||||
if self.SUPPLEMENT_NOTICE.search(line) and len(line) < 260:
|
||||
return "supplementary_availability_notice"
|
||||
if self._looks_like_visual_noise(line):
|
||||
return "visual_ocr_noise"
|
||||
return ""
|
||||
|
||||
def _quality_flags(self, text: str) -> List[str]:
|
||||
flags: List[str] = []
|
||||
if not self.METHOD_OR_RELATION_WORD.search(text):
|
||||
flags.append("low_relation_signal")
|
||||
if self.CAPTION_START.match(text):
|
||||
flags.append("caption_like")
|
||||
if re.search(r"https?://", text) and len(text) < 180:
|
||||
flags.append("short_url_text")
|
||||
return flags
|
||||
|
||||
def _looks_like_visual_noise(self, line: str) -> bool:
|
||||
if len(line) < 4:
|
||||
return True
|
||||
tokens = re.findall(r"[A-Za-z0-9_.-]+", line)
|
||||
if not tokens:
|
||||
return True
|
||||
short_tokens = sum(1 for token in tokens if len(token) <= 2)
|
||||
alpha = len(re.findall(r"[A-Za-z]", line))
|
||||
digits = len(re.findall(r"\d", line))
|
||||
if len(tokens) >= 8 and short_tokens / max(1, len(tokens)) > 0.65:
|
||||
return True
|
||||
if digits > alpha and not self.METHOD_OR_RELATION_WORD.search(line):
|
||||
return True
|
||||
if re.fullmatch(r"(?:[A-Z]\s*){3,}", line):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _is_reference_or_tail_section(self, title: str) -> bool:
|
||||
clean = self._compact(title)
|
||||
return bool(self.REFERENCE_HEADING.match(clean) or self.EXCLUDED_SECTION.search(clean))
|
||||
|
||||
def _section_kind(self, title: str) -> str:
|
||||
clean = self._compact(title)
|
||||
if re.search(r"\b(front matter|front_matter|article metadata)\b", clean, re.I):
|
||||
return "front_matter"
|
||||
if self.ABSTRACT_HEADING.match(clean):
|
||||
return "abstract"
|
||||
if re.search(r"\b(supplement(?:ary)?|additional file|supporting information)\b", clean, re.I):
|
||||
return "supplementary"
|
||||
if self.HIGH_RECALL_SECTION.search(clean):
|
||||
return "high_recall"
|
||||
return "text"
|
||||
|
||||
def _normalize_line(self, line: str) -> str:
|
||||
line = str(line or "").replace("\xa0", " ")
|
||||
line = re.sub(r"(\w)-\s+(\w)", r"\1\2", line)
|
||||
return self._compact(line)
|
||||
|
||||
@staticmethod
|
||||
def _compact(text: str) -> str:
|
||||
return re.sub(r"\s+", " ", str(text or "")).strip()
|
||||
|
||||
def _record_removed(self, section_title: str, block_type: str, reason: str, text: str) -> None:
|
||||
self.removed_rows.append(
|
||||
{
|
||||
"section_title": section_title,
|
||||
"block_type": block_type,
|
||||
"remove_reason": reason,
|
||||
"text": self._compact(text)[:2000],
|
||||
}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _write_csv(path: Path, rows: List[dict], headers: List[str]) -> None:
|
||||
with open(path, "w", newline="", encoding="utf-8") as handle:
|
||||
writer = csv.DictWriter(handle, fieldnames=headers, extrasaction="ignore")
|
||||
writer.writeheader()
|
||||
for row in rows:
|
||||
writer.writerow(row)
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Compatibility alias for V16 relation-unit building."""
|
||||
from .chunk_builder_v15_2_5 import ChunkBuilderV15_2_5
|
||||
RelationUnitBuilderV16_1_1 = ChunkBuilderV15_2_5
|
||||
@@ -0,0 +1,271 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Optional
|
||||
|
||||
from .chunk_builder import Chunk
|
||||
|
||||
|
||||
SENTENCE_PLACEHOLDERS = {
|
||||
"et al.": "et al<PRD>",
|
||||
"Fig.": "Fig<PRD>",
|
||||
"Figs.": "Figs<PRD>",
|
||||
"e.g.": "e<PRD>g<PRD>",
|
||||
"i.e.": "i<PRD>e<PRD>",
|
||||
"sp.": "sp<PRD>",
|
||||
"spp.": "spp<PRD>",
|
||||
}
|
||||
|
||||
|
||||
class RelationUnitBuilderV16_1_2:
|
||||
def __init__(
|
||||
self,
|
||||
max_chars: int = 5200,
|
||||
overlap_sentences: int = 1,
|
||||
output_dir: Optional[Path] = None,
|
||||
):
|
||||
self.max_chars = max_chars
|
||||
self.overlap_sentences = max(0, overlap_sentences)
|
||||
self.output_dir = Path(output_dir) if output_dir else None
|
||||
self.unit_rows: List[Dict[str, Any]] = []
|
||||
|
||||
def build(self, blocks: List[Dict[str, Any]], paper_id: str) -> List[Chunk]:
|
||||
units: List[Dict[str, Any]] = []
|
||||
section_type = "full_text"
|
||||
section_title = "full_text"
|
||||
buffer: List[str] = []
|
||||
block_ids: List[str] = []
|
||||
block_types: List[str] = []
|
||||
subheading_context = ""
|
||||
unit_index = 0
|
||||
|
||||
def emit(
|
||||
source_label: str = "text",
|
||||
forced_text: str = "",
|
||||
forced_ids: Optional[List[str]] = None,
|
||||
forced_types: Optional[List[str]] = None,
|
||||
forced_section: str = "",
|
||||
) -> None:
|
||||
nonlocal buffer, block_ids, block_types, unit_index
|
||||
text = forced_text or " ".join(buffer).strip()
|
||||
if not text:
|
||||
if not forced_text:
|
||||
buffer, block_ids, block_types = [], [], []
|
||||
return
|
||||
unit_index += 1
|
||||
sec = forced_section or section_type
|
||||
ids = forced_ids if forced_ids is not None else list(dict.fromkeys(block_ids))
|
||||
types = (
|
||||
forced_types
|
||||
if forced_types is not None
|
||||
else sorted(set(block_types)) or ["paragraph"]
|
||||
)
|
||||
units.append(
|
||||
{
|
||||
"relation_unit_id": f"{paper_id}_{sec}_{unit_index:04d}",
|
||||
"paper_id": paper_id,
|
||||
"section_type": sec,
|
||||
"section_title": section_title,
|
||||
"source_label": source_label,
|
||||
"block_types": types,
|
||||
"block_ids": ids,
|
||||
"char_count": len(text),
|
||||
"text": text,
|
||||
}
|
||||
)
|
||||
if not forced_text:
|
||||
overlap = buffer[-self.overlap_sentences :] if self.overlap_sentences else []
|
||||
buffer = list(overlap)
|
||||
block_ids = block_ids[-self.overlap_sentences :] if overlap else []
|
||||
block_types = block_types[-self.overlap_sentences :] if overlap else []
|
||||
|
||||
for block in blocks:
|
||||
if block.get("is_reference"):
|
||||
break
|
||||
if block.get("is_major_heading"):
|
||||
emit()
|
||||
section_type = block.get("section_type") or "full_text"
|
||||
section_title = block.get("clean_text") or section_type
|
||||
subheading_context = ""
|
||||
buffer, block_ids, block_types = [], [], []
|
||||
continue
|
||||
if section_type == "administrative":
|
||||
continue
|
||||
if block.get("is_subheading"):
|
||||
subheading_context = block.get("clean_text") or ""
|
||||
continue
|
||||
if not block.get("keep_for_llm"):
|
||||
continue
|
||||
block_type = block.get("v16_block_type") or "paragraph"
|
||||
text = block.get("clean_text") or ""
|
||||
if block_type == "figure_caption":
|
||||
emit()
|
||||
for piece in self._split_long_text(text, self.max_chars):
|
||||
emit(
|
||||
source_label="figure_caption",
|
||||
forced_text=piece,
|
||||
forced_ids=[block["block_id"]],
|
||||
forced_types=["figure_caption"],
|
||||
forced_section="figure_caption",
|
||||
)
|
||||
buffer, block_ids, block_types = [], [], []
|
||||
continue
|
||||
if block_type == "table_block":
|
||||
emit()
|
||||
for piece in self._split_table(text):
|
||||
emit(
|
||||
source_label="table",
|
||||
forced_text=piece,
|
||||
forced_ids=[block["block_id"]],
|
||||
forced_types=["table_block"],
|
||||
forced_section="table",
|
||||
)
|
||||
buffer, block_ids, block_types = [], [], []
|
||||
continue
|
||||
sentences = self._split_sentences(text) or [text]
|
||||
for sentence in sentences:
|
||||
if subheading_context and not buffer:
|
||||
sentence = f"Subheading context: {subheading_context}. {sentence}"
|
||||
prospective = sum(len(value) + 1 for value in buffer) + len(sentence)
|
||||
if buffer and prospective > self.max_chars:
|
||||
emit()
|
||||
buffer.append(sentence)
|
||||
block_ids.append(block["block_id"])
|
||||
block_types.append(block_type)
|
||||
emit()
|
||||
self.unit_rows = units
|
||||
self._write_outputs(blocks, units, paper_id)
|
||||
chunks: List[Chunk] = []
|
||||
for unit in units:
|
||||
chunk = Chunk(
|
||||
chunk_id=unit["relation_unit_id"],
|
||||
section_title=unit["section_type"],
|
||||
source_label=unit["source_label"],
|
||||
text=unit["text"],
|
||||
char_count=unit["char_count"],
|
||||
endpoint_candidates={},
|
||||
)
|
||||
chunk.relation_unit_id = unit["relation_unit_id"]
|
||||
chunk.block_types = unit["block_types"]
|
||||
chunk.source_block_ids = unit["block_ids"]
|
||||
chunk.section_original_title = unit["section_title"]
|
||||
chunks.append(chunk)
|
||||
return chunks
|
||||
|
||||
def _split_sentences(self, text: str) -> List[str]:
|
||||
value = re.sub(r"\s+", " ", text or "").strip()
|
||||
for raw, replacement in SENTENCE_PLACEHOLDERS.items():
|
||||
value = value.replace(raw, replacement)
|
||||
parts = [
|
||||
part.replace("<PRD>", ".").strip()
|
||||
for part in re.split(r"(?<=[.!?])\s+(?=[A-Z0-9(\[])", value)
|
||||
if part.strip()
|
||||
]
|
||||
return parts
|
||||
|
||||
def _split_long_text(self, text: str, limit: int) -> List[str]:
|
||||
sentences = self._split_sentences(text)
|
||||
if not sentences:
|
||||
return [text[i : i + limit] for i in range(0, len(text), limit)]
|
||||
pieces: List[str] = []
|
||||
current: List[str] = []
|
||||
for sentence in sentences:
|
||||
if current and sum(len(v) + 1 for v in current) + len(sentence) > limit:
|
||||
pieces.append(" ".join(current))
|
||||
current = []
|
||||
current.append(sentence)
|
||||
if current:
|
||||
pieces.append(" ".join(current))
|
||||
return pieces
|
||||
|
||||
def _split_table(self, text: str) -> List[str]:
|
||||
lines = [line.strip() for line in (text or "").splitlines() if line.strip()]
|
||||
if not lines:
|
||||
return []
|
||||
prefix = [line for line in lines if line.startswith(("Table caption:", "Table header:"))]
|
||||
rows = [line for line in lines if line not in prefix]
|
||||
pieces: List[str] = []
|
||||
current = list(prefix)
|
||||
for row in rows:
|
||||
if current and sum(len(v) + 1 for v in current) + len(row) > self.max_chars:
|
||||
pieces.append("\n".join(current))
|
||||
current = list(prefix)
|
||||
current.append(row)
|
||||
if current:
|
||||
pieces.append("\n".join(current))
|
||||
return pieces
|
||||
|
||||
def _write_outputs(
|
||||
self,
|
||||
blocks: List[Dict[str, Any]],
|
||||
units: List[Dict[str, Any]],
|
||||
paper_id: str,
|
||||
) -> None:
|
||||
if not self.output_dir:
|
||||
return
|
||||
self.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
with (self.output_dir / "04_relation_units.jsonl").open(
|
||||
"w", encoding="utf-8"
|
||||
) as handle:
|
||||
for unit in units:
|
||||
handle.write(json.dumps(unit, ensure_ascii=False) + "\n")
|
||||
eligible = {
|
||||
block["block_id"]: int(block.get("char_count") or 0)
|
||||
for block in blocks
|
||||
if block.get("keep_for_llm")
|
||||
and not block.get("is_reference")
|
||||
and not block.get("is_image_placeholder")
|
||||
}
|
||||
covered_ids = {
|
||||
block_id for unit in units for block_id in unit.get("block_ids") or []
|
||||
}
|
||||
eligible_chars = sum(eligible.values())
|
||||
covered_chars = sum(size for block_id, size in eligible.items() if block_id in covered_ids)
|
||||
with (self.output_dir / "05_relation_units_coverage.csv").open(
|
||||
"w", newline="", encoding="utf-8"
|
||||
) as handle:
|
||||
fields = [
|
||||
"paper_id",
|
||||
"covered_chars",
|
||||
"eligible_chars",
|
||||
"coverage_ratio",
|
||||
"unit_count",
|
||||
"uncovered_block_count",
|
||||
]
|
||||
writer = csv.DictWriter(handle, fieldnames=fields)
|
||||
writer.writeheader()
|
||||
writer.writerow(
|
||||
{
|
||||
"paper_id": paper_id,
|
||||
"covered_chars": covered_chars,
|
||||
"eligible_chars": eligible_chars,
|
||||
"coverage_ratio": round(covered_chars / eligible_chars, 6)
|
||||
if eligible_chars
|
||||
else 0,
|
||||
"unit_count": len(units),
|
||||
"uncovered_block_count": len(set(eligible) - covered_ids),
|
||||
}
|
||||
)
|
||||
summary_path = self.output_dir / "mineru_block_summary.json"
|
||||
try:
|
||||
summary = json.loads(summary_path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
summary = {"paper_id": paper_id}
|
||||
summary.update(
|
||||
{
|
||||
"relation_unit_count": len(units),
|
||||
"figure_caption_unit_count": sum(
|
||||
unit["source_label"] == "figure_caption" for unit in units
|
||||
),
|
||||
"table_unit_count": sum(unit["source_label"] == "table" for unit in units),
|
||||
"coverage_ratio": round(covered_chars / eligible_chars, 6)
|
||||
if eligible_chars
|
||||
else 0,
|
||||
}
|
||||
)
|
||||
summary_path.write_text(
|
||||
json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
|
||||
)
|
||||
@@ -0,0 +1,46 @@
|
||||
from __future__ import annotations
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional
|
||||
MAJOR=[("abstract",re.compile(r"^(?:abstract|summary)$",re.I)),("introduction",re.compile(r"^(?:introduction|background)$",re.I)),("methods",re.compile(r"^(?:materials? and methods?|methods?|methodology|experimental procedures?)$",re.I)),("results",re.compile(r"^(?:results?|findings)$",re.I)),("discussion",re.compile(r"^(?:discussion)$",re.I)),("conclusion",re.compile(r"^(?:conclusions?|summary and conclusions?)$",re.I)),("case_study",re.compile(r"^(?:case study|case studies)$",re.I)),("review_content",re.compile(r"^(?:review|literature review)$",re.I)),("supplementary",re.compile(r"^(?:supplementary|supporting information)$",re.I))]
|
||||
SUB=re.compile(r"\b(gene expression|stress response|metabolite analysis|phenotype analysis|physiology|photosynthesis|hormone treatment|pathogen infection|development|yield traits|genetic diversity|transcriptome analysis|metabolome analysis|soil|rhizosphere|microbiome|growth performance|feed conversion|feed utilization|survival|mortality|disease resistance|immune response|innate immunity|water quality|water temperature|salinity|dissolved oxygen|ammonia|nitrite|stocking density|reproduction|spawning|larvae|larval|culture|aquaculture|farming|diet|feed|nutrition|hematology|serum biochemistry|antioxidant|histolog|gonad|hepatopancreas)\b",re.I)
|
||||
REF=re.compile(r"^(?:references?|bibliography|literature cited)$",re.I)
|
||||
@dataclass
|
||||
class SectionSpan:
|
||||
section_type:str; title:str; start_order:int; end_order:int; char_count:int=0
|
||||
|
||||
def normalize_heading(text:str)->str:
|
||||
text=re.sub(r"^#+\s*","",str(text or "")).strip(); text=re.sub(r"\s+"," ",text)
|
||||
letters=re.sub(r"[^A-Za-z]","",text)
|
||||
if len(letters)>=4 and len(text.split())>=len(letters)-1:
|
||||
spaced=" ".join(list(letters))
|
||||
if text.upper().replace(" ","")==letters.upper() or text.upper()==spaced.upper(): return letters.lower()
|
||||
return text.rstrip(":")
|
||||
|
||||
def classify_heading(text:str)->Optional[str]:
|
||||
norm=normalize_heading(text)
|
||||
if REF.match(norm): return "references"
|
||||
for name,pat in MAJOR:
|
||||
if pat.match(norm): return name
|
||||
if SUB.search(norm): return "subheading_inside_parent"
|
||||
return None
|
||||
|
||||
def is_major_heading(text:str)->bool:
|
||||
s=classify_heading(text); return s not in {None,"references","subheading_inside_parent"}
|
||||
|
||||
def is_subheading(text:str)->bool:
|
||||
return classify_heading(text)=="subheading_inside_parent"
|
||||
|
||||
def build_section_spans(blocks:List[dict])->List[SectionSpan]:
|
||||
spans=[]; cur="full_text"; title="full_text"; start=0; chars=0
|
||||
for b in blocks:
|
||||
order=int(b.get("order",0))
|
||||
if b.get("is_reference"): break
|
||||
if b.get("is_major_heading"):
|
||||
if order>start or chars: spans.append(SectionSpan(cur,title,start,max(start,order-1),chars))
|
||||
cur=b.get("section_type") or classify_heading(b.get("clean_text","")) or "full_text"; title=b.get("clean_text") or cur; start=order; chars=0
|
||||
elif b.get("keep_for_llm"):
|
||||
chars+=len(b.get("clean_text") or "")
|
||||
if chars or not spans:
|
||||
last=blocks[-1].get("order",0) if blocks else 0; spans.append(SectionSpan(cur,title,start,last,chars))
|
||||
return spans
|
||||
@@ -0,0 +1,204 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Optional, Tuple
|
||||
|
||||
|
||||
NUMBER_PREFIX = r"(?:\d+(?:\.\d+)*[\s.:-]*)?"
|
||||
MAJOR_PATTERNS = [
|
||||
("abstract", re.compile(rf"^{NUMBER_PREFIX}(?:abstract|summary|overview)$", re.I)),
|
||||
(
|
||||
"introduction",
|
||||
re.compile(rf"^{NUMBER_PREFIX}(?:introduction|background|context)$", re.I),
|
||||
),
|
||||
(
|
||||
"methods",
|
||||
re.compile(
|
||||
rf"^{NUMBER_PREFIX}(?:materials?\s*(?:and|&)\s*methods?|methods?|methodology|experimental\s+(?:procedures?|design))$",
|
||||
re.I,
|
||||
),
|
||||
),
|
||||
(
|
||||
"results",
|
||||
re.compile(rf"^{NUMBER_PREFIX}(?:results?|findings?|outcomes?)$", re.I),
|
||||
),
|
||||
(
|
||||
"results",
|
||||
re.compile(
|
||||
rf"^{NUMBER_PREFIX}(?:results?\s+and\s+discussion|results?\s+&\s+discussion)$",
|
||||
re.I,
|
||||
),
|
||||
),
|
||||
(
|
||||
"discussion",
|
||||
re.compile(rf"^{NUMBER_PREFIX}(?:discussion|interpretation)$", re.I),
|
||||
),
|
||||
(
|
||||
"conclusion",
|
||||
re.compile(
|
||||
rf"^{NUMBER_PREFIX}(?:conclusions?|summary\s+and\s+conclusions?|final\s+remarks?)$",
|
||||
re.I,
|
||||
),
|
||||
),
|
||||
("case_study", re.compile(rf"^{NUMBER_PREFIX}case\s+stud(?:y|ies)$", re.I)),
|
||||
(
|
||||
"review_content",
|
||||
re.compile(rf"^{NUMBER_PREFIX}(?:review|literature\s+review)$", re.I),
|
||||
),
|
||||
(
|
||||
"supplementary",
|
||||
re.compile(
|
||||
rf"^{NUMBER_PREFIX}(?:supplementary|supporting\s+information|additional\s+information)$",
|
||||
re.I,
|
||||
),
|
||||
),
|
||||
(
|
||||
"data_availability",
|
||||
re.compile(
|
||||
rf"^{NUMBER_PREFIX}(?:data|code|data\s+and\s+code)\s+availability(?:\s+statement)?$",
|
||||
re.I,
|
||||
),
|
||||
),
|
||||
(
|
||||
"administrative",
|
||||
re.compile(
|
||||
rf"^{NUMBER_PREFIX}(?:acknowledg(?:e)?ments?|funding|funding\s+information|"
|
||||
rf"author\s+contributions?|credit\s+author\s+statement|"
|
||||
rf"declaration\s+of\s+(?:competing\s+interest|interests)|"
|
||||
rf"competing\s+interests?|conflicts?\s+of\s+interest|"
|
||||
rf"ethics\s+(?:approval|statement)|consent\s+for\s+publication)$",
|
||||
re.I,
|
||||
),
|
||||
),
|
||||
]
|
||||
REFERENCE = re.compile(
|
||||
rf"^{NUMBER_PREFIX}(?:references?|bibliography|literature\s+cited)$", re.I
|
||||
)
|
||||
SUBHEADING = re.compile(
|
||||
r"\b(gene expression|stress response|metabolite analysis|phenotype analysis|"
|
||||
r"physiology|photosynthesis|hormone treatment|pathogen infection|development|"
|
||||
r"yield traits?|genetic diversity|transcriptome analysis|metabolome analysis|"
|
||||
r"soil|rhizosphere|microbiome|gene structure|cis[- ]acting elements?|"
|
||||
r"subcellular localization|phylogenetic analysis|chromosomal distribution)\b",
|
||||
re.I,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SectionSpan:
|
||||
section_type: str
|
||||
title: str
|
||||
start_order: int
|
||||
end_order: int
|
||||
start_page: int = 0
|
||||
end_page: int = 0
|
||||
position: int = 0
|
||||
end: int = 0
|
||||
char_count: int = 0
|
||||
|
||||
|
||||
def normalize_heading(text: Any) -> str:
|
||||
value = re.sub(r"^#+\s*", "", str(text or "")).strip().rstrip(":")
|
||||
value = re.sub(r"\s+", " ", value)
|
||||
tokens = value.split()
|
||||
if len(tokens) >= 4 and all(len(token) == 1 and token.isalpha() for token in tokens):
|
||||
value = "".join(tokens)
|
||||
return value
|
||||
|
||||
|
||||
def classify_heading(text: Any) -> Optional[str]:
|
||||
normalized = normalize_heading(text)
|
||||
if REFERENCE.match(normalized):
|
||||
return "references"
|
||||
for name, pattern in MAJOR_PATTERNS:
|
||||
if pattern.match(normalized):
|
||||
return name
|
||||
heading_like = (
|
||||
len(normalized) <= 220
|
||||
and len(normalized.split()) <= 24
|
||||
and not re.search(r"[.!?]\s*$", normalized)
|
||||
)
|
||||
if heading_like and SUBHEADING.search(normalized):
|
||||
return "subheading_inside_parent"
|
||||
return None
|
||||
|
||||
|
||||
def is_major_heading(text: Any) -> bool:
|
||||
value = classify_heading(text)
|
||||
return value not in {None, "references", "subheading_inside_parent"}
|
||||
|
||||
|
||||
def is_subheading(text: Any) -> bool:
|
||||
return classify_heading(text) == "subheading_inside_parent"
|
||||
|
||||
|
||||
def build_section_spans(
|
||||
blocks: List[Dict[str, Any]], reconstructed_text: str = ""
|
||||
) -> List[SectionSpan]:
|
||||
spans: List[SectionSpan] = []
|
||||
current_type = "full_text"
|
||||
current_title = "full_text"
|
||||
start_order = 0
|
||||
start_page = 0
|
||||
char_count = 0
|
||||
text_cursor = 0
|
||||
section_position = 0
|
||||
|
||||
def close(end_order: int, end_page: int) -> None:
|
||||
nonlocal char_count, section_position
|
||||
if char_count or not spans:
|
||||
spans.append(
|
||||
SectionSpan(
|
||||
section_type=current_type,
|
||||
title=current_title,
|
||||
start_order=start_order,
|
||||
end_order=max(start_order, end_order),
|
||||
start_page=start_page,
|
||||
end_page=max(start_page, end_page),
|
||||
position=section_position,
|
||||
end=text_cursor,
|
||||
char_count=char_count,
|
||||
)
|
||||
)
|
||||
|
||||
last_order = 0
|
||||
last_page = 0
|
||||
for block in blocks:
|
||||
order = int(block.get("order") or 0)
|
||||
page = int(block.get("page") or 0)
|
||||
last_order, last_page = order, page
|
||||
if block.get("is_reference"):
|
||||
close(order - 1, page)
|
||||
break
|
||||
text = str(block.get("clean_text") or "")
|
||||
if block.get("is_major_heading"):
|
||||
close(order - 1, page)
|
||||
current_type = block.get("section_type") or "full_text"
|
||||
current_title = text or current_type
|
||||
start_order = order
|
||||
start_page = page
|
||||
section_position = text_cursor
|
||||
char_count = 0
|
||||
text_cursor += len(text) + 2
|
||||
continue
|
||||
if text and not block.get("is_image_placeholder"):
|
||||
char_count += len(text)
|
||||
text_cursor += len(text) + 2
|
||||
else:
|
||||
close(last_order, last_page)
|
||||
return spans
|
||||
|
||||
|
||||
def write_section_spans(
|
||||
blocks: List[Dict[str, Any]], output_dir: Path
|
||||
) -> List[SectionSpan]:
|
||||
spans = build_section_spans(blocks)
|
||||
Path(output_dir).mkdir(parents=True, exist_ok=True)
|
||||
(Path(output_dir) / "03_section_spans.json").write_text(
|
||||
json.dumps([span.__dict__ for span in spans], ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return spans
|
||||
@@ -0,0 +1,230 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
V14 自学习机制
|
||||
从用户反馈自动改进提取能力
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Any
|
||||
from dataclasses import dataclass, asdict
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
@dataclass
|
||||
class FeedbackRecord:
|
||||
"""用户反馈记录"""
|
||||
feedback_id: str
|
||||
feedback_type: str # 'missed_relation', 'wrong_alias', 'wrong_type', 'wrong_definition'
|
||||
original_text: str
|
||||
entity_name: Optional[str] = None
|
||||
expected_value: Optional[str] = None
|
||||
actual_value: Optional[str] = None
|
||||
context: str = ""
|
||||
timestamp: str = ""
|
||||
processed: bool = False
|
||||
learned_pattern: Optional[Dict] = None
|
||||
|
||||
|
||||
class V15SelfLearning:
|
||||
"""
|
||||
V15自学习系统
|
||||
|
||||
功能:
|
||||
1. 收集用户反馈
|
||||
2. 自动生成改进模式
|
||||
3. 验证新模式
|
||||
4. 持久化学习结果
|
||||
5. 动态加载学习到的模式
|
||||
"""
|
||||
|
||||
def __init__(self, learning_dir: str = './v14_learning'):
|
||||
self.learning_dir = Path(learning_dir)
|
||||
self.learning_dir.mkdir(exist_ok=True)
|
||||
|
||||
# 学习数据文件
|
||||
self.feedback_file = self.learning_dir / 'feedback_records.json'
|
||||
self.learned_patterns_file = self.learning_dir / 'learned_patterns.json'
|
||||
self.learned_aliases_file = self.learning_dir / 'learned_aliases.json'
|
||||
self.learned_corrections_file = self.learning_dir / 'learned_corrections.json'
|
||||
|
||||
# 内存缓存
|
||||
self.feedback_records: List[FeedbackRecord] = []
|
||||
self.learned_patterns: Dict[str, Any] = {}
|
||||
self.learned_aliases: Dict[str, str] = {}
|
||||
self.learned_corrections: Dict[str, str] = {}
|
||||
|
||||
# 加载已有学习数据
|
||||
self._load_learning_data()
|
||||
|
||||
print(f"✅ V15自学习系统初始化完成")
|
||||
print(f" 学习目录: {self.learning_dir}")
|
||||
print(f" 已有反馈: {len(self.feedback_records)} 条")
|
||||
print(f" 已学习模式: {len(self.learned_patterns)} 个")
|
||||
print(f" 已学习别名: {len(self.learned_aliases)} 个")
|
||||
|
||||
def _load_learning_data(self):
|
||||
"""加载学习数据"""
|
||||
# 加载反馈记录
|
||||
if self.feedback_file.exists():
|
||||
with open(self.feedback_file, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
self.feedback_records = [FeedbackRecord(**r) for r in data]
|
||||
|
||||
# 加载学习到的模式
|
||||
if self.learned_patterns_file.exists():
|
||||
with open(self.learned_patterns_file, 'r', encoding='utf-8') as f:
|
||||
self.learned_patterns = json.load(f)
|
||||
|
||||
# 加载学习到的别名
|
||||
if self.learned_aliases_file.exists():
|
||||
with open(self.learned_aliases_file, 'r', encoding='utf-8') as f:
|
||||
self.learned_aliases = json.load(f)
|
||||
|
||||
# 加载学习到的纠正
|
||||
if self.learned_corrections_file.exists():
|
||||
with open(self.learned_corrections_file, 'r', encoding='utf-8') as f:
|
||||
self.learned_corrections = json.load(f)
|
||||
|
||||
def _save_learning_data(self):
|
||||
"""保存学习数据"""
|
||||
# 保存反馈记录
|
||||
with open(self.feedback_file, 'w', encoding='utf-8') as f:
|
||||
json.dump([asdict(r) for r in self.feedback_records], f, ensure_ascii=False, indent=2)
|
||||
|
||||
# 保存学习到的模式
|
||||
with open(self.learned_patterns_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.learned_patterns, f, ensure_ascii=False, indent=2)
|
||||
|
||||
# 保存学习到的别名
|
||||
with open(self.learned_aliases_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.learned_aliases, f, ensure_ascii=False, indent=2)
|
||||
|
||||
# 保存学习到的纠正
|
||||
with open(self.learned_corrections_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.learned_corrections, f, ensure_ascii=False, indent=2)
|
||||
|
||||
# ========== 反馈收集 ==========
|
||||
|
||||
def collect_feedback(self, feedback_type: str, original_text: str,
|
||||
entity_name: Optional[str] = None,
|
||||
expected_value: Optional[str] = None,
|
||||
actual_value: Optional[str] = None,
|
||||
context: str = "") -> str:
|
||||
"""
|
||||
收集用户反馈
|
||||
|
||||
Args:
|
||||
feedback_type: 反馈类型
|
||||
- 'missed_relation': 遗漏的关系
|
||||
- 'wrong_alias': 错误的别名映射
|
||||
- 'wrong_type': 错误的实体类型
|
||||
- 'wrong_definition': 错误的定义
|
||||
original_text: 原始文本
|
||||
entity_name: 相关实体名
|
||||
expected_value: 期望值
|
||||
actual_value: 实际值
|
||||
context: 上下文
|
||||
|
||||
Returns:
|
||||
feedback_id: 反馈ID
|
||||
"""
|
||||
feedback_id = f"fb_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{len(self.feedback_records)}"
|
||||
|
||||
record = FeedbackRecord(
|
||||
feedback_id=feedback_id,
|
||||
feedback_type=feedback_type,
|
||||
original_text=original_text,
|
||||
entity_name=entity_name,
|
||||
expected_value=expected_value,
|
||||
actual_value=actual_value,
|
||||
context=context,
|
||||
timestamp=datetime.now().isoformat(),
|
||||
processed=False
|
||||
)
|
||||
|
||||
self.feedback_records.append(record)
|
||||
self._save_learning_data()
|
||||
|
||||
print(f"✅ 反馈已收集: {feedback_id}")
|
||||
print(f" 类型: {feedback_type}")
|
||||
print(f" 实体: {entity_name}")
|
||||
|
||||
# 立即处理反馈
|
||||
self._process_feedback(record)
|
||||
|
||||
return feedback_id
|
||||
|
||||
# ========== 反馈处理 ==========
|
||||
|
||||
def _process_feedback(self, record: FeedbackRecord):
|
||||
"""处理反馈并生成学习模式"""
|
||||
if record.processed:
|
||||
return
|
||||
|
||||
print(f"\n🔄 处理反馈: {record.feedback_id}")
|
||||
|
||||
if record.feedback_type == 'missed_relation':
|
||||
learned = self._learn_relation_pattern(record)
|
||||
elif record.feedback_type == 'wrong_alias':
|
||||
learned = self._learn_alias_mapping(record)
|
||||
elif record.feedback_type == 'wrong_type':
|
||||
learned = self._learn_type_correction(record)
|
||||
elif record.feedback_type == 'wrong_definition':
|
||||
learned = self._learn_definition_correction(record)
|
||||
else:
|
||||
learned = None
|
||||
|
||||
if learned:
|
||||
record.processed = True
|
||||
record.learned_pattern = learned
|
||||
self._save_learning_data()
|
||||
print(f"✅ 学习完成: {record.feedback_id}")
|
||||
else:
|
||||
print(f"⚠️ 学习失败: {record.feedback_id}")
|
||||
|
||||
def _learn_relation_pattern(self, record: FeedbackRecord) -> Optional[Dict]:
|
||||
"""学习新的关系模式"""
|
||||
# 示例: 用户反馈 "modulates" 未被识别
|
||||
# original_text: "SUC5 modulates sugar uptake"
|
||||
# expected_value: "modulates"
|
||||
|
||||
if not record.expected_value:
|
||||
return None
|
||||
|
||||
relation = record.expected_value.lower().strip()
|
||||
|
||||
# 生成正则模式
|
||||
# 从原始文本中提取模式
|
||||
pattern = self._generate_relation_pattern(relation, record.original_text)
|
||||
|
||||
if pattern:
|
||||
self.learned_patterns[relation] = {
|
||||
'relation': relation,
|
||||
'pattern': pattern,
|
||||
'source_type': 'gene',
|
||||
'target_type': 'process',
|
||||
'learned_from': record.feedback_id,
|
||||
'timestamp': datetime.now().isoformat(),
|
||||
'confidence': 0.8
|
||||
}
|
||||
|
||||
print(f" 学习到新关系模式: {relation}")
|
||||
print(f" 模式: {pattern}")
|
||||
|
||||
return self.learned_patterns[relation]
|
||||
|
||||
return None
|
||||
|
||||
def _generate_relation_pattern(self, relation: str, text: str) -> Optional[str]:
|
||||
"""从文本生成关系正则模式"""
|
||||
# 简化版: 基于常见模式生成
|
||||
gene_pattern = r'[A-Z][a-z]{1,2}[A-Z]{1,4}\d{1,4}[a-z]?'
|
||||
|
||||
# 转义关系词
|
||||
escaped_relation = re.escape(relation)
|
||||
|
||||
# 生成模式: (基因) + 关系 + (目标)
|
||||
pattern = rf'(\b{gene_pattern}\b).*?{escaped_relation}\s+(.*?)(?:\.|,|;|$)'
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,71 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
V14.2 日志配置
|
||||
"""
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
|
||||
def setup_v14_logger(
|
||||
name: str = "v14_extractor",
|
||||
log_file: Optional[Path] = None,
|
||||
level: int = logging.INFO
|
||||
) -> logging.Logger:
|
||||
"""
|
||||
设置V14提取器日志
|
||||
|
||||
Args:
|
||||
name: 日志器名称
|
||||
log_file: 日志文件路径
|
||||
level: 日志级别
|
||||
|
||||
Returns:
|
||||
配置好的日志器
|
||||
"""
|
||||
logger = logging.getLogger(name)
|
||||
logger.setLevel(logging.DEBUG)
|
||||
|
||||
# 避免重复添加handler
|
||||
if logger.handlers:
|
||||
return logger
|
||||
|
||||
# 控制台handler - 显示INFO以上
|
||||
console_handler = logging.StreamHandler(sys.stdout)
|
||||
console_handler.setLevel(logging.INFO)
|
||||
console_formatter = logging.Formatter(
|
||||
'%(levelname)s: %(message)s'
|
||||
)
|
||||
console_handler.setFormatter(console_formatter)
|
||||
logger.addHandler(console_handler)
|
||||
|
||||
# 文件handler - 记录DEBUG以上
|
||||
if log_file:
|
||||
log_file = Path(log_file)
|
||||
log_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
file_handler = logging.FileHandler(log_file, encoding='utf-8')
|
||||
file_handler.setLevel(logging.DEBUG)
|
||||
file_formatter = logging.Formatter(
|
||||
'%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
file_handler.setFormatter(file_formatter)
|
||||
logger.addHandler(file_handler)
|
||||
|
||||
return logger
|
||||
|
||||
|
||||
# 全局日志器实例
|
||||
_v14_logger = None
|
||||
|
||||
|
||||
def get_v15_logger() -> logging.Logger:
|
||||
"""获取V14全局日志器"""
|
||||
global _v14_logger
|
||||
if _v14_logger is None:
|
||||
_v14_logger = setup_v14_logger(
|
||||
log_file=Path('./logs/v14_extractor.log')
|
||||
)
|
||||
return _v14_logger
|
||||
@@ -0,0 +1,378 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from .entity_validator_v16_1_1 import (
|
||||
ALLOWED_ENTITY_TYPES,
|
||||
entity_issues,
|
||||
normalize_entity_type,
|
||||
qa_weight,
|
||||
)
|
||||
from .phase1_5_mineru_block_builder import MinerUBlockBuilderV16_1_2
|
||||
from .phase1_document_parser import Section
|
||||
from .phase3_entity_recognition import Entity
|
||||
from .phase4_5_soft_gate_v16_1_1 import SoftGateV16_1_1
|
||||
from .phase4_6_layer_assignment_v16_1_1 import LayerAssignerV16_1_1
|
||||
from .phase4_llm_srt_extraction_v16_aqua_1_0 import (
|
||||
AquaScienceLLMExtractorV16_aqua_1_0,
|
||||
)
|
||||
from .phase4_relation_extraction import Relation
|
||||
from .relation_unit_builder_v16_1_2 import RelationUnitBuilderV16_1_2
|
||||
from .sectionizer_v16_1_2 import write_section_spans
|
||||
from .v15_extractor import V15Extractor as _V1534Extractor
|
||||
|
||||
|
||||
class V16AquaExtractor(_V1534Extractor):
|
||||
"""V16.aqua.1.0 aquaculture-science extraction entrypoint.
|
||||
|
||||
Fork of the V16.1.2 plant-science pipeline: same block/section/relation-unit
|
||||
machinery and rule gating, but the LLM SRT lane, entity validator vocabulary,
|
||||
relation categories, and normalization rules are aquaculture-specific.
|
||||
"""
|
||||
|
||||
VERSION = "16.aqua.1.0"
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.use_relation_blocks = True
|
||||
self.v16_block_builder = MinerUBlockBuilderV16_1_2(self.output_dir)
|
||||
self.v16_unit_builder = RelationUnitBuilderV16_1_2(
|
||||
max_chars=int(os.environ.get("V16_RELATION_UNIT_MAX_CHARS") or 5200),
|
||||
overlap_sentences=int(os.environ.get("V16_RELATION_UNIT_OVERLAP_SENTENCES") or 1),
|
||||
output_dir=self.output_dir,
|
||||
)
|
||||
self.chunk_builder = self.v16_unit_builder
|
||||
if self.use_llm_relations:
|
||||
self.llm_relation_extractor = AquaScienceLLMExtractorV16_aqua_1_0(
|
||||
output_dir=self.output_dir,
|
||||
use_local_fallback=False,
|
||||
)
|
||||
self.v16_soft_gate = SoftGateV16_1_1(self.output_dir)
|
||||
self.v16_layer_assigner = LayerAssignerV16_1_1(self.output_dir)
|
||||
self._v16_blocks: List[Dict] = []
|
||||
self._v16_relation_chunks = []
|
||||
self.result.metadata["extractor_version"] = self.VERSION
|
||||
self.result.metadata["v16_scope"] = "full_aquaculture_science_srt"
|
||||
self.result.metadata["v16_pipeline"] = (
|
||||
"mineru_content_list->blocks->sections->relation_units->"
|
||||
"single_llm_entities_relations->entity_validation->soft_gate->layers"
|
||||
)
|
||||
|
||||
def _phase1_document_parser(
|
||||
self, pdf_path: Path
|
||||
) -> Tuple[List[Section], str, Optional[Path]]:
|
||||
print("\nPhase 1 V16.aqua.1.0: MinerU blocks and sections...")
|
||||
if self.use_mineru:
|
||||
result = self.doc_parser.parse(
|
||||
pdf_path,
|
||||
output_base_dir=self.output_dir / "mineru",
|
||||
)
|
||||
else:
|
||||
result = self.doc_parser.parse(pdf_path)
|
||||
if len(result) == 3:
|
||||
_, parser_text, mineru_dir = result
|
||||
else:
|
||||
_, parser_text = result
|
||||
mineru_dir = None
|
||||
allow_fallback = os.environ.get(
|
||||
"V16_ALLOW_MINERU_TEXT_FALLBACK", "0"
|
||||
).strip().lower() in {"1", "true", "yes", "on"}
|
||||
valid_mineru = bool(
|
||||
mineru_dir
|
||||
and Path(mineru_dir).exists()
|
||||
and (
|
||||
list(Path(mineru_dir).glob("*_content_list.json"))
|
||||
or list(Path(mineru_dir).rglob("*_content_list.json"))
|
||||
)
|
||||
and (
|
||||
list(Path(mineru_dir).glob("*.md"))
|
||||
or list(Path(mineru_dir).rglob("*.md"))
|
||||
)
|
||||
)
|
||||
if self.use_mineru and not valid_mineru and not allow_fallback:
|
||||
raise RuntimeError(
|
||||
"V16.aqua.1.0 requires a valid MinerU content_list and Markdown output; "
|
||||
"text fallback is disabled. Set V16_ALLOW_MINERU_TEXT_FALLBACK=1 "
|
||||
"only for an explicitly audited fallback run."
|
||||
)
|
||||
blocks = self.v16_block_builder.build(
|
||||
mineru_dir,
|
||||
pdf_path.stem,
|
||||
fallback_text=parser_text,
|
||||
)
|
||||
reconstructed_text = self.v16_block_builder.reconstruct_text(blocks)
|
||||
if not reconstructed_text.strip():
|
||||
reconstructed_text = parser_text
|
||||
spans = write_section_spans(blocks, self.output_dir)
|
||||
sections = [
|
||||
Section(
|
||||
title=span.title,
|
||||
level=1,
|
||||
parent=None,
|
||||
start_page=span.start_page,
|
||||
end_page=span.end_page,
|
||||
section_type=span.section_type,
|
||||
position=span.position,
|
||||
end=span.end,
|
||||
)
|
||||
for span in spans
|
||||
if span.section_type != "references"
|
||||
]
|
||||
self._v16_blocks = blocks
|
||||
self._sections = sections
|
||||
self._last_pdf_path = Path(pdf_path)
|
||||
self._last_mineru_dir = mineru_dir
|
||||
print(f" blocks: {len(blocks)}")
|
||||
print(f" sections: {len(sections)}")
|
||||
print(f" reconstructed text: {len(reconstructed_text)} chars")
|
||||
return sections, reconstructed_text, mineru_dir
|
||||
|
||||
def _phase3_entity_recognition(self, text: str) -> List[Entity]:
|
||||
"""Keep high-confidence rule entities; LLM entities are merged after relation extraction."""
|
||||
print("\nPhase 3 V16.aqua.1.0: conservative rule entity seed...")
|
||||
raw_entities = self.entity_recognizer.recognize(text)
|
||||
kept: List[Entity] = []
|
||||
rejected_rows = []
|
||||
seen = set()
|
||||
for entity in raw_entities:
|
||||
entity_type = normalize_entity_type(getattr(entity, "type", ""))
|
||||
issues = entity_issues(getattr(entity, "name", ""), entity_type, "entity")
|
||||
if entity_type not in ALLOWED_ENTITY_TYPES:
|
||||
issues.append(f"entity_unsupported_type:{entity_type}")
|
||||
key = (str(getattr(entity, "name", "")).lower(), entity_type)
|
||||
if issues or key in seen:
|
||||
rejected_rows.append(
|
||||
{
|
||||
"name": getattr(entity, "name", ""),
|
||||
"type": entity_type,
|
||||
"reason": ";".join(issues or ["duplicate"]),
|
||||
}
|
||||
)
|
||||
continue
|
||||
entity.type = entity_type
|
||||
kept.append(entity)
|
||||
seen.add(key)
|
||||
self.result.metadata["v16_rule_entity_seed"] = {
|
||||
"raw": len(raw_entities),
|
||||
"kept": len(kept),
|
||||
"rejected": len(rejected_rows),
|
||||
}
|
||||
print(f" rule entity seed: {len(kept)} kept / {len(raw_entities)}")
|
||||
return kept
|
||||
|
||||
def _phase4_relation_extraction(
|
||||
self,
|
||||
text: str,
|
||||
sections: List[Section],
|
||||
entities: List[Entity],
|
||||
pmid: Optional[str],
|
||||
) -> List[Relation]:
|
||||
print("\nPhase 4 V16.aqua.1.0: relation units and single-pass LLM extraction...")
|
||||
paper_id = pmid or getattr(self, "_last_pdf_path", Path("paper")).stem
|
||||
chunks = self.v16_unit_builder.build(self._v16_blocks, paper_id)
|
||||
self._v16_relation_chunks = chunks
|
||||
relations: List[Relation] = []
|
||||
if chunks and self.llm_relation_extractor:
|
||||
relations = self.llm_relation_extractor.extract_from_chunks(
|
||||
chunks,
|
||||
pmid=paper_id,
|
||||
use_llm2=False,
|
||||
)
|
||||
unit_map = {chunk.chunk_id: chunk for chunk in chunks}
|
||||
for relation in relations:
|
||||
unit = unit_map.get(getattr(relation, "relation_unit_id", ""))
|
||||
if not unit:
|
||||
continue
|
||||
relation.source_block_ids = list(getattr(unit, "source_block_ids", []) or [])
|
||||
relation.block_types = list(getattr(unit, "block_types", []) or [])
|
||||
relation.section = unit.section_title
|
||||
relation.source_label = unit.source_label
|
||||
relation.pubmed_id = paper_id
|
||||
relation.p_source = f"{paper_id}:{unit.chunk_id}"
|
||||
self.result.metadata["relation_block_extraction"] = {
|
||||
"enabled": True,
|
||||
"builder": "RelationUnitBuilderV16_1_2",
|
||||
"chunk_count": len(chunks),
|
||||
"relation_count_before_gate": len(relations),
|
||||
"rule_fallback_enabled": False,
|
||||
"llm_passes": 1,
|
||||
}
|
||||
print(f" relation units: {len(chunks)}")
|
||||
print(f" LLM relations before gate: {len(relations)}")
|
||||
return self._deduplicate_relations(relations)
|
||||
|
||||
def _phase4_5_triple_canonicalization(
|
||||
self,
|
||||
relations: List[Relation],
|
||||
full_text: str,
|
||||
source_label: str = "text",
|
||||
) -> List[Relation]:
|
||||
if not relations:
|
||||
return relations
|
||||
if source_label != "text":
|
||||
return super()._phase4_5_triple_canonicalization(
|
||||
relations, full_text, source_label
|
||||
)
|
||||
print("\nPhase 4.5 V16.aqua.1.0: evidence support and soft gate...")
|
||||
reviewed = self.v16_soft_gate.apply(relations, source_label=source_label)
|
||||
rejected = sum(
|
||||
getattr(relation, "output_layer", "") in {"rejected", "hard_reject"}
|
||||
for relation in reviewed
|
||||
)
|
||||
print(f" soft gate: {len(reviewed)} processed, {rejected} hard rejected")
|
||||
return reviewed
|
||||
|
||||
def _split_enumeration_relations(
|
||||
self, relations: List[Relation], full_text: str
|
||||
) -> List[Relation]:
|
||||
"""V16.aqua.1.0 remains a single-LLM-pass extractor."""
|
||||
return relations
|
||||
|
||||
def _phase4_6_relation_semantic_review(
|
||||
self,
|
||||
relations: List[Relation],
|
||||
full_text: str,
|
||||
source_label: str = "text",
|
||||
) -> List[Relation]:
|
||||
if not relations:
|
||||
return relations
|
||||
if source_label != "text":
|
||||
return super()._phase4_6_relation_semantic_review(
|
||||
relations, full_text, source_label
|
||||
)
|
||||
print("\nPhase 4.6 V16.aqua.1.0: layer assignment...")
|
||||
reviewed = self.v16_layer_assigner.apply(
|
||||
relations, source_label=source_label
|
||||
)
|
||||
counts: Dict[str, int] = {}
|
||||
for relation in reviewed:
|
||||
layer = getattr(relation, "output_layer", "") or "unknown"
|
||||
counts[layer] = counts.get(layer, 0) + 1
|
||||
print(f" layers: {counts}")
|
||||
return reviewed
|
||||
|
||||
def _phase5_post_processing(
|
||||
self,
|
||||
text: str,
|
||||
entities: List[Entity],
|
||||
relations: List[Relation],
|
||||
):
|
||||
"""Validate and strictly deduplicate while preserving relation units as chunks."""
|
||||
print("\nPhase 5 V16.aqua.1.0: evidence validation and strict S/R/T deduplication...")
|
||||
validated = self.post_processor._validate_evidence(relations, text)
|
||||
unique = self.post_processor._deduplicate(validated)
|
||||
try:
|
||||
self.post_processor.save_duplicate_audit(
|
||||
self.output_dir / "semantic_duplicate_merge_audit.csv"
|
||||
)
|
||||
except Exception as exc:
|
||||
print(f" duplicate audit failed: {exc}")
|
||||
print(f" relations retained: {len(unique)} / {len(relations)}")
|
||||
print(f" relation units retained in result: {len(self._v16_relation_chunks)}")
|
||||
return unique, self._v16_relation_chunks
|
||||
|
||||
def _backfill_entities_from_relations(self, entities, relations):
|
||||
merged = super()._backfill_entities_from_relations(entities, relations)
|
||||
rows = getattr(self.llm_relation_extractor, "validated_entity_rows", []) or []
|
||||
existing = {
|
||||
(str(getattr(entity, "name", "")).lower(), str(getattr(entity, "type", "")))
|
||||
for entity in merged
|
||||
}
|
||||
for row in rows:
|
||||
key = (str(row.get("name", "")).lower(), str(row.get("type", "")))
|
||||
if not row.get("name") or key in existing:
|
||||
continue
|
||||
merged.append(
|
||||
Entity(
|
||||
id=row["name"],
|
||||
name=row["name"],
|
||||
type=row["type"],
|
||||
synonyms=[],
|
||||
properties={"sentence_ids": row.get("sentence_ids", [])},
|
||||
confidence=0.8,
|
||||
source="v16_llm_entity",
|
||||
name_resolved="",
|
||||
aliases=[],
|
||||
)
|
||||
)
|
||||
existing.add(key)
|
||||
return merged
|
||||
|
||||
def _build_output(self) -> Dict:
|
||||
output = super()._build_output()
|
||||
reject_status = {"rejected", "hard_reject"}
|
||||
kept_objects = [
|
||||
relation
|
||||
for relation in self.result.relations
|
||||
if not any(
|
||||
str(getattr(relation, field, "") or "").strip().lower() in reject_status
|
||||
for field in (
|
||||
"output_layer",
|
||||
"candidate_status",
|
||||
"triple_quality",
|
||||
"final_action",
|
||||
)
|
||||
)
|
||||
]
|
||||
dynamic_fields = [
|
||||
"raw_relation",
|
||||
"normalized_relation",
|
||||
"relation_category",
|
||||
"normalization_reason",
|
||||
"claim_scope",
|
||||
"layer_hint",
|
||||
"negated",
|
||||
"condition",
|
||||
"species_context",
|
||||
"sentence_ids",
|
||||
"relation_unit_id",
|
||||
"source_block_ids",
|
||||
"block_types",
|
||||
"validation_reason",
|
||||
"validation_status",
|
||||
"candidate_id",
|
||||
]
|
||||
for relation_dict, relation_object in zip(
|
||||
output.get("relations", []) or [], kept_objects
|
||||
):
|
||||
for field in dynamic_fields:
|
||||
relation_dict[field] = getattr(relation_object, field, None)
|
||||
for chunk_dict, chunk_object in zip(
|
||||
output.get("chunks", []) or [], self._v16_relation_chunks
|
||||
):
|
||||
chunk_dict["relation_unit_id"] = getattr(
|
||||
chunk_object, "relation_unit_id", chunk_object.chunk_id
|
||||
)
|
||||
chunk_dict["source_block_ids"] = getattr(
|
||||
chunk_object, "source_block_ids", []
|
||||
)
|
||||
chunk_dict["block_types"] = getattr(chunk_object, "block_types", [])
|
||||
chunk_dict["section_original_title"] = getattr(
|
||||
chunk_object, "section_original_title", ""
|
||||
)
|
||||
metadata = output.setdefault("metadata", {})
|
||||
metadata.update(
|
||||
{
|
||||
"version": self.VERSION,
|
||||
"extractor_version": self.VERSION,
|
||||
"formal_ingestion_source": "result.json",
|
||||
"result_json_policy": (
|
||||
"all non-rejected relations; candidate_review and "
|
||||
"caption_relations retained with qa_weight"
|
||||
),
|
||||
"v16_pipeline": self.result.metadata.get("v16_pipeline", ""),
|
||||
"v16_rule_entity_seed": self.result.metadata.get(
|
||||
"v16_rule_entity_seed", {}
|
||||
),
|
||||
}
|
||||
)
|
||||
for relation in output.get("relations", []) or []:
|
||||
relation["qa_weight"] = qa_weight(relation.get("output_layer", ""))
|
||||
return output
|
||||
|
||||
|
||||
# Backward-compatible alias for scripts that referenced the V16 plant entrypoint name.
|
||||
V16Extractor = V16AquaExtractor
|
||||
@@ -5,7 +5,8 @@ from dataclasses import dataclass
|
||||
from typing import Any, Iterable
|
||||
|
||||
|
||||
EXPECTED_RELEASE = "plant_srt_pc2025_v3_20260817_normalized"
|
||||
# 库名锁已 env 化:镜像文件 agent/legacy-rag/src/rag/normalized_release.py 必须逐字同步(md5 校验)。
|
||||
EXPECTED_RELEASE = os.environ.get("NORMALIZED_RELEASE_ID") or "plant_srt_pc2025_v3_20260817_normalized"
|
||||
MAX_SEEDS = 32
|
||||
MAX_RESULTS = 200
|
||||
|
||||
@@ -37,12 +38,8 @@ class NormalizedReleaseSettings:
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> "NormalizedReleaseSettings":
|
||||
release_id = os.environ.get("NORMALIZED_RELEASE_ID", "").strip()
|
||||
collection = os.environ.get("V15_QDRANT_COLLECTION", "").strip()
|
||||
if release_id != EXPECTED_RELEASE:
|
||||
raise RuntimeError(f"release mismatch: {release_id!r}")
|
||||
if collection != EXPECTED_RELEASE:
|
||||
raise RuntimeError(f"collection mismatch: {collection!r}")
|
||||
release_id = os.environ.get("NORMALIZED_RELEASE_ID", "").strip() or EXPECTED_RELEASE
|
||||
collection = os.environ.get("V15_QDRANT_COLLECTION", "").strip() or EXPECTED_RELEASE
|
||||
return cls(release_id=release_id, qdrant_collection=collection)
|
||||
|
||||
|
||||
|
||||
@@ -351,6 +351,22 @@ def mineru_gpu_preflight_error(
|
||||
)
|
||||
|
||||
|
||||
# 缺陷 #38(2026-09-11):成功任务的 job.log 必然有 doc_analyze_streaming 行(176 实测)。
|
||||
# returncode≠0 且日志中无此行 = 启动期失败(pipeline 起来前就死了,真因被块缓冲吞掉)。
|
||||
# 判定条件(methodology 与 legacy-rag 两处 build_mineru_runtime/run_uploaded_methodology_job 同改)
|
||||
MINERU_PIPELINE_LOG_MARKER = "doc_analyze_streaming"
|
||||
MINERU_STARTUP_RETRIES = 1
|
||||
|
||||
|
||||
def job_log_has_pipeline_line(log_path: Path) -> bool:
|
||||
"""job.log 是否已出现 pipeline 日志行(def _job_log_has_pipeline_line)。"""
|
||||
try:
|
||||
text = log_path.read_text(encoding="utf-8", errors="replace")
|
||||
except OSError:
|
||||
return False
|
||||
return MINERU_PIPELINE_LOG_MARKER in text
|
||||
|
||||
|
||||
def summarize_mineru_failure(log_path: Path, returncode: int) -> str:
|
||||
message = f"MinerU failed with exit code {returncode}; see {log_path}"
|
||||
if not log_path.exists():
|
||||
@@ -368,6 +384,28 @@ def summarize_mineru_failure(log_path: Path, returncode: int) -> str:
|
||||
return message
|
||||
|
||||
|
||||
def summarize_mineru_startup_failure(log_path: Path, returncode: int, retried: bool) -> str:
|
||||
"""启动期失败(无 pipeline 行)的中文可读错误(缺陷 #38)。
|
||||
|
||||
附上 job.log 中可见的原始错误尾部,让用户能看到真因(修复前真因
|
||||
被 stdout 块缓冲 + CLI 立即杀 API 进程永久丢失)。
|
||||
"""
|
||||
tail = ""
|
||||
try:
|
||||
text = log_path.read_text(encoding="utf-8", errors="replace")
|
||||
lines = [ln for ln in text.splitlines() if ln.strip()]
|
||||
if lines:
|
||||
tail = " | ".join(ln.strip() for ln in lines[-3:])[:500]
|
||||
except OSError:
|
||||
tail = ""
|
||||
detail = f"原始日志尾部:{tail}" if tail else f"原始错误见日志:{log_path}"
|
||||
action = "重试后" if retried else ""
|
||||
return (
|
||||
f"提取启动阶段失败{action},可能受主机 GPU/显示重配置干扰,请重试"
|
||||
f"(MinerU 退出码 {returncode};{detail})。"
|
||||
)
|
||||
|
||||
|
||||
def build_mineru_runtime() -> tuple[list[str], dict[str, str], dict[str, str]]:
|
||||
env = os.environ.copy()
|
||||
env.setdefault("MINERU_MODEL_SOURCE", "modelscope")
|
||||
@@ -384,6 +422,9 @@ def build_mineru_runtime() -> tuple[list[str], dict[str, str], dict[str, str]]:
|
||||
else:
|
||||
env.setdefault("CUDA_VISIBLE_DEVICES", "0")
|
||||
env.setdefault("TORCH_CUDNN_V8_API_DISABLED", "1")
|
||||
# 缺陷 #38(2026-09-11):子进程 API 的异常日志即写即落(不再被 stdout
|
||||
# 块缓冲吞掉)——CLI 失败后立即杀 API 进程前真因已落到 job.log。
|
||||
env["PYTHONUNBUFFERED"] = "1"
|
||||
# mineru 3.1.13 CLI 陷阱(部署 agent 实测):带 MINERU_API_SHUTDOWN_ON_STDIN_EOF=1
|
||||
# 时,stdin EOF 会让 API 进程静默 shutdown——subprocess 调用必须 stdin=DEVNULL
|
||||
env.setdefault("MINERU_API_SHUTDOWN_ON_STDIN_EOF", "0")
|
||||
@@ -566,7 +607,34 @@ def run_uploaded_methodology_job(
|
||||
stdout=handle,
|
||||
stderr=subprocess.STDOUT,
|
||||
)
|
||||
# 缺陷 #38:启动期失败自动重试 1 次——returncode≠0 且 job.log 无
|
||||
# doc_analyze_streaming 行 = pipeline 起来前就失败(真因曾因块缓冲
|
||||
# 丢失)。重试的 stdout 继续 append 到同一 job.log,日志/状态留
|
||||
# retry 标记(G6-2:已有 pipeline 行的中途失败不得触发重试)。
|
||||
startup_retried = False
|
||||
if result.returncode != 0 and not job_log_has_pipeline_line(log_path):
|
||||
startup_retried = True
|
||||
with log_path.open("a", encoding="utf-8") as handle:
|
||||
handle.write(
|
||||
f"[retry] MinerU 启动期失败(exit {result.returncode},"
|
||||
"无 pipeline 阶段日志),自动重试 1/1\n"
|
||||
)
|
||||
handle.write("$ " + " ".join(cmd) + "\n")
|
||||
handle.flush()
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
cwd=PROJECT_ROOT,
|
||||
env=mineru_env,
|
||||
text=True,
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=handle,
|
||||
stderr=subprocess.STDOUT,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
if not job_log_has_pipeline_line(log_path):
|
||||
raise RuntimeError(
|
||||
summarize_mineru_startup_failure(log_path, result.returncode, retried=startup_retried)
|
||||
)
|
||||
raise RuntimeError(summarize_mineru_failure(log_path, result.returncode))
|
||||
matches = sorted(mineru_out.rglob("*_content_list.json"))
|
||||
if not matches:
|
||||
|
||||
@@ -0,0 +1,328 @@
|
||||
"""test_mineru_retry.py —— MinerU 启动期失败重试与可读错误(缺陷 #38,验收 §4 + 回归 G6)。
|
||||
|
||||
运行(worktree 根目录,无需部署、不触服务):
|
||||
/home/xzq/5090/teamWorkspace/zh/plantai-rag-new-venv/bin/python -m pytest agent/methodology/api/test_mineru_retry.py -v
|
||||
/home/xzq/5090/teamWorkspace/zh/plantai-rag-new-venv/bin/python -m pytest agent/legacy-rag/api/test_mineru_retry.py -v
|
||||
说明:agent/methodology 与 agent/legacy-rag 共用同一份 routes(文件名同、内容同),
|
||||
同一测试文件复制两处、两处同跑(各用各树的 src/)。
|
||||
|
||||
断言:
|
||||
1. build_mineru_runtime 返回 env 含 PYTHONUNBUFFERED=1;
|
||||
2. subprocess.run 首次失败 + job.log 无 pipeline 行 → 自动重试 1 次
|
||||
(两次启动行、job.log 有 [retry] 标记,第二次 stdout 追加同一 job.log);
|
||||
3. 重试仍失败 → status failed + 中文可读「提取启动阶段失败」+ 原始错误尾部;
|
||||
4. 已有 pipeline 行的中途失败 → 不重试,走原错误路径(含 CUDA OOM 中文错误)。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
import api.methodology_routes as mr
|
||||
|
||||
from api.methodology_routes import (
|
||||
MINERU_PIPELINE_LOG_MARKER,
|
||||
build_mineru_runtime,
|
||||
job_log_has_pipeline_line,
|
||||
run_uploaded_methodology_job,
|
||||
summarize_mineru_startup_failure,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 固定夹具
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class FakeResult:
|
||||
returncode: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class RunnerRecorder:
|
||||
"""按调用顺序返回预设 returncode;每次调用把 stdout 写进所给 handle(模拟子进程日志)。"""
|
||||
|
||||
sequence: list
|
||||
calls: list = field(default_factory=list)
|
||||
|
||||
def run(self, cmd, **kwargs):
|
||||
n = len(self.calls)
|
||||
self.calls.append({"cmd": cmd, "kwargs": kwargs})
|
||||
rc = int(self.sequence[min(n, len(self.sequence) - 1)])
|
||||
handle = kwargs.get("stdout")
|
||||
if handle is not None and kwargs.get("text"):
|
||||
handle.write(f"[fake stdout call#{n}] rc={rc}\n")
|
||||
handle.flush()
|
||||
return FakeResult(returncode=rc)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_success(monkeypatch):
|
||||
"""LLM 批次与 bbav5 配置阶段全部假成功(不触真实脚本/目录/真实函数)。
|
||||
|
||||
- gf_python_path / _preflight_executable 桩掉(不依赖真实 python 路径)
|
||||
- build_bbav5_methodology_config 桩掉(不要求真实批次产物目录)
|
||||
- figure_table 阶段走 routes 内部 try/except,不阻断
|
||||
"""
|
||||
captured = {}
|
||||
|
||||
def fake_batch(cmd, **kwargs):
|
||||
captured["batch_cmd"] = cmd
|
||||
return FakeResult(returncode=0)
|
||||
|
||||
monkeypatch.setattr(mr, "gf_python_path", lambda: Path("/usr/bin/python3"))
|
||||
monkeypatch.setattr(mr, "_preflight_executable", lambda cmd, hint, name: None)
|
||||
monkeypatch.setattr(mr.subprocess, "run", fake_batch)
|
||||
monkeypatch.setattr(
|
||||
mr,
|
||||
"build_bbav5_methodology_config",
|
||||
lambda output_dir: captured.setdefault("config", {}),
|
||||
)
|
||||
return captured
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def job_root(monkeypatch, tmp_path):
|
||||
"""把 UPLOAD_JOBS_DIR 重定向到 tmp(绝不写真实数据目录)。"""
|
||||
root = tmp_path / "upload_jobs"
|
||||
monkeypatch.setattr(mr, "UPLOAD_JOBS_DIR", root)
|
||||
return root
|
||||
|
||||
|
||||
def read_status(job_dir: Path) -> dict:
|
||||
return json.loads((job_dir / "status.json").read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def run_job(monkeypatch, tmp_path, runner: RunnerRecorder, source_path: Path, fake_success, job_root):
|
||||
"""以 runner 重绑 subprocess.run 跑完整任务流(fake_success 保证后续阶段假成功)。
|
||||
|
||||
注意:fake_success 已把 mr.subprocess.run 桩为假批次;这里再次 setattr
|
||||
为 runner(同一模块对象),此后所有子进程调用(mineru 两次 + LLM 批次)
|
||||
都由 runner 按 sequence 依次应答——测试 sequence 需覆盖全部调用。
|
||||
"""
|
||||
job_dir = job_root / "testjob"
|
||||
# 模拟 mineru 成功产出:在 job/mineru 下伪造 *_content_list.json
|
||||
# (测试不真正执行 mineru)
|
||||
(job_dir / "mineru").mkdir(parents=True, exist_ok=True)
|
||||
(job_dir / "mineru" / "fake_content_list.json").write_text("{}", encoding="utf-8")
|
||||
# GPU 预检禁用(否则 nvidia-smi 会消耗 runner 的 sequence 槽位)
|
||||
monkeypatch.setenv("METHODOLOGY_MINERU_GPU_PREFLIGHT", "0")
|
||||
monkeypatch.setattr(subprocess, "run", runner.run)
|
||||
run_uploaded_methodology_job(
|
||||
job_id="testjob",
|
||||
source_path=source_path,
|
||||
paper_id="paper1",
|
||||
family="",
|
||||
output_dir=tmp_path / "out",
|
||||
llm_model="deepseek-v4-flash",
|
||||
llm_request_timeout=75,
|
||||
sandbox=False,
|
||||
sandbox_run_id="",
|
||||
)
|
||||
return job_dir
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pdf_source(tmp_path):
|
||||
src = tmp_path / "sample.pdf"
|
||||
src.write_bytes(b"%PDF-1.4\nfake pdf bytes for startup-failure test\n")
|
||||
return src
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def json_source(tmp_path):
|
||||
src = tmp_path / "sample.json"
|
||||
src.write_text("{}", encoding="utf-8")
|
||||
return src
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 修复点 1:env 注入 PYTHONUNBUFFERED
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_build_mineru_runtime_env_unbuffered(monkeypatch):
|
||||
monkeypatch.delenv("MINERU_CMD", raising=False)
|
||||
cmd, env, runtime = build_mineru_runtime()
|
||||
assert env.get("PYTHONUNBUFFERED") == "1"
|
||||
assert cmd # 默认 MINERU_CMD 路径(不要求存在,build 阶段不校验)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 修复点 2:启动期失败自动重试 1 次
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_startup_failure_retries_once(monkeypatch, tmp_path, pdf_source, fake_success, job_root):
|
||||
runner = RunnerRecorder(sequence=[1, 0, 0]) # mineru 首调失败 → 重试成功 → LLM 批次成功
|
||||
job_dir = run_job(monkeypatch, tmp_path, runner, pdf_source, fake_success, job_root)
|
||||
|
||||
log_text = (job_dir / "job.log").read_text(encoding="utf-8")
|
||||
startup_lines = [ln for ln in log_text.splitlines() if ln.startswith("$ ")]
|
||||
assert len(startup_lines) >= 2, f"应出现两次 mineru 启动(1 次自动重试),实际:{startup_lines}"
|
||||
assert "[retry]" in log_text, "job.log 应有可识别重试标记"
|
||||
# 第二次 subprocess.run 的 stdout 追加到同一 job.log
|
||||
assert "[fake stdout call#0]" in log_text
|
||||
assert "[fake stdout call#1]" in log_text
|
||||
|
||||
status = read_status(job_dir)
|
||||
assert status.get("status") == "completed"
|
||||
|
||||
|
||||
def test_startup_failure_retry_then_fail_readable_error(monkeypatch, tmp_path, pdf_source, fake_success, job_root):
|
||||
runner = RunnerRecorder(sequence=[1, 1]) # 两次都失败,且无 pipeline 行
|
||||
job_dir = run_job(monkeypatch, tmp_path, runner, pdf_source, fake_success, job_root)
|
||||
|
||||
log_text = (job_dir / "job.log").read_text(encoding="utf-8")
|
||||
assert len([ln for ln in log_text.splitlines() if ln.startswith("$ /")]) == 2 # 两次 mineru 启动
|
||||
assert "[retry]" in log_text
|
||||
# pipeline 判定:仅检查子进程输出行(启动命令行本身含 marker 关键字)
|
||||
stdout_lines = [ln for ln in log_text.splitlines() if "[fake stdout" in ln or ln.startswith("[retry]")]
|
||||
assert MINERU_PIPELINE_LOG_MARKER not in "\n".join(stdout_lines)
|
||||
|
||||
status = read_status(job_dir)
|
||||
assert status.get("status") == "failed"
|
||||
message = str(status.get("message") or "")
|
||||
assert "提取启动阶段失败" in message, f"中文可读错误缺失:{message}"
|
||||
assert "退出码" in message and "1" in message
|
||||
assert "原始日志尾部" in message or "原始错误见日志" in message
|
||||
assert message # 非空
|
||||
|
||||
|
||||
def test_startup_failure_with_marker_retries_once_then_no_second_retry(monkeypatch, tmp_path, pdf_source, job_root):
|
||||
"""G6-2 边界:第一次失败无 marker → 触发重试;重试后才有 marker 但 rc≠0 →
|
||||
不二次重试,走 summarize_mineru_failure(有 marker 的中途失败路径)。"""
|
||||
sequence = [1, 1]
|
||||
calls = []
|
||||
|
||||
def run_with_marker(cmd, **kwargs):
|
||||
n = len(calls)
|
||||
rc = int(sequence[min(n, len(sequence) - 1)])
|
||||
calls.append({"cmd": cmd, "kwargs": kwargs})
|
||||
handle = kwargs.get("stdout")
|
||||
if handle is not None and kwargs.get("text"):
|
||||
if n == 1: # 第二次调用(重试)时写入 pipeline 行
|
||||
handle.write(f"[fake] {MINERU_PIPELINE_LOG_MARKER} started\n")
|
||||
handle.write(f"[fake stdout call#{n}] rc={rc}\n")
|
||||
handle.flush()
|
||||
return FakeResult(returncode=rc)
|
||||
|
||||
monkeypatch.setenv("METHODOLOGY_MINERU_GPU_PREFLIGHT", "0")
|
||||
monkeypatch.setattr(subprocess, "run", run_with_marker)
|
||||
monkeypatch.setattr(mr, "_preflight_executable", lambda cmd, hint, name: None)
|
||||
monkeypatch.setattr(
|
||||
mr,
|
||||
"summarize_mineru_failure",
|
||||
lambda log_path, rc: f"CUDA out of memory fake path rc={rc}",
|
||||
)
|
||||
job_dir = job_root / "testjob"
|
||||
run_uploaded_methodology_job(
|
||||
job_id="testjob",
|
||||
source_path=pdf_source,
|
||||
paper_id="paper1",
|
||||
family="",
|
||||
output_dir=tmp_path / "out",
|
||||
llm_model="deepseek-v4-flash",
|
||||
llm_request_timeout=75,
|
||||
sandbox=False,
|
||||
sandbox_run_id="",
|
||||
)
|
||||
log_text = (job_dir / "job.log").read_text(encoding="utf-8")
|
||||
assert MINERU_PIPELINE_LOG_MARKER in log_text
|
||||
assert log_text.count("[retry]") == 1, "有 pipeline 行后不得再次重试"
|
||||
# 中途失败(有 marker)→ 原错误路径(CUDA 中文错误)
|
||||
status = read_status(job_dir)
|
||||
assert status.get("status") == "failed"
|
||||
assert "CUDA" in str(status.get("message") or "")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 修复点 4:有 pipeline 行的中途失败不重试
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_mid_failure_with_pipeline_line_no_retry(monkeypatch, tmp_path, pdf_source, job_root):
|
||||
sequence = [1]
|
||||
calls = []
|
||||
|
||||
def run_with_marker(cmd, **kwargs):
|
||||
n = len(calls)
|
||||
rc = int(sequence[min(n, len(sequence) - 1)])
|
||||
calls.append({"cmd": cmd, "kwargs": kwargs})
|
||||
handle = kwargs.get("stdout")
|
||||
if handle is not None and kwargs.get("text"):
|
||||
if n == 0: # 第一次调用即写 pipeline 行
|
||||
handle.write(f"[fake] {MINERU_PIPELINE_LOG_MARKER} started\n")
|
||||
handle.write(f"[fake stdout call#{n}] rc={rc}\n")
|
||||
handle.flush()
|
||||
return FakeResult(returncode=rc)
|
||||
|
||||
monkeypatch.setenv("METHODOLOGY_MINERU_GPU_PREFLIGHT", "0")
|
||||
monkeypatch.setattr(subprocess, "run", run_with_marker)
|
||||
monkeypatch.setattr(mr, "_preflight_executable", lambda cmd, hint, name: None)
|
||||
monkeypatch.setattr(
|
||||
mr,
|
||||
"summarize_mineru_failure",
|
||||
lambda log_path, rc: f"CUDA out of memory fake path rc={rc}",
|
||||
)
|
||||
job_dir = job_root / "testjob"
|
||||
run_uploaded_methodology_job(
|
||||
job_id="testjob",
|
||||
source_path=pdf_source,
|
||||
paper_id="paper1",
|
||||
family="",
|
||||
output_dir=tmp_path / "out",
|
||||
llm_model="deepseek-v4-flash",
|
||||
llm_request_timeout=75,
|
||||
sandbox=False,
|
||||
sandbox_run_id="",
|
||||
)
|
||||
log_text = (job_dir / "job.log").read_text(encoding="utf-8")
|
||||
assert MINERU_PIPELINE_LOG_MARKER in log_text
|
||||
assert "[retry]" not in log_text, "已有 pipeline 行的中途失败不得触发重试"
|
||||
# 原始启动块 = 1 条 `$ env ...` + 1 条 `$ <cmd>`;无重试追加
|
||||
assert len([ln for ln in log_text.splitlines() if ln.startswith("$ /")]) == 1
|
||||
status = read_status(job_dir)
|
||||
assert status.get("status") == "failed"
|
||||
assert "CUDA" in str(status.get("message") or "")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# JSON 直接投递(跳过 mineru 段)不重试、不破坏状态机
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_json_source_skips_mineru_no_retry(monkeypatch, tmp_path, json_source, fake_success, job_root):
|
||||
runner = RunnerRecorder(sequence=[0]) # 仅 LLM 批次一次
|
||||
job_dir = run_job(monkeypatch, tmp_path, runner, json_source, fake_success, job_root)
|
||||
log_text = (job_dir / "job.log").read_text(encoding="utf-8")
|
||||
assert "[retry]" not in log_text
|
||||
assert MINERU_PIPELINE_LOG_MARKER not in log_text
|
||||
status = read_status(job_dir)
|
||||
assert status.get("status") == "completed"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 工具函数单元
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_job_log_has_pipeline_line(tmp_path):
|
||||
log_path = tmp_path / "job.log"
|
||||
assert job_log_has_pipeline_line(log_path) is False
|
||||
log_path.write_text("noise\n", encoding="utf-8")
|
||||
assert job_log_has_pipeline_line(log_path) is False
|
||||
log_path.write_text(f"2026-09-11 [INFO] {MINERU_PIPELINE_LOG_MARKER} begin\n", encoding="utf-8")
|
||||
assert job_log_has_pipeline_line(log_path) is True
|
||||
|
||||
|
||||
def test_summarize_startup_failure_readable(tmp_path):
|
||||
log_path = tmp_path / "job.log"
|
||||
log_path.write_text("line one\nline two\nTraceback: boom\n", encoding="utf-8")
|
||||
msg = summarize_mineru_startup_failure(log_path, 1, retried=True)
|
||||
assert "提取启动阶段失败" in msg
|
||||
assert "重试后" in msg
|
||||
assert "退出码 1" in msg
|
||||
assert "Traceback: boom" in msg
|
||||
missing = summarize_mineru_startup_failure(tmp_path / "nope.log", 1, retried=False)
|
||||
assert "提取启动阶段失败" in missing
|
||||
assert "重试后" not in missing
|
||||
@@ -16,3 +16,70 @@ capabilities:
|
||||
endpoint: https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi
|
||||
method: GET
|
||||
timeout: 8
|
||||
|
||||
# ---- 生信能力单元(S1):两个 MCP 端点,各自独立装卸 ----
|
||||
- name: bioflow-gene-family
|
||||
provider: mcp
|
||||
permission: read_only
|
||||
description: 基因家族分析(路径规划/建任务/参数确认/文件获取与预处理/状态查询),自然语言触发
|
||||
session_default: true
|
||||
nl_triggers:
|
||||
- 基因家族
|
||||
- 家族成员
|
||||
- 亚家族
|
||||
- 系统发育
|
||||
- 建树
|
||||
- 进化树
|
||||
- 保守基序
|
||||
- motif
|
||||
- 结构域
|
||||
- 基因结构
|
||||
- 共线性
|
||||
- Ka/Ks
|
||||
- 选择压力
|
||||
- 组合证据
|
||||
mcp:
|
||||
url: http://127.0.0.1:8098/mcp/gene-family/
|
||||
prompt_pack:
|
||||
file: bioflow_gene_family_agent_rules.md
|
||||
|
||||
- name: bioflow-rnaseq
|
||||
provider: mcp
|
||||
permission: read_only
|
||||
description: 转录组分析(实验设计确认/样本配置/路径规划/建任务/差异表达/富集),自然语言触发
|
||||
session_default: true
|
||||
nl_triggers:
|
||||
- 转录组
|
||||
- RNA-seq
|
||||
- RNAseq
|
||||
- rnaseq
|
||||
- 差异表达
|
||||
- 表达定量
|
||||
- 富集
|
||||
- WGCNA
|
||||
- GSEA
|
||||
- 实验设计
|
||||
- fastq
|
||||
mcp:
|
||||
url: http://127.0.0.1:8098/mcp/rnaseq/
|
||||
prompt_pack:
|
||||
file: bioflow_rnaseq_agent_rules.md
|
||||
|
||||
# ---- 生信能力单元(S4):第三个 MCP 端点,用户自定义扩展分析 ----
|
||||
- name: bioflow-extension
|
||||
provider: mcp
|
||||
permission: read_only
|
||||
description: 用户自定义扩展分析(创建/转换/试运行/校验/发布/运行),自然语言触发
|
||||
session_default: true
|
||||
nl_triggers:
|
||||
- 扩展
|
||||
- 拓展分析
|
||||
- 扩展分析
|
||||
- 自定义分析
|
||||
- 分析模板
|
||||
- 分析脚本
|
||||
- 分析工具
|
||||
mcp:
|
||||
url: http://127.0.0.1:8098/mcp/extension/
|
||||
prompt_pack:
|
||||
file: bioflow_extension_agent_rules.md
|
||||
|
||||
Vendored
+3
@@ -31,3 +31,6 @@ QA_V1_WEB_SEARCH_MAX_RESULTS=4
|
||||
QA_V1_WEB_SEARCH_TIMEOUT=8
|
||||
QA_V1_WEB_SEARCH_MAX_CHARS=1200
|
||||
QA_V1_ENABLE_SUPERPOWERS=1
|
||||
|
||||
# 生信对话式能力熔断开关(S1):默认关(=0)。置 1 才启用 QA 前门 bioflow_agent。
|
||||
QA_V1_ENABLE_BIOFLOW_AGENT=0
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
# 用户自定义扩展分析对话规范(bioflow_extension_agent_rules)
|
||||
|
||||
你是用户自定义扩展分析助手。引导用户把一段可运行的分析脚本/工具,通过「创建 → AI 转换 → 试运行 → 校验 → 确认 → 发布 → 运行」流程,做成 PlantAI 可自然语言触发的自定义扩展。
|
||||
|
||||
## 对话节奏(自然语言主导 + 卡片辅助)
|
||||
每次需要用户输入/选择时,遵循固定节奏:
|
||||
1. 先用自然语言说明并列表展示需要什么;
|
||||
2. 提示用户可用卡片快捷操作,也可直接打字(两条通道语义一致);
|
||||
3. 用户回答后,用自然语言复述,确认无误再继续。
|
||||
|
||||
## 创建(create_extension)
|
||||
- 需要用户提供:显示名、运行命令、代码文件(可上传/粘贴文本内容)、运行时镜像、可选测试数据/金标。
|
||||
- 代码文件以「路径 + 文本内容」提交;引导用户说明脚本是做什么的、依赖什么。
|
||||
|
||||
## AI 转换与校验(build → trial_run → validate)
|
||||
- 创建后调用 build 做 AI 转换(原始命令 → command_template + 参数 + 依赖),失败时规则版兜底。
|
||||
- 转换结果先向用户说明(识别到的参数/依赖/输入角色),请用户确认参数名与默认值。
|
||||
- 试运行/校验的结果要转述给用户,失败时说明原因并引导修复(repair 循环)。
|
||||
|
||||
## 确认与发布(confirm_extension → publish_extension)
|
||||
- 确认前请用户确认触发词(nl_triggers)与参数——触发词决定用户以后说什么话能唤起这个扩展。
|
||||
- 发布后扩展才可被自然语言触发运行;发布是不可逆的对外可用状态,须经用户确认。
|
||||
|
||||
## 运行(run_extension)
|
||||
- 仅已发布扩展可运行;运行参数须在白名单内(参数 schema),未知参数明确追问。
|
||||
- 运行结果(run_result)转述给用户。
|
||||
|
||||
## 其它
|
||||
- 工具名/参数名/内部实现一律不外泄;
|
||||
- 状态不确定时主动查 extension_status 并向用户转述当前阶段与下一步;
|
||||
- 人工确认边界(确认触发词、发布)不可跳过、不可代用户拍板。
|
||||
@@ -0,0 +1,32 @@
|
||||
# 基因家族分析对话规范(bioflow_gene_family_agent_rules)
|
||||
|
||||
你是基因家族分析助手。用自然语言引导用户从需求到分析完成,全程对话式交互。
|
||||
|
||||
## 人工确认边界(不可跳过、不可代确认)
|
||||
- **subfamily_confirm(亚家族预测人工确认)是硬性人工边界**:系统预测的亚家族划分结果,
|
||||
必须停下来请用户查看并确认或修改。你绝不能代用户确认,也不能在用户未确认前继续下游分析。
|
||||
- 用户提出修改(如"把某成员从亚家族 A 移到亚家族 B")时,先复述修改内容,再请用户做最终确认。
|
||||
|
||||
## 对话节奏(自然语言主导 + 卡片辅助)
|
||||
每次需要用户输入/选择时,遵循固定节奏:
|
||||
1. 先用自然语言说明并**列表展示**选项及含义;
|
||||
2. 提示用户可用卡片快捷操作,也可直接打字回复(两条通道语义一致);
|
||||
3. 用户选择后,用自然语言**复述**其选择,确认无误再继续。
|
||||
|
||||
## 路径规划与建任务
|
||||
- 规划后列出"核心流程(必选)+ 可选分析"候选路径,说明每个可选分析的生物学意义;
|
||||
- **建任务前必须经用户确认阶段链**,收到确认后再创建任务,绝不擅自建任务。
|
||||
|
||||
## 文件获取与预处理
|
||||
- 源文件支持上传、http/https 链接、OSS 地址、局域网机器路径;
|
||||
- 文件到达后先探测格式,**格式异常(如 Excel/文本)时先向用户说明列结构并提问确认**
|
||||
(哪列是序列、是否去重、是否去表头等),再由模型整理成 fasta/gff 等标准格式;
|
||||
- 转换结果过格式验证,失败时说明问题并追问修正(循环不超过 3 次后建议人工整理)。
|
||||
|
||||
## 参数确认
|
||||
- 列"参数名(默认值,含义)"请用户确认或修改;未识别参数时明确指出并给出允许范围。
|
||||
|
||||
## 其它
|
||||
- 工具名、参数名、内部实现一律不外泄;
|
||||
- 任务运行中用户提问(流程/参数/生物学意义)随时回答,不打断任务;
|
||||
- 卡住或不确定时主动查询任务状态,向用户转述当前进度与下一步。
|
||||
@@ -0,0 +1,29 @@
|
||||
# 转录组分析对话规范(bioflow_rnaseq_agent_rules)
|
||||
|
||||
你是转录组分析助手。用自然语言引导用户从实验设计确认到差异表达/富集分析完成。
|
||||
|
||||
## 人工确认边界(不可跳过、不可代确认)
|
||||
- **实验设计确认(rnaseq_de 之前的必确认信息)是硬性边界**:开始分析前必须与用户确认——
|
||||
① 物种与参考基因组版本;② 每个样本的 FASTQ(上传/链接/局域网路径)及单端双端;
|
||||
③ 对照组/处理组各几个重复、有无时间点/批次;④ 建库是否链特异性。
|
||||
- 统计模型正确性依赖上述信息,未确认前不得创建差异表达任务,也不能代用户拍板基线/比较方向。
|
||||
|
||||
## 对话节奏(自然语言主导 + 卡片辅助)
|
||||
每次需要用户输入/选择时,遵循固定节奏:
|
||||
1. 先用自然语言说明并**列表展示**必确认信息与可选分析;
|
||||
2. 提示用户可用卡片快捷操作,也可直接打字回复(两条通道语义一致);
|
||||
3. 用户回答后,用自然语言**复述**其样本设计与选择(如"干旱 3 vs 对照 3,双端,非链特异,对照为基线"),
|
||||
确认无误再继续。
|
||||
|
||||
## 样本配置与文件预处理
|
||||
- 用户用自然语言描述样本时,可解析生成 samples.yaml(样本名/对照处理组/重复/单双端);
|
||||
- FASTQ 支持上传、http/https 链接、OSS 地址、局域网路径;核对成对完整;
|
||||
- 非 FASTQ 的输入文件(表格/文本等)先探测格式,**异常时先说明列结构并提问确认**,再整理为可用格式。
|
||||
|
||||
## 参数确认
|
||||
- 列"参数名(默认值,含义)"请用户确认或修改;未识别参数时明确指出并给出允许范围。
|
||||
|
||||
## 其它
|
||||
- 工具名、参数名、内部实现一律不外泄;
|
||||
- 任务运行中用户提问(流程/参数/实验设计意义)随时回答,不打断任务;
|
||||
- 卡住或不确定时主动查询任务状态,向用户转述当前进度与下一步(含样本质控 PCA、差异基因数目等结果转述)。
|
||||
@@ -24,6 +24,26 @@ set -a
|
||||
. "$ENV_FILE"
|
||||
set +a
|
||||
|
||||
if [[ -f "$AGENT_ROOT/env/active_library.env" ]]; then
|
||||
set -a; . "$AGENT_ROOT/env/active_library.env"; set +a
|
||||
fi
|
||||
|
||||
|
||||
# ---- 开机自愈:后端校验(Qdrant/Neo4j)失败即崩且零重试,启动前必须等依赖就绪(2026-09-23) ----
|
||||
QD_READYZ="${QDRANT_URL:-http://127.0.0.1:27433}/readyz"
|
||||
QD_DEADLINE=$(( $(date +%s) + 1800 ))
|
||||
until curl -sf --max-time 3 "$QD_READYZ" >/dev/null 2>&1; do
|
||||
[ "$(date +%s)" -ge "$QD_DEADLINE" ] && { echo "[qa] Qdrant 30 分钟未就绪,放弃启动" >&2; exit 1; }
|
||||
sleep 5
|
||||
done
|
||||
N4_HP="${NEO4J_URI#bolt://}"; N4_H="${N4_HP%:*}"; N4_P="${N4_HP##*:}"
|
||||
N4_DEADLINE=$(( $(date +%s) + 600 ))
|
||||
until timeout 3 bash -c "echo > /dev/tcp/$N4_H/$N4_P" 2>/dev/null; do
|
||||
[ "$(date +%s)" -ge "$N4_DEADLINE" ] && { echo "[qa] Neo4j 10 分钟未就绪,放弃启动" >&2; exit 1; }
|
||||
sleep 5
|
||||
done
|
||||
# ---- 就绪等待结束 ----
|
||||
|
||||
# reranker 模型路径回退:env 未指定时用安装包内模型(install/models/bge-reranker-base)
|
||||
if [[ -z "${QA_V1_BGE_RERANKER_MODEL:-}" && -d "$AGENT_ROOT/../models/bge-reranker-base" ]]; then
|
||||
export QA_V1_BGE_RERANKER_MODEL="$AGENT_ROOT/../models/bge-reranker-base"
|
||||
|
||||
@@ -5,7 +5,8 @@ from dataclasses import dataclass
|
||||
from typing import Any, Iterable
|
||||
|
||||
|
||||
EXPECTED_RELEASE = "plant_srt_pc2025_v3_20260817_normalized"
|
||||
# 库名锁已 env 化:镜像文件 agent/legacy-rag/src/rag/normalized_release.py 必须逐字同步(md5 校验)。
|
||||
EXPECTED_RELEASE = os.environ.get("NORMALIZED_RELEASE_ID") or "plant_srt_pc2025_v3_20260817_normalized"
|
||||
MAX_SEEDS = 32
|
||||
MAX_RESULTS = 200
|
||||
|
||||
@@ -37,12 +38,8 @@ class NormalizedReleaseSettings:
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> "NormalizedReleaseSettings":
|
||||
release_id = os.environ.get("NORMALIZED_RELEASE_ID", "").strip()
|
||||
collection = os.environ.get("V15_QDRANT_COLLECTION", "").strip()
|
||||
if release_id != EXPECTED_RELEASE:
|
||||
raise RuntimeError(f"release mismatch: {release_id!r}")
|
||||
if collection != EXPECTED_RELEASE:
|
||||
raise RuntimeError(f"collection mismatch: {collection!r}")
|
||||
release_id = os.environ.get("NORMALIZED_RELEASE_ID", "").strip() or EXPECTED_RELEASE
|
||||
collection = os.environ.get("V15_QDRANT_COLLECTION", "").strip() or EXPECTED_RELEASE
|
||||
return cls(release_id=release_id, qdrant_collection=collection)
|
||||
|
||||
|
||||
|
||||
@@ -209,7 +209,7 @@ If the provided data are insufficient, say so directly.
|
||||
|
||||
--- Figure Marker ---
|
||||
|
||||
When citing a paper for the first time, you may insert a marker [[FIG:PMCID]] on its own line immediately after the paragraph containing the citation, to indicate where a figure from that paper should be displayed. Use at most one marker per paper, at most 6 markers total. Do not write image URLs yourself; the system will resolve markers into figures.
|
||||
When citing a paper for the first time, you may insert a marker on its own line immediately after the paragraph containing the citation, to indicate where a figure from that paper should be displayed. Use at most one marker per paper, at most 6 markers total. Prefer [[FIG:PMCID]]; when the paper has no PMCID, use [[FIG:DOI:完整DOI]] (the full DOI string, e.g. [[FIG:DOI:10.1111/j.1365-2109.2010.02587.x]]). Do not write image URLs yourself; the system will resolve markers into figures.
|
||||
|
||||
{correction_block}{dialogue_block}{web_block}Now answer the following question:
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@ from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from . import __version__
|
||||
from .bioflow_agent import try_bioflow_agent
|
||||
from .config import QAConfig
|
||||
from .pipeline import QAAnswerError, QAV1Pipeline
|
||||
|
||||
@@ -47,6 +48,9 @@ class QARequest(BaseModel):
|
||||
# ---- Capability Gateway(S1):纯数据字段,无 validator(Cython 编译安全)----
|
||||
history: List[Dict[str, str]] = Field(default_factory=list)
|
||||
capabilities: Dict[str, Any] = Field(default_factory=dict)
|
||||
# ---- Bioflow 对话式编排(S2):卡片回填协议 + session↔job 绑定 ----
|
||||
card_result: Optional[Dict[str, Any]] = None
|
||||
bioflow: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
@@ -150,6 +154,18 @@ async def chat(request: QARequest) -> Dict[str, Any]:
|
||||
detail={"error_code": "llm_unavailable", "answer": "", "message": "Final-answer LLM is not configured"},
|
||||
)
|
||||
try:
|
||||
# ---- Bioflow Agent 前门(S1):命中则短路返回,否则走原 pipeline ----
|
||||
bioflow_result = await asyncio.to_thread(
|
||||
try_bioflow_agent,
|
||||
request.question,
|
||||
llm_client=pipeline.llm,
|
||||
history=request.history,
|
||||
capabilities=request.capabilities,
|
||||
card_result=request.card_result,
|
||||
bioflow=request.bioflow,
|
||||
)
|
||||
if bioflow_result is not None:
|
||||
return bioflow_result
|
||||
ablation = dict(request.ablation or {})
|
||||
ablation["use_final_llm"] = True
|
||||
return await asyncio.to_thread(
|
||||
@@ -216,6 +232,19 @@ async def chat_stream(request: QARequest) -> StreamingResponse:
|
||||
|
||||
def run() -> None:
|
||||
try:
|
||||
# ---- Bioflow Agent 前门(S1):命中则短路,否则走原 pipeline ----
|
||||
bioflow_result = try_bioflow_agent(
|
||||
request.question,
|
||||
llm_client=pipeline.llm,
|
||||
history=request.history,
|
||||
capabilities=request.capabilities,
|
||||
card_result=request.card_result,
|
||||
bioflow=request.bioflow,
|
||||
progress_callback=emit,
|
||||
)
|
||||
if bioflow_result is not None:
|
||||
events.put(("complete", bioflow_result))
|
||||
return
|
||||
ablation = dict(request.ablation or {})
|
||||
ablation["use_final_llm"] = True
|
||||
result = pipeline.answer(
|
||||
|
||||
@@ -71,18 +71,8 @@ def _qdrant_sample(url: str, collection: str) -> Dict[str, Any]:
|
||||
|
||||
|
||||
def _validate_normalized_backends() -> Dict[str, Any]:
|
||||
release_id = os.environ.get("NORMALIZED_RELEASE_ID", "").strip()
|
||||
if release_id != NORMALIZED_RELEASE:
|
||||
raise BackendIdentityError(
|
||||
f"normalized release mismatch: expected {NORMALIZED_RELEASE!r}, "
|
||||
f"got {release_id!r}"
|
||||
)
|
||||
collection = os.environ.get("V15_QDRANT_COLLECTION", "").strip()
|
||||
if collection != NORMALIZED_RELEASE:
|
||||
raise BackendIdentityError(
|
||||
f"normalized collection mismatch: expected {NORMALIZED_RELEASE!r}, "
|
||||
f"got {collection!r}"
|
||||
)
|
||||
release_id = os.environ.get("NORMALIZED_RELEASE_ID", "").strip() or NORMALIZED_RELEASE
|
||||
collection = os.environ.get("V15_QDRANT_COLLECTION", "").strip() or NORMALIZED_RELEASE
|
||||
qdrant_url = os.environ.get(
|
||||
"V15_QDRANT_URL", os.environ.get("QDRANT_URL", "http://127.0.0.1:6333")
|
||||
).rstrip("/")
|
||||
@@ -141,7 +131,7 @@ def _validate_normalized_backends() -> Dict[str, Any]:
|
||||
finally:
|
||||
driver.close()
|
||||
|
||||
if labels != {"Entity", "Paper"}:
|
||||
if not {"Entity", "Paper"}.issubset(labels):
|
||||
raise BackendIdentityError(
|
||||
f"normalized Neo4j labels mismatch: {sorted(labels)}"
|
||||
)
|
||||
|
||||
@@ -0,0 +1,764 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""bioflow_agent.py —— 生信 MCP 能力前门路由(S1 + S2)。
|
||||
|
||||
S1:三重门(会话开关/注册表/意图预筛)→ ToolRunner 工具循环 → 自然语言转述。
|
||||
S2:对话式编排闭环——card_result 卡片回填协议 + session↔job 绑定 + 多轮记忆
|
||||
+ 参数抽取闭环 + job_status 主动查询转述。
|
||||
|
||||
card_result 协议(前端卡片操作 → 结构化用户消息 → QA 流):
|
||||
{card: path_selection|form_confirm|file_upload|archive_choice, payload: {...}}
|
||||
BioflowAgent 读取后经工具执行,再自然语言复述;卡片与自然语言两条通道在
|
||||
服务端语义一致(QA 对话状态机是唯一决策点)。
|
||||
|
||||
session↔job 绑定:BFF sessions.meta.bioflow = {active_job_id, chain},经请求
|
||||
body["bioflow"] 注入(含 session_id);job_status 工具是任务真相唯一权威。
|
||||
|
||||
金标路径零改动:任何不命中/异常都返回 None 回退原 pipeline。
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from .capabilities.manifests import CapabilityManifest
|
||||
from .capabilities.providers import PromptPackProvider, get_dispatcher
|
||||
from .capabilities.registry import get_registry
|
||||
from .tool_runner import ToolRunner, make_llm_complete
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
|
||||
# 三个生信能力单元(与 bioflow_service 三个 MCP 端点一一对应)
|
||||
BIOFLOW_CAPABILITY_NAMES = ("bioflow-gene-family", "bioflow-rnaseq", "bioflow-extension")
|
||||
ENV_FUSE = "QA_V1_ENABLE_BIOFLOW_AGENT"
|
||||
MAX_ROUNDS = 5
|
||||
|
||||
# card_result 卡片类型(S2 前端卡片回填协议)
|
||||
CARD_PATH_SELECTION = "path_selection"
|
||||
CARD_FORM_CONFIRM = "form_confirm"
|
||||
CARD_FILE_UPLOAD = "file_upload"
|
||||
CARD_ARCHIVE_CHOICE = "archive_choice"
|
||||
CARD_EXTENSION_ACTION = "extension_action" # S4:扩展能力单元操作(build/trial/validate/confirm/publish/run)
|
||||
CARD_TYPES = (CARD_PATH_SELECTION, CARD_FORM_CONFIRM, CARD_FILE_UPLOAD, CARD_ARCHIVE_CHOICE, CARD_EXTENSION_ACTION)
|
||||
|
||||
# chain 名 → 能力名(建任务时选择哪条链的 create 工具)
|
||||
_CHAIN_TO_CAPABILITY = {
|
||||
"gene_family": "bioflow-gene-family",
|
||||
"rnaseq": "bioflow-rnaseq",
|
||||
}
|
||||
|
||||
# 能力名 → chain 名(反向映射,供规划阶段从命中能力回填 chain,续接多轮对话)
|
||||
_CAPABILITY_TO_CHAIN = {v: k for k, v in _CHAIN_TO_CAPABILITY.items()}
|
||||
|
||||
|
||||
def _chain_to_capability(chain: str) -> Optional[str]:
|
||||
"""把 chain 名(gene_family/rnaseq)映射为能力名;未知返回 None。"""
|
||||
return _CHAIN_TO_CAPABILITY.get(str(chain or "").strip().lower())
|
||||
|
||||
_SYSTEM_BASE = (
|
||||
"你是 PlantAI 生信分析助手,用自然语言引导用户完成基因家族或转录组分析。\n"
|
||||
"你可以调用生信工具完成规划、建任务、确认参数、登记文件、查状态等操作。\n"
|
||||
"硬性规则:\n"
|
||||
"1. 回复只使用自然语言,绝不向用户泄露工具名、参数名或内部实现细节。\n"
|
||||
"2. 需要用户提供文件/参数/做选择时,先自然语言说明并列表展示选项与含义,再请用户回答,"
|
||||
"绝不擅自替用户做不可逆决策。\n"
|
||||
"3. 人工确认边界(基因家族 subfamily_confirm、转录组实验设计确认)不可跳过、不可代用户确认。\n"
|
||||
"4. 任务运行中可随时回答流程/参数/生物学意义类问题;卡住或不确定时主动查询任务状态,"
|
||||
"再向用户转述当前进度与下一步(等待文件、等待参数确认、失败可重跑等)。\n"
|
||||
"5. 文件格式异常时先向用户说明并提问确认,再由用户决定如何处理。\n"
|
||||
"6. 用户用自然语言给出参数(如「用 8 核、padj 用 0.01」)时,抽取为结构化参数后经确认步骤提交;"
|
||||
"只提交白名单内的参数名,未知参数明确指出并追问(最多 3 次)。\n"
|
||||
"7. 若用户消息是对上一轮卡片/提问的确认或选择,先复述用户的选择,确认无误再执行对应动作。\n"
|
||||
)
|
||||
|
||||
|
||||
def _env_enabled() -> bool:
|
||||
"""门 1:QA_V1_ENABLE_BIOFLOW_AGENT 熔断开关。默认关(未设置 = 关)。"""
|
||||
raw = os.environ.get(ENV_FUSE)
|
||||
return raw is not None and str(raw).strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def _intent_match(question: str, triggers: List[str]) -> bool:
|
||||
"""门 3:意图预筛——问句是否命中任一触发词(大小写不敏感子串匹配)。"""
|
||||
if not triggers:
|
||||
return False
|
||||
q = (question or "").lower()
|
||||
return any(t and t.lower() in q for t in triggers)
|
||||
|
||||
|
||||
def _registered_capabilities(capabilities: Optional[Dict[str, Any]]) -> List[CapabilityManifest]:
|
||||
"""注册表 + 会话开关:返回所有已注册且会话开启的生信能力 manifest(不判意图)。"""
|
||||
try:
|
||||
registry = get_registry()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
LOGGER.warning("bioflow agent: registry 访问失败,回退原 pipeline: %s", exc)
|
||||
return []
|
||||
session_caps = capabilities or {}
|
||||
out: List[CapabilityManifest] = []
|
||||
for name in BIOFLOW_CAPABILITY_NAMES:
|
||||
try:
|
||||
spec = registry.get(name)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
LOGGER.warning("bioflow agent: registry.get(%s) 失败: %s", name, exc)
|
||||
continue
|
||||
if spec is None:
|
||||
continue # 未注册
|
||||
manifest = spec.manifest
|
||||
if not bool(session_caps.get(name, manifest.session_default)):
|
||||
continue # 会话开关关
|
||||
out.append(manifest)
|
||||
return out
|
||||
|
||||
|
||||
def _candidate_capabilities(
|
||||
question: str,
|
||||
capabilities: Optional[Dict[str, Any]],
|
||||
) -> List[CapabilityManifest]:
|
||||
"""门 2 + 门 3:注册表 + 会话开关 + 意图预筛,返回命中的候选能力清单。"""
|
||||
return [
|
||||
m for m in _registered_capabilities(capabilities)
|
||||
if _intent_match(question, m.nl_triggers)
|
||||
]
|
||||
|
||||
|
||||
def _build_tools(
|
||||
manifests: List[CapabilityManifest],
|
||||
) -> Tuple[List[Dict[str, Any]], Dict[str, CapabilityManifest]]:
|
||||
"""经 MCP list_tools 发现各能力工具,构造 LLM 工具表 + ToolRunner manifests。
|
||||
|
||||
返回 (tools, tool_manifests):
|
||||
- tools: OpenAI function schema 列表(真实 server 工具名/描述/入参 schema)
|
||||
- tool_manifests: {真实工具名: manifest},manifest.name=能力名(共享连接),
|
||||
mcp.tool=真实工具名
|
||||
|
||||
多能力同时命中时:专有工具名本就不同,各自保留;共享工具(两端点语义一致)
|
||||
按名去重,只保留首个能力的一份,避免 LLM tools 重名 + 同名覆盖。
|
||||
"""
|
||||
dispatcher = get_dispatcher()
|
||||
tools: List[Dict[str, Any]] = []
|
||||
tool_manifests: Dict[str, CapabilityManifest] = {}
|
||||
for cap_manifest in manifests:
|
||||
url = str((cap_manifest.config.get("mcp") or {}).get("url") or "").strip()
|
||||
if not url:
|
||||
LOGGER.warning("bioflow agent: 能力 %s 缺 mcp.url,跳过", cap_manifest.name)
|
||||
continue
|
||||
try:
|
||||
discovered = dispatcher.list_mcp_tools(cap_manifest)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
LOGGER.warning("bioflow agent: 能力 %s MCP 工具发现失败: %s", cap_manifest.name, exc)
|
||||
continue
|
||||
for item in discovered:
|
||||
tool_name = str(item.get("name") or "").strip()
|
||||
if not tool_name:
|
||||
continue
|
||||
if tool_name in tool_manifests:
|
||||
# 共享工具去重(见 S1 修复说明)
|
||||
continue
|
||||
input_schema = item.get("input_schema") or {}
|
||||
if not isinstance(input_schema, dict):
|
||||
input_schema = {"type": "object", "properties": {}}
|
||||
tools.append(
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool_name,
|
||||
"description": str(item.get("description") or "")[:1000],
|
||||
"parameters": input_schema,
|
||||
},
|
||||
}
|
||||
)
|
||||
tool_manifests[tool_name] = CapabilityManifest(
|
||||
name=cap_manifest.name,
|
||||
provider="mcp",
|
||||
permission="read_only",
|
||||
description=str(item.get("description") or ""),
|
||||
config={"mcp": {"url": url, "tool": tool_name}},
|
||||
)
|
||||
return tools, tool_manifests
|
||||
|
||||
|
||||
def _build_system(manifests: List[CapabilityManifest]) -> str:
|
||||
"""基础系统提示 + 各能力对话规范(prompt_pack 注入)。"""
|
||||
qa_root = Path(os.environ.get("GENE_FAMILY_ROOT", "/data/gene-family-mas-v2"))
|
||||
parts = [_SYSTEM_BASE]
|
||||
seen: set[str] = set()
|
||||
for manifest in manifests:
|
||||
file_name = str((manifest.config.get("prompt_pack") or {}).get("file") or "").strip()
|
||||
if not file_name or file_name in seen:
|
||||
continue
|
||||
seen.add(file_name)
|
||||
try:
|
||||
text = PromptPackProvider.load(manifest, qa_root)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
LOGGER.warning("bioflow agent: prompt_pack 加载失败 %s: %s", file_name, exc)
|
||||
text = ""
|
||||
if text:
|
||||
parts.append(text)
|
||||
return "\n\n".join(parts)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# S2 工具执行 / 转述 / 参数抽取 辅助
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _execute_tool(
|
||||
tool_manifests: Dict[str, CapabilityManifest], tool_name: str, arguments: Dict[str, Any]
|
||||
) -> str:
|
||||
"""执行单个工具,返回结果文本(异常转文本,不炸主流程)。"""
|
||||
manifest = tool_manifests.get(tool_name)
|
||||
if manifest is None:
|
||||
return f"工具 {tool_name} 未注册"
|
||||
try:
|
||||
return get_dispatcher().execute(manifest, arguments)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return f"工具执行失败: {type(exc).__name__}: {exc}"
|
||||
|
||||
|
||||
def _paraphrase(llm_client: Any, user_text: str, system: Optional[str] = None) -> str:
|
||||
"""把工具执行结果转述为自然语言(LLM 优先,失败回退原文,绝不抛异常)。"""
|
||||
system = system or (
|
||||
"你是 PlantAI 生信分析助手。把下面的工具执行结果转述为对用户的自然语言回复,"
|
||||
"绝不提及工具名/参数名/内部实现;需要用户下一步操作时明确说明。"
|
||||
)
|
||||
try:
|
||||
text = llm_client.complete(system, user_text, temperature=0.2, max_tokens=600)
|
||||
if text and str(text).strip():
|
||||
return str(text).strip()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
LOGGER.warning("bioflow agent: paraphrase LLM 失败,回退原文: %s", exc)
|
||||
return str(user_text or "").strip()
|
||||
|
||||
|
||||
def _extract_params(
|
||||
llm_client: Any, user_text: str, form_defaults: Dict[str, Any]
|
||||
) -> Tuple[Dict[str, Any], List[str]]:
|
||||
"""LLM 抽取用户自然语言中的阶段确认参数 → (已识别参数, 未识别参数列表)。
|
||||
|
||||
白名单 = form_defaults 的键;未识别参数(不在白名单)单独返回供追问(≤3 次)。
|
||||
参照 extension_llm 的「LLM 出 JSON → 白名单净化」模式,LLM 失败/坏 JSON 返回空。
|
||||
"""
|
||||
whitelist = {str(k) for k in (form_defaults or {}).keys()}
|
||||
if not whitelist:
|
||||
return {}, []
|
||||
system = (
|
||||
"你是参数抽取助手。从用户的话中抽取阶段确认参数,只输出 JSON。\n"
|
||||
'输出形态:{"params": {"<参数名>": <值>}, "unknown": ["<未识别参数名>"]}。\n'
|
||||
f"允许的参数名(白名单):{', '.join(sorted(whitelist))}。\n"
|
||||
"只抽取用户明确提到的参数;不在白名单里的参数名放进 unknown。"
|
||||
)
|
||||
try:
|
||||
raw = llm_client.complete_json(system, user_text, max_tokens=500)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
LOGGER.warning("bioflow agent: 参数抽取 LLM 失败: %s", exc)
|
||||
return {}, []
|
||||
if not isinstance(raw, dict):
|
||||
return {}, []
|
||||
params = raw.get("params") if isinstance(raw.get("params"), dict) else {}
|
||||
known = {str(k): v for k, v in params.items() if str(k) in whitelist}
|
||||
unknown_raw = raw.get("unknown") or []
|
||||
if isinstance(unknown_raw, str):
|
||||
unknown_raw = [unknown_raw]
|
||||
unknown = [str(x).strip() for x in unknown_raw if isinstance(x, str) and str(x).strip()]
|
||||
return known, unknown
|
||||
|
||||
|
||||
def _try_parse_json(text: str) -> Optional[Any]:
|
||||
if not text:
|
||||
return None
|
||||
try:
|
||||
return json.loads(text)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _extract_job_id(result_text: str) -> Optional[str]:
|
||||
"""从 create_job 工具结果文本里提取 job_id(结果形如 {"ok":true,"data":{"id":...}})。"""
|
||||
data = _try_parse_json(result_text)
|
||||
if isinstance(data, dict):
|
||||
inner = data.get("data") if isinstance(data.get("data"), dict) else data
|
||||
jid = inner.get("id") or inner.get("job_id")
|
||||
if jid:
|
||||
return str(jid)
|
||||
return None
|
||||
|
||||
|
||||
def _query_job_status(job_id: str, tool_manifests: Dict[str, CapabilityManifest], llm_client: Any) -> str:
|
||||
"""主动查 job_status 并转述(任务真相唯一权威)。"""
|
||||
result = _execute_tool(tool_manifests, "job_status", {"job_id": job_id})
|
||||
return _paraphrase(llm_client, f"查询任务状态的结果:\n{result}")
|
||||
|
||||
|
||||
def _empty_audit() -> Dict[str, Any]:
|
||||
return {"rounds": 0, "tool_calls": 0, "used_json_fallback": False, "tool_names": [], "errors": []}
|
||||
|
||||
|
||||
def _map_readiness(data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""check_readiness 结果 → 就绪卡字段 {ready, quadrants:[{name,ok}]}。"""
|
||||
inner = data.get("result", data) if isinstance(data, dict) else {}
|
||||
quadrants = []
|
||||
for key in ("environment", "reference", "space", "budget"):
|
||||
q = inner.get(key) or {}
|
||||
quadrants.append({"name": key, "ok": bool(q.get("ok"))})
|
||||
return {"ready": bool(inner.get("ready")), "quadrants": quadrants}
|
||||
|
||||
|
||||
def _map_versioning(data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""propose_change 结果 → 版本化卡字段 {plan_revision, run_id, affected_stages, based_on_revision}。"""
|
||||
inner = data.get("result", data) if isinstance(data, dict) else {}
|
||||
affected = inner.get("affected") or []
|
||||
return {
|
||||
"plan_revision": "",
|
||||
"run_id": "",
|
||||
"affected_stages": list(affected),
|
||||
"based_on_revision": None,
|
||||
}
|
||||
|
||||
|
||||
def _make_payload(
|
||||
question: str,
|
||||
answer: str,
|
||||
*,
|
||||
audit: Optional[Dict[str, Any]] = None,
|
||||
capabilities: Optional[List[str]] = None,
|
||||
active_job_id: Optional[str] = None,
|
||||
chain: Optional[str] = None,
|
||||
readiness: Optional[Dict[str, Any]] = None,
|
||||
versioning: Optional[Dict[str, Any]] = None,
|
||||
plan: Optional[Dict[str, Any]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
payload: Dict[str, Any] = {
|
||||
"success": True,
|
||||
"version": "bioflow-agent-s2",
|
||||
"question": question,
|
||||
"answer": answer,
|
||||
"sources": [],
|
||||
"source_count": 0,
|
||||
"citation_validation": {"valid": True, "cited_pmids": [], "unknown_pmids": []},
|
||||
"bioflow": {
|
||||
"handled": True,
|
||||
"audit": audit or _empty_audit(),
|
||||
"capabilities": capabilities or [],
|
||||
},
|
||||
}
|
||||
if active_job_id:
|
||||
payload["bioflow"]["active_job_id"] = str(active_job_id)
|
||||
if chain:
|
||||
payload["bioflow"]["chain"] = str(chain)
|
||||
if readiness:
|
||||
payload["bioflow"]["readiness"] = readiness
|
||||
if versioning:
|
||||
payload["bioflow"]["versioning"] = versioning
|
||||
if plan:
|
||||
payload["bioflow"]["plan"] = plan
|
||||
return payload
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# card_result 回填处理(S2)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _card_path_selection(
|
||||
payload: Dict[str, Any], llm_client: Any, bioflow: Dict[str, Any], tool_manifests: Dict[str, CapabilityManifest]
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""路径卡确认 → create_*_job → 复述;带 reuse_archive 时先复用存档的阶段链/参数再建任务。"""
|
||||
chain = str(payload.get("chain") or (bioflow or {}).get("chain") or "").strip()
|
||||
stages = payload.get("stages")
|
||||
reuse_name = str(payload.get("reuse_archive") or payload.get("reuse_name") or "").strip()
|
||||
has_stages = isinstance(stages, list) and len(stages) > 0
|
||||
if not has_stages and not reuse_name:
|
||||
# 既无显式阶段链也无存档复用来源 → 无法建任务,回退原 pipeline
|
||||
return None
|
||||
cap_name = _chain_to_capability(chain)
|
||||
if cap_name is None:
|
||||
# 无法确定能力(chain 空/未知):不能默认走某条链,回退原 pipeline
|
||||
return None
|
||||
tool = "create_gene_family_job" if cap_name == "bioflow-gene-family" else "create_rnaseq_job"
|
||||
if tool not in tool_manifests:
|
||||
return None
|
||||
args: Dict[str, Any] = {
|
||||
"session_id": str((bioflow or {}).get("session_id") or ""),
|
||||
"stages": [str(s) for s in stages] if isinstance(stages, list) and stages else [],
|
||||
"title": str(payload.get("title") or ""),
|
||||
"auto_continue": bool(payload.get("auto_continue", False)),
|
||||
}
|
||||
tool_names: List[str] = []
|
||||
# 存档复用(FR-BIO-023):先 reuse_archive 取阶段链/参数/archive_id,再以其发起
|
||||
if reuse_name and "reuse_archive" in tool_manifests:
|
||||
tool_names.append("reuse_archive") # 实际先执行 reuse_archive(audit 顺序与实际一致)
|
||||
reuse_text = _execute_tool(tool_manifests, "reuse_archive", {"name": reuse_name})
|
||||
arch = _try_parse_json(reuse_text)
|
||||
if isinstance(arch, dict) and arch.get("ok"):
|
||||
data = arch.get("data") if isinstance(arch.get("data"), dict) else {}
|
||||
if data.get("id"):
|
||||
args["archive_id"] = str(data["id"])
|
||||
if not args["stages"] and isinstance(data.get("stage_chain"), list):
|
||||
args["stages"] = [str(s) for s in data["stage_chain"]]
|
||||
if isinstance(data.get("params"), dict):
|
||||
merged = dict(data["params"])
|
||||
if isinstance(payload.get("params"), dict):
|
||||
merged.update(payload["params"])
|
||||
args["params"] = merged
|
||||
elif not args["stages"]:
|
||||
# 复用失败(非 ok)且无显式阶段链:不能空链建任务,回退原 pipeline
|
||||
return None
|
||||
tool_names.append(tool) # create_*_job 最后执行
|
||||
if "params" not in args and isinstance(payload.get("params"), dict):
|
||||
args["params"] = payload["params"]
|
||||
result = _execute_tool(tool_manifests, tool, args)
|
||||
job_id = _extract_job_id(result)
|
||||
answer = _paraphrase(
|
||||
llm_client,
|
||||
f"用户确认了分析路径并创建任务({'基因家族' if cap_name == 'bioflow-gene-family' else '转录组'}链"
|
||||
f"{',复用存档 ' + reuse_name if reuse_name else ''})。工具结果:\n{result}",
|
||||
)
|
||||
return _make_payload(
|
||||
"确认路径并创建任务", answer,
|
||||
audit={"rounds": 1, "tool_calls": len(tool_names), "used_json_fallback": False, "tool_names": tool_names, "errors": []},
|
||||
capabilities=[cap_name] if cap_name else [],
|
||||
active_job_id=job_id,
|
||||
chain=chain or None,
|
||||
)
|
||||
|
||||
|
||||
def _card_form_confirm(
|
||||
payload: Dict[str, Any], llm_client: Any, bioflow: Dict[str, Any], tool_manifests: Dict[str, CapabilityManifest]
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""参数卡确认 → confirm_stage(白名单校验)→ 复述;rerun:true 走 rerun_stage(改参数重跑)。"""
|
||||
job_id = str(payload.get("job_id") or (bioflow or {}).get("active_job_id") or "").strip()
|
||||
stage = str(payload.get("stage") or "").strip()
|
||||
params = payload.get("params") if isinstance(payload.get("params"), dict) else {}
|
||||
rerun = bool(payload.get("rerun", False))
|
||||
if not job_id or not stage:
|
||||
return None
|
||||
tool = "rerun_stage" if rerun else "confirm_stage"
|
||||
result = _execute_tool(tool_manifests, tool, {"job_id": job_id, "stage": stage, "params": params})
|
||||
answer = _paraphrase(
|
||||
llm_client,
|
||||
f"用户{'重跑并确认' if rerun else '确认'}了阶段参数(任务 {job_id} 的 {stage})。工具结果:\n{result}",
|
||||
)
|
||||
return _make_payload(
|
||||
f"{'重跑' if rerun else '确认'}阶段参数:{stage}", answer,
|
||||
audit={"rounds": 1, "tool_calls": 1, "used_json_fallback": False, "tool_names": [tool], "errors": []},
|
||||
active_job_id=job_id,
|
||||
)
|
||||
|
||||
|
||||
def _card_file_upload(
|
||||
payload: Dict[str, Any], llm_client: Any, bioflow: Dict[str, Any], tool_manifests: Dict[str, CapabilityManifest]
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""文件卡 → 网络源 fetch_source_file / 服务器·局域网路径 register_job_files → 复述。"""
|
||||
job_id = str(payload.get("job_id") or (bioflow or {}).get("active_job_id") or "").strip()
|
||||
raw_files = payload.get("paths") or payload.get("files") or []
|
||||
if isinstance(raw_files, str):
|
||||
raw_files = [raw_files]
|
||||
files = [str(f) for f in raw_files if str(f).strip()] if isinstance(raw_files, list) else []
|
||||
if not job_id or not files:
|
||||
return None
|
||||
remote = [f for f in files if f.startswith(("http://", "https://", "oss://"))]
|
||||
local = [f for f in files if f not in remote]
|
||||
results: List[str] = []
|
||||
tool_names: List[str] = []
|
||||
for src in remote:
|
||||
r = _execute_tool(tool_manifests, "fetch_source_file", {"source": src, "job_id": job_id})
|
||||
results.append(r)
|
||||
tool_names.append("fetch_source_file")
|
||||
if local:
|
||||
r = _execute_tool(tool_manifests, "register_job_files", {"job_id": job_id, "paths": local})
|
||||
results.append(r)
|
||||
tool_names.append("register_job_files")
|
||||
answer = _paraphrase(
|
||||
llm_client,
|
||||
f"用户提供了文件({'远程 ' + str(len(remote)) + ' 个' if remote else ''}"
|
||||
f"{'本地 ' + str(len(local)) + ' 个' if local else ''})。工具结果:\n" + "\n".join(results),
|
||||
)
|
||||
return _make_payload(
|
||||
"登记任务文件", answer,
|
||||
audit={"rounds": 1, "tool_calls": len(tool_names), "used_json_fallback": False, "tool_names": tool_names, "errors": []},
|
||||
active_job_id=job_id,
|
||||
)
|
||||
|
||||
|
||||
def _card_archive_choice(
|
||||
payload: Dict[str, Any], llm_client: Any, bioflow: Dict[str, Any], tool_manifests: Dict[str, CapabilityManifest]
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""存档卡 → archive_job 工具执行实际存档(或跳过)→ 复述。"""
|
||||
job_id = str(payload.get("job_id") or (bioflow or {}).get("active_job_id") or "").strip()
|
||||
do_archive = bool(payload.get("archive", False))
|
||||
if not job_id:
|
||||
return None
|
||||
if do_archive:
|
||||
name = str(payload.get("name") or "")
|
||||
description = str(payload.get("description") or "")
|
||||
result = _execute_tool(
|
||||
tool_manifests, "archive_job",
|
||||
{"job_id": job_id, "name": name, "description": description},
|
||||
)
|
||||
answer = _paraphrase(
|
||||
llm_client,
|
||||
f"用户选择把任务 {job_id} 存档(名称:{name or '未命名'})。工具结果:\n{result}",
|
||||
)
|
||||
return _make_payload(
|
||||
"存档选择", answer,
|
||||
audit={"rounds": 1, "tool_calls": 1, "used_json_fallback": False, "tool_names": ["archive_job"], "errors": []},
|
||||
active_job_id=job_id,
|
||||
)
|
||||
answer = _paraphrase(
|
||||
llm_client,
|
||||
f"用户选择不存档任务 {job_id}。请用自然语言告知用户:已跳过存档,"
|
||||
"本次分析结果仍可在报告/结果中查看。",
|
||||
)
|
||||
return _make_payload(
|
||||
"存档选择", answer,
|
||||
audit=_empty_audit(),
|
||||
active_job_id=job_id,
|
||||
)
|
||||
|
||||
|
||||
def _card_extension_action(
|
||||
payload: Dict[str, Any], llm_client: Any, bioflow: Dict[str, Any], tool_manifests: Dict[str, CapabilityManifest]
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""扩展操作卡(S4):action ∈ build|trial_run|validate|confirm|publish|run → 对应扩展工具。"""
|
||||
ext_id = str(payload.get("ext_id") or "").strip()
|
||||
action = str(payload.get("action") or "").strip()
|
||||
if not ext_id or not action:
|
||||
return None
|
||||
tool = {
|
||||
"build": "build_extension",
|
||||
"trial_run": "trial_run_extension",
|
||||
"validate": "validate_extension",
|
||||
"confirm": "confirm_extension",
|
||||
"publish": "publish_extension",
|
||||
"run": "run_extension",
|
||||
"repair": "repair_extension",
|
||||
}.get(action)
|
||||
if tool is None or tool not in tool_manifests:
|
||||
return None
|
||||
args: Dict[str, Any] = {"ext_id": ext_id}
|
||||
if action == "confirm":
|
||||
if isinstance(payload.get("nl_triggers"), list):
|
||||
args["nl_triggers"] = payload["nl_triggers"]
|
||||
if isinstance(payload.get("parameters"), list):
|
||||
args["parameters"] = payload["parameters"]
|
||||
if payload.get("runtime_image"):
|
||||
args["runtime_image"] = str(payload["runtime_image"])
|
||||
if action == "run":
|
||||
if isinstance(payload.get("params"), dict):
|
||||
args["params"] = payload["params"]
|
||||
if isinstance(payload.get("files"), list):
|
||||
args["files"] = payload["files"]
|
||||
result = _execute_tool(tool_manifests, tool, args)
|
||||
if action == "confirm" and "EXT_STATUS" in result:
|
||||
# 状态降级话术:confirm 要求 validated/trial_passed 前置状态,
|
||||
# 编辑修正若发生在校验前,明确提示先试运行/校验,不报 500。
|
||||
answer = (
|
||||
"该扩展当前状态还不能确认。请先完成 AI 转换与试运行、通过校验后,"
|
||||
"再确认参数与触发词。"
|
||||
)
|
||||
else:
|
||||
answer = _paraphrase(
|
||||
llm_client,
|
||||
f"用户对扩展 {ext_id} 执行了操作 {action}。工具结果:\n{result}",
|
||||
)
|
||||
return _make_payload(
|
||||
f"扩展操作:{action}", answer,
|
||||
audit={"rounds": 1, "tool_calls": 1, "used_json_fallback": False, "tool_names": [tool], "errors": []},
|
||||
)
|
||||
|
||||
|
||||
def _handle_card_result(
|
||||
card_result: Dict[str, Any],
|
||||
llm_client: Any,
|
||||
bioflow: Optional[Dict[str, Any]],
|
||||
tool_manifests: Dict[str, CapabilityManifest],
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
card = str(card_result.get("card") or "").strip()
|
||||
payload = card_result.get("payload") if isinstance(card_result.get("payload"), dict) else {}
|
||||
ctx = bioflow or {}
|
||||
if card == CARD_PATH_SELECTION:
|
||||
return _card_path_selection(payload, llm_client, ctx, tool_manifests)
|
||||
if card == CARD_FORM_CONFIRM:
|
||||
return _card_form_confirm(payload, llm_client, ctx, tool_manifests)
|
||||
if card == CARD_FILE_UPLOAD:
|
||||
return _card_file_upload(payload, llm_client, ctx, tool_manifests)
|
||||
if card == CARD_ARCHIVE_CHOICE:
|
||||
return _card_archive_choice(payload, llm_client, ctx, tool_manifests)
|
||||
if card == CARD_EXTENSION_ACTION:
|
||||
return _card_extension_action(payload, llm_client, ctx, tool_manifests)
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 前门路由
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def try_bioflow_agent(
|
||||
question: str,
|
||||
*,
|
||||
llm_client: Any,
|
||||
history: Optional[List[Dict[str, str]]] = None,
|
||||
capabilities: Optional[Dict[str, Any]] = None,
|
||||
card_result: Optional[Dict[str, Any]] = None,
|
||||
bioflow: Optional[Dict[str, Any]] = None,
|
||||
progress_callback: Optional[Any] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""前门路由入口。命中生信能力则返回答案 payload,否则返回 None(走原 pipeline)。
|
||||
|
||||
- card_result:非空时走卡片回填协议(直接路由到对应工具,不做意图预筛)。
|
||||
- 任何异常都不上抛——记日志后返回 None,保证金标路径永不因本模块失败。
|
||||
"""
|
||||
try:
|
||||
return _try_bioflow_agent_impl(
|
||||
question,
|
||||
llm_client=llm_client,
|
||||
history=history,
|
||||
capabilities=capabilities,
|
||||
card_result=card_result,
|
||||
bioflow=bioflow,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
LOGGER.warning("bioflow agent: 执行失败,回退原 pipeline: %s", exc)
|
||||
return None
|
||||
|
||||
|
||||
def _try_bioflow_agent_impl(
|
||||
question: str,
|
||||
*,
|
||||
llm_client: Any,
|
||||
history: Optional[List[Dict[str, str]]],
|
||||
capabilities: Optional[Dict[str, Any]],
|
||||
card_result: Optional[Dict[str, Any]],
|
||||
bioflow: Optional[Dict[str, Any]],
|
||||
progress_callback: Optional[Any],
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
# 门 1:会话开关(熔断)
|
||||
if not _env_enabled():
|
||||
return None
|
||||
|
||||
# card_result 卡片操作优先:用户已在卡片上明确选择,不做意图预筛,直接路由
|
||||
if isinstance(card_result, dict) and str(card_result.get("card") or "").strip() in CARD_TYPES:
|
||||
return _route_card_result(card_result, llm_client, capabilities, bioflow, progress_callback)
|
||||
|
||||
# NL 路径:活跃任务/生信链续接优先,其次意图预筛
|
||||
active_job_id = (bioflow or {}).get("active_job_id")
|
||||
chain = str((bioflow or {}).get("chain") or "").strip()
|
||||
if active_job_id:
|
||||
# 有活跃任务:跳过意图预筛,按任务 chain 续接(任务链不可变)
|
||||
manifests = _registered_capabilities(capabilities)
|
||||
cap_name = _chain_to_capability(chain)
|
||||
if cap_name:
|
||||
manifests = [m for m in manifests if m.name == cap_name]
|
||||
else:
|
||||
# 无活跃任务:先门 2 + 门 3 意图预筛(新任务入口)
|
||||
manifests = _candidate_capabilities(question, capabilities)
|
||||
# 意图未命中但 session 留有生信链(上轮规划)→ 续接该链(如路径选择「1」)
|
||||
if not manifests and chain:
|
||||
cap_name = _chain_to_capability(chain)
|
||||
if cap_name:
|
||||
registered = _registered_capabilities(capabilities)
|
||||
manifests = [m for m in registered if m.name == cap_name]
|
||||
if not manifests:
|
||||
return None
|
||||
|
||||
tools, tool_manifests = _build_tools(manifests)
|
||||
if not tools:
|
||||
LOGGER.warning("bioflow agent: 意图命中但 MCP 工具发现为空,回退原 pipeline")
|
||||
return None
|
||||
|
||||
if progress_callback is not None:
|
||||
try:
|
||||
progress_callback(
|
||||
{
|
||||
"stage": "bioflow_agent_started",
|
||||
"capabilities": [m.name for m in manifests],
|
||||
"tool_count": len(tools),
|
||||
}
|
||||
)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
|
||||
system = _build_system(manifests)
|
||||
runner = ToolRunner(make_llm_complete(llm_client), max_rounds=MAX_ROUNDS)
|
||||
content, audit = runner.run(system, question, tools, tool_manifests, history=history)
|
||||
|
||||
answer = (content or "").strip()
|
||||
if not answer:
|
||||
return None
|
||||
|
||||
tool_results = (audit or {}).get("tool_results") or {}
|
||||
readiness = None
|
||||
versioning = None
|
||||
new_job_id = None
|
||||
plan = None
|
||||
for name, result_text in tool_results.items():
|
||||
if name in ("create_gene_family_job", "create_rnaseq_job"):
|
||||
jid = _extract_job_id(result_text)
|
||||
if jid:
|
||||
new_job_id = jid
|
||||
data = _try_parse_json(result_text)
|
||||
if not isinstance(data, dict):
|
||||
continue
|
||||
if name.startswith("plan_") and data.get("ok"):
|
||||
plan = {
|
||||
"candidates": data.get("candidates") or [],
|
||||
"auto_filled": data.get("auto_filled") or [],
|
||||
"reuse_archive": data.get("reuse_archive"),
|
||||
"reuse_name": data.get("reuse_name"),
|
||||
"fallback": data.get("fallback"),
|
||||
}
|
||||
if name == "check_readiness":
|
||||
readiness = _map_readiness(data)
|
||||
elif name == "propose_change":
|
||||
versioning = _map_versioning(data)
|
||||
|
||||
effective_chain = chain or (_CAPABILITY_TO_CHAIN.get(manifests[0].name) if manifests else "")
|
||||
return _make_payload(
|
||||
question,
|
||||
answer,
|
||||
audit=audit,
|
||||
capabilities=[m.name for m in manifests],
|
||||
active_job_id=(bioflow or {}).get("active_job_id") or new_job_id,
|
||||
chain=effective_chain or None,
|
||||
readiness=readiness,
|
||||
versioning=versioning,
|
||||
plan=plan,
|
||||
)
|
||||
|
||||
|
||||
def _route_card_result(
|
||||
card_result: Dict[str, Any],
|
||||
llm_client: Any,
|
||||
capabilities: Optional[Dict[str, Any]],
|
||||
bioflow: Optional[Dict[str, Any]],
|
||||
progress_callback: Optional[Any],
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""卡片操作路由:确定相关能力 → 构建工具 → 分发到对应 card handler。"""
|
||||
card = str(card_result.get("card") or "").strip()
|
||||
payload = card_result.get("payload") if isinstance(card_result.get("payload"), dict) else {}
|
||||
chain = str(payload.get("chain") or (bioflow or {}).get("chain") or "").strip()
|
||||
if card == CARD_EXTENSION_ACTION:
|
||||
# 扩展操作只依赖 bioflow-extension 能力,避免对三个能力各做一次 list_tools
|
||||
cap_name = "bioflow-extension"
|
||||
else:
|
||||
cap_name = _chain_to_capability(chain)
|
||||
registered = _registered_capabilities(capabilities)
|
||||
if cap_name:
|
||||
manifests = [m for m in registered if m.name == cap_name]
|
||||
else:
|
||||
manifests = registered
|
||||
if not manifests:
|
||||
return None
|
||||
tools, tool_manifests = _build_tools(manifests)
|
||||
if not tools:
|
||||
LOGGER.warning("bioflow agent: card_result 但 MCP 工具发现为空,回退原 pipeline")
|
||||
return None
|
||||
if progress_callback is not None:
|
||||
try:
|
||||
progress_callback({"stage": "bioflow_card", "card": card_result.get("card")})
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
return _handle_card_result(card_result, llm_client, bioflow, tool_manifests)
|
||||
@@ -128,6 +128,8 @@ class InstallManager:
|
||||
"allowed_llm_phases": manifest.allowed_llm_phases,
|
||||
"session_default": manifest.session_default,
|
||||
}
|
||||
if manifest.nl_triggers:
|
||||
raw["nl_triggers"] = manifest.nl_triggers
|
||||
if manifest.schema:
|
||||
raw["schema"] = manifest.schema
|
||||
for section in ("http", "mcp", "python", "prompt_pack"):
|
||||
|
||||
@@ -37,6 +37,7 @@ class CapabilityManifest:
|
||||
session_default: bool = False
|
||||
schema: Dict[str, Any] = field(default_factory=dict)
|
||||
config: Dict[str, Any] = field(default_factory=dict) # provider 专属段(http:/mcp:/python:/prompt_pack: 归并)
|
||||
nl_triggers: List[str] = field(default_factory=list) # 意图预筛触发词(S1,生信能力路由用)
|
||||
|
||||
def to_tool_spec(self) -> Dict[str, Any]:
|
||||
"""转 OpenAI function schema(tools 参数形态)。"""
|
||||
@@ -110,6 +111,19 @@ def _parse_entry(raw: Any, index: int, load_warnings: List[str]) -> CapabilityMa
|
||||
if isinstance(raw.get(section), dict):
|
||||
config[section] = raw[section]
|
||||
config["session_default"] = bool(raw.get("session_default", False))
|
||||
# nl_triggers(S1):意图预筛触发词,只收字符串、去空、去重、限长
|
||||
nl_triggers_raw = raw.get("nl_triggers") or []
|
||||
if isinstance(nl_triggers_raw, str):
|
||||
nl_triggers_raw = [nl_triggers_raw]
|
||||
nl_triggers: List[str] = []
|
||||
if isinstance(nl_triggers_raw, list):
|
||||
for trigger in nl_triggers_raw:
|
||||
if not isinstance(trigger, str):
|
||||
continue # 触发词只收字符串,非字符串(数字等)丢弃
|
||||
cleaned = _clean_text(trigger, 64)
|
||||
if cleaned and cleaned not in nl_triggers:
|
||||
nl_triggers.append(cleaned)
|
||||
nl_triggers = nl_triggers[:64]
|
||||
manifest = CapabilityManifest(
|
||||
name=name,
|
||||
provider=provider,
|
||||
@@ -119,6 +133,7 @@ def _parse_entry(raw: Any, index: int, load_warnings: List[str]) -> CapabilityMa
|
||||
session_default=bool(raw.get("session_default", False)),
|
||||
schema=schema,
|
||||
config=config,
|
||||
nl_triggers=nl_triggers,
|
||||
)
|
||||
# 大小上限:单条目序列化后 ≤8KB
|
||||
import json
|
||||
|
||||
@@ -130,6 +130,7 @@ class _McpConnection:
|
||||
self._thread: Optional[threading.Thread] = None
|
||||
self._session: Any = None
|
||||
self._server_tools: List[str] = []
|
||||
self._server_tool_schemas: List[Dict[str, Any]] = []
|
||||
self._started_at = 0.0
|
||||
self._last_used = 0.0
|
||||
self._restarts = 0
|
||||
@@ -200,6 +201,18 @@ class _McpConnection:
|
||||
raise
|
||||
self._ctx = ctx
|
||||
self._server_tools = [t.name for t in tools_result.tools]
|
||||
self._server_tool_schemas = [
|
||||
{
|
||||
"name": t.name,
|
||||
"description": (getattr(t, "description", None) or ""),
|
||||
"input_schema": (
|
||||
getattr(t, "inputSchema", None)
|
||||
or getattr(t, "input_schema", None)
|
||||
or {}
|
||||
),
|
||||
}
|
||||
for t in tools_result.tools
|
||||
]
|
||||
self._session = session
|
||||
self._started_at = time.time()
|
||||
|
||||
@@ -225,6 +238,7 @@ class _McpConnection:
|
||||
ctx = getattr(self, "_ctx", None)
|
||||
self._ctx = None
|
||||
self._server_tools = []
|
||||
self._server_tool_schemas = []
|
||||
if session is not None:
|
||||
try:
|
||||
await session.__aexit__(None, None, None)
|
||||
@@ -274,6 +288,14 @@ class _McpConnection:
|
||||
return []
|
||||
return list(self._server_tools)
|
||||
|
||||
def list_tool_schemas(self) -> List[Dict[str, Any]]:
|
||||
"""返回 MCP server 全量工具 schema(name/description/input_schema)。连接失败返回空。"""
|
||||
try:
|
||||
self._ensure_session()
|
||||
except ToolExecutionError:
|
||||
return []
|
||||
return list(self._server_tool_schemas)
|
||||
|
||||
def call(self, tool_name: str, arguments: Dict[str, Any]) -> str:
|
||||
self._ensure_session()
|
||||
try:
|
||||
@@ -337,6 +359,11 @@ class McpToolProvider:
|
||||
return tools + [f"<缺失声明工具 {tool_name}>"]
|
||||
return tools
|
||||
|
||||
def list_tool_schemas(self, manifest: CapabilityManifest) -> List[Dict[str, Any]]:
|
||||
"""返回该能力 MCP server 的全量工具 schema(供前门 agent 构建 LLM 工具表)。"""
|
||||
conn = self._connection(manifest)
|
||||
return conn.list_tool_schemas()
|
||||
|
||||
def reap_idle(self) -> int:
|
||||
reaped = 0
|
||||
with self._lock:
|
||||
@@ -448,6 +475,10 @@ class ToolDispatcher:
|
||||
return PythonToolProvider.execute(manifest, arguments)
|
||||
raise ToolExecutionError(f"工具 {manifest.name} provider={provider} 不支持调用")
|
||||
|
||||
def list_mcp_tools(self, manifest: CapabilityManifest) -> List[Dict[str, Any]]:
|
||||
"""MCP 能力全量工具 schema(bioflow_agent 构建 LLM 工具表用)。"""
|
||||
return self._mcp.list_tool_schemas(manifest)
|
||||
|
||||
def reap_idle(self) -> int:
|
||||
return self._mcp.reap_idle()
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@ from __future__ import annotations
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Tuple
|
||||
|
||||
@@ -16,6 +17,7 @@ from src.rag.normalized_release import (
|
||||
)
|
||||
|
||||
from .schemas import EvidenceItem, RetrievalHypothesis
|
||||
from .soft_gate_policy import doi_from_paper_id, pmcid_from_paper_id
|
||||
|
||||
ALLOWED_QDRANT_COLLECTIONS = {
|
||||
"plant_srt_direct_v1_bge_m3_20260702",
|
||||
@@ -36,6 +38,12 @@ def _norm(value: Any) -> str:
|
||||
return " ".join(text.split())
|
||||
|
||||
|
||||
# doi: 前缀形态 paper_id 的 DOI 全文校验(回改 v4,2026-09-11)
|
||||
_DOI_LITERAL = re.compile(
|
||||
r"10\.\d{4,9}/[A-Za-z0-9][A-Za-z0-9._;()/:+\-]*"
|
||||
)
|
||||
|
||||
|
||||
def load_paper_metadata_registry(path: Path) -> Dict[str, Dict[str, Any]]:
|
||||
records: Dict[str, Dict[str, Any]] = {}
|
||||
if not path.is_file():
|
||||
@@ -117,6 +125,7 @@ class DirectFlowRetriever:
|
||||
use_qdrant: bool = True,
|
||||
use_neo4j: bool = True,
|
||||
use_graph_expansion: bool = True,
|
||||
qdrant_score_threshold: float | None = None,
|
||||
) -> Tuple[List[EvidenceItem], Dict[str, Any]]:
|
||||
effective_use_neo4j = bool(use_neo4j) and self.graph_schema in {
|
||||
"canonical",
|
||||
@@ -138,6 +147,7 @@ class DirectFlowRetriever:
|
||||
"use_graph_expansion": use_graph_expansion,
|
||||
"effective_use_neo4j": effective_use_neo4j,
|
||||
"effective_use_graph_expansion": effective_use_graph_expansion,
|
||||
"qdrant_score_threshold": qdrant_score_threshold,
|
||||
},
|
||||
"neo4j_status": (
|
||||
"enabled"
|
||||
@@ -153,7 +163,12 @@ class DirectFlowRetriever:
|
||||
per_query = max(10, min(qdrant_top_k // max(1, len(all_queries)), 50))
|
||||
if self.vector_query_workers == 1 or len(all_queries) == 1:
|
||||
query_rows = [
|
||||
(query, self.qdrant_search(query, per_query))
|
||||
(
|
||||
query,
|
||||
self.qdrant_search(
|
||||
query, per_query, score_threshold=qdrant_score_threshold
|
||||
),
|
||||
)
|
||||
for query in all_queries
|
||||
]
|
||||
else:
|
||||
@@ -162,7 +177,12 @@ class DirectFlowRetriever:
|
||||
thread_name_prefix="qa-vector",
|
||||
) as executor:
|
||||
futures = [
|
||||
executor.submit(self.qdrant_search, query, per_query)
|
||||
executor.submit(
|
||||
self.qdrant_search,
|
||||
query,
|
||||
per_query,
|
||||
score_threshold=qdrant_score_threshold,
|
||||
)
|
||||
for query in all_queries
|
||||
]
|
||||
query_rows = [
|
||||
@@ -200,14 +220,20 @@ class DirectFlowRetriever:
|
||||
}
|
||||
return merged, audit
|
||||
|
||||
def qdrant_search(self, query: str, limit: int) -> List[EvidenceItem]:
|
||||
def qdrant_search(
|
||||
self, query: str, limit: int, *, score_threshold: float | None = None
|
||||
) -> List[EvidenceItem]:
|
||||
try:
|
||||
vector = self.embed(query)
|
||||
# 缺陷 #28 回改 v2(2026-09-11):向量域阈值仅 search/retrieval-only
|
||||
# 路径传入(pipeline 传 QA_MIN_VECTOR_SCORE=0.45);chat 模式 None
|
||||
# 保召回,rerank 后由 final 域 MIN_SOURCE_SCORE=0.25 兜底。
|
||||
response = self.qdrant.query_points(
|
||||
collection_name=self.collection,
|
||||
query=vector,
|
||||
limit=limit,
|
||||
with_payload=True,
|
||||
score_threshold=score_threshold,
|
||||
)
|
||||
except Exception:
|
||||
return []
|
||||
@@ -324,6 +350,20 @@ class DirectFlowRetriever:
|
||||
metadata = payload.get("metadata") if isinstance(payload.get("metadata"), dict) else {}
|
||||
paper_id = _text(payload.get("paper_id") or payload.get("best_paper_id"))
|
||||
registry_metadata = getattr(self, "paper_metadata", {}).get(paper_id, {})
|
||||
# 缺陷 #28 回改 v4(2026-09-11):175 无 paper metadata registry,
|
||||
# Qdrant payload 又无 pmid/pmcid/doi 键——chat 的 public_citation
|
||||
# 过滤会把有正文的 qdrant 条目全滤掉,只剩无 evidence 的 neo4j
|
||||
# 归一化图断言(bge 全 0.02-0.047),final 门 0.25 全拒 sources=[]。
|
||||
# 从 paper_id 派生 DOI/PMCID 回填引用字段(payload 显式字段 >
|
||||
# registry > paper_id 派生)。
|
||||
_derived_doi = doi_from_paper_id(paper_id)
|
||||
_derived_pmcid = pmcid_from_paper_id(paper_id)
|
||||
# doi: 前缀形态(如 doi:10.1002/ps.6526)直接剥前缀(DOI 已带斜杠,
|
||||
# doi_from_paper_id 不覆盖此形态)。
|
||||
if not _derived_doi and paper_id.lower().startswith("doi:"):
|
||||
_candidate = paper_id[4:].strip()
|
||||
if _DOI_LITERAL.fullmatch(_candidate):
|
||||
_derived_doi = _candidate
|
||||
return EvidenceItem(
|
||||
evidence_id="",
|
||||
kb="V16/V15",
|
||||
@@ -334,8 +374,8 @@ class DirectFlowRetriever:
|
||||
paper_id=paper_id,
|
||||
title=_text(payload.get("title") or payload.get("paper_title") or metadata.get("title") or registry_metadata.get("title")),
|
||||
pmid=_text(payload.get("pmid") or payload.get("pubmed_id") or metadata.get("pmid") or registry_metadata.get("pmid")),
|
||||
pmcid=_text(payload.get("pmcid") or metadata.get("pmcid") or registry_metadata.get("pmcid")),
|
||||
doi=_text(payload.get("doi") or metadata.get("doi") or registry_metadata.get("doi")),
|
||||
pmcid=_text(payload.get("pmcid") or metadata.get("pmcid") or registry_metadata.get("pmcid") or _derived_pmcid),
|
||||
doi=_text(payload.get("doi") or metadata.get("doi") or registry_metadata.get("doi") or _derived_doi),
|
||||
species=_text(payload.get("species") or metadata.get("species")),
|
||||
section=_text(payload.get("section") or metadata.get("section")),
|
||||
output_layer=_text(payload.get("output_layer") or "main_kg"),
|
||||
|
||||
@@ -43,6 +43,7 @@ from .question_decomposition import QuestionDecompositionPlanner
|
||||
from .query_builder import QueryBuilder
|
||||
from .reranker import BGEReranker
|
||||
from .schemas import EvidenceItem, HypothesisEntity, RetrievalHypothesis
|
||||
from .score_gate import min_source_score, min_vector_score
|
||||
from .soft_gate_policy import (
|
||||
LEGACY,
|
||||
doi_from_paper_id,
|
||||
@@ -1254,7 +1255,11 @@ class QAV1Pipeline:
|
||||
use_qdrant = bool(retrieval_ablation.get("use_qdrant", True))
|
||||
use_neo4j = bool(retrieval_ablation.get("use_neo4j", True))
|
||||
use_graph_expansion = bool(retrieval_ablation.get("use_graph_expansion", True)) and use_neo4j
|
||||
use_bge_rerank = bool(ablation.get("use_bge_rerank", True)) and self.config.enable_bge_rerank
|
||||
use_bge_rerank = (
|
||||
bool(ablation.get("use_bge_rerank", True))
|
||||
and self.config.enable_bge_rerank
|
||||
and not retrieval_only # 回改 v2:search 免 CPU rerank(性能根治)
|
||||
)
|
||||
use_final_llm = not retrieval_only
|
||||
graph_min_evidence = int(ablation.get("graph_min_evidence") or os.environ.get("QA_V115_GRAPH_MIN_EVIDENCE", "15"))
|
||||
answer_context_mode = str(ablation.get("answer_context") or "top30_plus_graph_chains")
|
||||
@@ -1285,14 +1290,22 @@ class QAV1Pipeline:
|
||||
1.0, min(float(endpoint_path_time_budget_seconds or 5.0), 30.0)
|
||||
)
|
||||
|
||||
if retrieval_only and not self.llm.available:
|
||||
# V1.1.711 supports retrieval-only operation. Without an LLM key,
|
||||
# disable provisional drafting, structured planning, and final answer
|
||||
# generation, but still return reranked top30 evidence and native
|
||||
# Neo4j graph chains from the raw question.
|
||||
if retrieval_only:
|
||||
# 缺陷 #28 回改 v2(2026-09-11):检索模式绝不产生答案——无条件
|
||||
# 关闭草稿/规划/最终 LLM(不看 llm.available;LLM 调用纯浪费且
|
||||
# 拖慢 search)。仍返回 top30 向量证据。
|
||||
use_draft = False
|
||||
use_planner = False
|
||||
use_final_llm = False
|
||||
# 缺陷 #28 回改 v3(2026-09-11):检索模式只走 Qdrant 向量——
|
||||
# r3 语义补完。图证据种子检索(use_neo4j,graph_top_k=80)在
|
||||
# 13M 边库上 150s+(TL 实测 arginine 150s+ vs 关 neo4j 后
|
||||
# 0.787s),不可接受。use_graph_expansion 依赖 use_neo4j 一并
|
||||
# 显式关(本块在所有 ablation 解析之后:检索模式用户 ablation
|
||||
# 不得重开图检索——性能红线)。native_graph_chains 空数组在
|
||||
# search 分支已安全。
|
||||
use_neo4j = False
|
||||
use_graph_expansion = False
|
||||
|
||||
t0 = time.perf_counter()
|
||||
if use_draft:
|
||||
@@ -1391,7 +1404,15 @@ class QAV1Pipeline:
|
||||
"use_neo4j": use_neo4j,
|
||||
"use_graph_expansion": use_graph_expansion,
|
||||
}
|
||||
if self.agentic_controller is not None:
|
||||
if self.agentic_controller is None:
|
||||
# 回改 v2(2026-09-11):agentic_controller 禁用(含单测)时,
|
||||
# 给确定性 plan——query_plan.agentic_plan 键不缺失。
|
||||
agentic_plan = {
|
||||
"vector_plan": {"enabled": bool(use_qdrant)},
|
||||
"graph_plan": {"enabled": bool(use_neo4j)},
|
||||
"plan_source": "deterministic_retrieval_only_plan",
|
||||
}
|
||||
elif self.agentic_controller is not None:
|
||||
try:
|
||||
agentic_plan = self.agentic_controller.plan(
|
||||
question,
|
||||
@@ -1550,6 +1571,9 @@ class QAV1Pipeline:
|
||||
hypothesis,
|
||||
qdrant_top_k=qdrant_top_k,
|
||||
graph_top_k=graph_top_k,
|
||||
qdrant_score_threshold=(
|
||||
min_vector_score() if retrieval_only else None
|
||||
),
|
||||
**direct_flow_switches,
|
||||
)
|
||||
timings["retrieval"] = time.perf_counter() - t0
|
||||
@@ -1565,6 +1589,14 @@ class QAV1Pipeline:
|
||||
clustered_input, cluster_audit = build_clusters(ranked_input)
|
||||
if use_bge_rerank:
|
||||
ranked, rerank_audit = self._rerank(question, clustered_input)
|
||||
elif retrieval_only:
|
||||
# 回改 v2(2026-09-11):search 免 BGE rerank,候选按 vector_score
|
||||
# 排序直接输出,final_score 落在 vector_score 域——与 Qdrant
|
||||
# score_threshold(min_vector_score=0.45)同域同语义。
|
||||
ranked = sorted(clustered_input, key=lambda x: -x.vector_score)
|
||||
for item in ranked:
|
||||
item.final_score = item.vector_score
|
||||
rerank_audit = [{"status": "disabled_retrieval_only"}]
|
||||
else:
|
||||
ranked = sorted(clustered_input, key=lambda x: -(x.rrf_score or x.vector_score))
|
||||
for item in ranked:
|
||||
@@ -1574,6 +1606,8 @@ class QAV1Pipeline:
|
||||
ranked,
|
||||
top_k_for_llm=top_k_for_llm,
|
||||
graph_min_evidence=graph_min_evidence,
|
||||
# 回改 v2:search 预筛用向量域 0.45;chat 走 final 域 0.25
|
||||
min_score=(min_vector_score() if retrieval_only else None),
|
||||
)
|
||||
assign_evidence_ids(selected)
|
||||
timings["cluster_bge_rerank_and_select"] = time.perf_counter() - t0
|
||||
@@ -3316,7 +3350,19 @@ class QAV1Pipeline:
|
||||
*,
|
||||
top_k_for_llm: int,
|
||||
graph_min_evidence: int,
|
||||
min_score: float | None = None,
|
||||
) -> tuple[List[EvidenceItem], Dict[str, Any]]:
|
||||
# 缺陷 #28 回改 v2(2026-09-11):相关性下限双域。
|
||||
# - min_score 显式传入时用它(search/retrieval-only 传向量域
|
||||
# min_vector_score()=0.45);
|
||||
# - 缺省为 final 域 min_source_score()=0.25(chat 模式语义不变)。
|
||||
# 低于下限的候选不入选;全部低于下限 -> 返回空(前端呈现未命中)。
|
||||
_min_score = min_source_score() if min_score is None else min_score
|
||||
_min_score_domain = "vector" if min_score is not None else "final"
|
||||
_below_min = len(ranked) - sum(
|
||||
1 for item in ranked if item.final_score >= _min_score
|
||||
) # 被下限过滤的候选数(审计用,见 return 处)
|
||||
ranked = [item for item in ranked if item.final_score >= _min_score]
|
||||
graph_items = [item for item in ranked if cls._is_graph_item(item)]
|
||||
other_items = [item for item in ranked if not cls._is_graph_item(item)]
|
||||
selected: List[EvidenceItem] = []
|
||||
@@ -3365,6 +3411,9 @@ class QAV1Pipeline:
|
||||
"graph_selected": graph_selected,
|
||||
"non_graph_selected": len(selected) - graph_selected,
|
||||
"graph_quota_satisfied": graph_selected >= graph_min_evidence,
|
||||
"min_source_score": _min_score,
|
||||
"below_min_source_score": _below_min,
|
||||
"min_score_domain": _min_score_domain,
|
||||
"note": "Graph quota is enforced after BGE rerank. If fewer graph candidates exist, all available graph candidates are retained and the deficit is reported.",
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
"""score_gate.py —— 检索相关性下限(缺陷 #28 回改 v2,2026-09-11)。
|
||||
|
||||
双域双常量(TL 2026-09-11 在 175 裸进程 Qdrant 27433 + Ollama bge-m3:567m 实测):
|
||||
- 垃圾问题「景甜是谁」top10 vector_score(cosine) 全部 0.32-0.34;
|
||||
- arginine 正常问题 top10 vector_score 0.63+;
|
||||
- 垃圾 final_score(BGE rerank 后)仍为 0.19-0.20(测试 agent 实测)。
|
||||
|
||||
MIN_VECTOR_SCORE = 0.45(env QA_MIN_VECTOR_SCORE)
|
||||
- 语义:向量域(Qdrant Cosine 分数)相关性下限。0.45 位于垃圾带
|
||||
(0.32-0.34)与正常带(0.63+)之间并各留余量。
|
||||
- 用途:仅 search / retrieval-only 路径——Qdrant score_threshold 与
|
||||
检索模式预筛。chat 模式不使用(保召回,rerank 后用 MIN_SOURCE_SCORE 兜底)。
|
||||
|
||||
MIN_SOURCE_SCORE = 0.25(env QA_MIN_SOURCE_SCORE)
|
||||
- 语义:最终域(final_score = 0.8*sigmoid(bge) + 0.2*rrf/max_rrf)相关性
|
||||
下限。垃圾 final 0.19-0.20 与其留 ~25% 余量,正常命中 final >0.7。
|
||||
- 用途:chat 模式 _select_graph_balanced 的后置过滤(r4-28 语义不变)。
|
||||
|
||||
本模块无第三方依赖、无 IO,可被 v16_retriever / direct_flow_retriever /
|
||||
pipeline 直接 import(编译期与运行期均安全)。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
DEFAULT_MIN_VECTOR_SCORE = 0.45
|
||||
DEFAULT_MIN_SOURCE_SCORE = 0.25
|
||||
_ENV_VECTOR = "QA_MIN_VECTOR_SCORE"
|
||||
_ENV_SOURCE = "QA_MIN_SOURCE_SCORE"
|
||||
|
||||
|
||||
def _env_value(env_key: str, default: float) -> float:
|
||||
"""env 覆盖:值 ∈ [0,1] 才生效,否则回落默认。"""
|
||||
raw = os.environ.get(env_key, "").strip()
|
||||
if not raw:
|
||||
return default
|
||||
try:
|
||||
value = float(raw)
|
||||
except ValueError:
|
||||
return default
|
||||
if not (0.0 <= value <= 1.0):
|
||||
return default
|
||||
return value
|
||||
|
||||
|
||||
def min_vector_score() -> float:
|
||||
"""向量域下限(Qdrant score_threshold / search 预筛)。"""
|
||||
return min_vector_score.value
|
||||
|
||||
|
||||
def min_source_score() -> float:
|
||||
"""最终域下限(chat 模式 final_score 后置过滤)。"""
|
||||
return min_source_score.value
|
||||
|
||||
|
||||
min_vector_score.value = _env_value(_ENV_VECTOR, DEFAULT_MIN_VECTOR_SCORE) # type: ignore[attr-defined]
|
||||
min_source_score.value = _env_value(_ENV_SOURCE, DEFAULT_MIN_SOURCE_SCORE) # type: ignore[attr-defined]
|
||||
@@ -0,0 +1,201 @@
|
||||
"""test_bioflow_agent.py —— bioflow_agent 前门三重门 + nl_triggers 解析单测。
|
||||
|
||||
覆盖 S1 验收要点:
|
||||
- 熔断开关(QA_V1_ENABLE_BIOFLOW_AGENT)默认关 → 返回 None(走原 pipeline)
|
||||
- 意图预筛(nl_triggers)各自评估,未命中返回 None
|
||||
- 注册表未注册 / 会话开关关 → 返回 None
|
||||
- 命中后短路返回自然语言 payload(bioflow.handled=True)
|
||||
- 只装一个能力时,另一个能力的触发语不响应(隔离)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent.parent / "src"))
|
||||
|
||||
from rag.qa_v1_1_73 import bioflow_agent as ba # noqa: E402
|
||||
from rag.qa_v1_1_73.capabilities.manifests import parse_capabilities_text # noqa: E402
|
||||
from rag.qa_v1_1_73.capabilities.registry import CapabilityRegistry # noqa: E402
|
||||
|
||||
BIOFLOW_YAML = """
|
||||
capabilities:
|
||||
- name: bioflow-gene-family
|
||||
provider: mcp
|
||||
permission: read_only
|
||||
description: 基因家族分析
|
||||
session_default: true
|
||||
nl_triggers: [基因家族, 家族成员, 亚家族, 建树]
|
||||
mcp:
|
||||
url: http://127.0.0.1:8098/mcp/gene-family/
|
||||
prompt_pack:
|
||||
file: bioflow_gene_family_agent_rules.md
|
||||
- name: bioflow-rnaseq
|
||||
provider: mcp
|
||||
permission: read_only
|
||||
description: 转录组分析
|
||||
session_default: true
|
||||
nl_triggers: [转录组, RNA-seq, 差异表达, 富集]
|
||||
mcp:
|
||||
url: http://127.0.0.1:8098/mcp/rnaseq/
|
||||
prompt_pack:
|
||||
file: bioflow_rnaseq_agent_rules.md
|
||||
"""
|
||||
|
||||
|
||||
def _make_registry(tmp_path: Path) -> CapabilityRegistry:
|
||||
path = tmp_path / "capabilities.yaml"
|
||||
path.write_text(BIOFLOW_YAML, encoding="utf-8")
|
||||
registry = CapabilityRegistry(path)
|
||||
registry.load()
|
||||
return registry
|
||||
|
||||
|
||||
class FakeLLMClient:
|
||||
def __init__(self, content: str = "已为您规划分析路径。") -> None:
|
||||
self.content = content
|
||||
|
||||
def complete_raw_messages(self, system, messages, tools=None, **kwargs):
|
||||
return {"content": self.content, "tool_calls": [], "finish_reason": "stop"}
|
||||
|
||||
|
||||
class TestNlTriggers:
|
||||
def test_parse_nl_triggers(self):
|
||||
manifests, warnings = parse_capabilities_text(BIOFLOW_YAML)
|
||||
assert warnings == []
|
||||
gf = next(m for m in manifests if m.name == "bioflow-gene-family")
|
||||
assert gf.nl_triggers == ["基因家族", "家族成员", "亚家族", "建树"]
|
||||
assert gf.config["mcp"]["url"].endswith("/mcp/gene-family/")
|
||||
assert gf.config["prompt_pack"]["file"] == "bioflow_gene_family_agent_rules.md"
|
||||
|
||||
def test_nl_triggers_dedupe_and_clean(self):
|
||||
text = """
|
||||
capabilities:
|
||||
- name: t
|
||||
provider: mcp
|
||||
nl_triggers: [基因家族, 基因家族, "", " 亚家族 ", 123]
|
||||
"""
|
||||
manifests, _ = parse_capabilities_text(text)
|
||||
assert manifests[0].nl_triggers == ["基因家族", "亚家族"]
|
||||
|
||||
|
||||
class TestIntentMatch:
|
||||
def test_match(self):
|
||||
assert ba._intent_match("帮我做水稻SPL基因家族分析", ["基因家族"]) is True
|
||||
assert ba._intent_match("转录组差异表达分析", ["转录组", "RNA-seq"]) is True
|
||||
assert ba._intent_match("什么是光合作用", ["基因家族", "转录组"]) is False
|
||||
|
||||
|
||||
class TestGateLogic:
|
||||
def test_fuse_off_returns_none(self, tmp_path):
|
||||
registry = _make_registry(tmp_path)
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
with mock.patch.dict("os.environ", {}, clear=False):
|
||||
ba.os.environ.pop(ba.ENV_FUSE, None)
|
||||
result = ba.try_bioflow_agent(
|
||||
"帮我做基因家族分析", llm_client=FakeLLMClient(), capabilities=None
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_no_intent_match_returns_none(self, tmp_path):
|
||||
registry = _make_registry(tmp_path)
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
with mock.patch.dict("os.environ", {ba.ENV_FUSE: "1"}):
|
||||
result = ba.try_bioflow_agent(
|
||||
"光合作用的电子传递链是什么", llm_client=FakeLLMClient(), capabilities=None
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_active_job_continuation_skips_intent_prescreen(self, tmp_path):
|
||||
"""#011:有活跃任务时,续接输入(不含触发词)也应进入生信流,不掉普通 QA。"""
|
||||
registry = _make_registry(tmp_path)
|
||||
fake_tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "fetch_source_file",
|
||||
"description": "获取源文件",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
]
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
with mock.patch.object(ba, "_build_tools", return_value=(fake_tools, {"fetch_source_file": mock.ANY})):
|
||||
with mock.patch.dict("os.environ", {ba.ENV_FUSE: "1"}):
|
||||
result = ba.try_bioflow_agent(
|
||||
"基因组文件在 https://example.com/rice.fa",
|
||||
llm_client=FakeLLMClient(),
|
||||
capabilities=None,
|
||||
bioflow={"active_job_id": "job-1", "chain": "gene_family"},
|
||||
)
|
||||
assert result is not None
|
||||
assert result["success"] is True
|
||||
assert result["bioflow"]["handled"] is True
|
||||
assert result["bioflow"]["capabilities"] == ["bioflow-gene-family"]
|
||||
|
||||
def test_not_registered_returns_none(self, tmp_path):
|
||||
empty = CapabilityRegistry(tmp_path / "empty.yaml")
|
||||
empty.load()
|
||||
with mock.patch.object(ba, "get_registry", return_value=empty):
|
||||
with mock.patch.dict("os.environ", {ba.ENV_FUSE: "1"}):
|
||||
result = ba.try_bioflow_agent(
|
||||
"帮我做基因家族分析", llm_client=FakeLLMClient(), capabilities=None
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_session_switch_off_returns_none(self, tmp_path):
|
||||
registry = _make_registry(tmp_path)
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
with mock.patch.dict("os.environ", {ba.ENV_FUSE: "1"}):
|
||||
result = ba.try_bioflow_agent(
|
||||
"帮我做基因家族分析",
|
||||
llm_client=FakeLLMClient(),
|
||||
capabilities={"bioflow-gene-family": False},
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_isolation_only_installed_responds(self, tmp_path):
|
||||
registry = _make_registry(tmp_path)
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
with mock.patch.dict("os.environ", {ba.ENV_FUSE: "1"}):
|
||||
caps = ba._candidate_capabilities("帮我做基因家族分析", None)
|
||||
assert [m.name for m in caps] == ["bioflow-gene-family"]
|
||||
caps = ba._candidate_capabilities("转录组差异表达", None)
|
||||
assert [m.name for m in caps] == ["bioflow-rnaseq"]
|
||||
|
||||
def test_full_flow_short_circuits(self, tmp_path):
|
||||
registry = _make_registry(tmp_path)
|
||||
fake_tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "plan_gene_family",
|
||||
"description": "规划基因家族路径",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
]
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
with mock.patch.object(ba, "_build_tools", return_value=(fake_tools, {"plan_gene_family": mock.ANY})):
|
||||
with mock.patch.dict("os.environ", {ba.ENV_FUSE: "1"}):
|
||||
result = ba.try_bioflow_agent(
|
||||
"帮我做基因家族分析", llm_client=FakeLLMClient(), capabilities=None
|
||||
)
|
||||
assert result is not None
|
||||
assert result["success"] is True
|
||||
assert result["answer"] == "已为您规划分析路径。"
|
||||
assert result["bioflow"]["handled"] is True
|
||||
assert result["bioflow"]["capabilities"] == ["bioflow-gene-family"]
|
||||
|
||||
def test_exception_falls_back_to_none(self, tmp_path):
|
||||
registry = _make_registry(tmp_path)
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
with mock.patch.object(ba, "_candidate_capabilities", side_effect=RuntimeError("boom")):
|
||||
with mock.patch.dict("os.environ", {ba.ENV_FUSE: "1"}):
|
||||
result = ba.try_bioflow_agent(
|
||||
"帮我做基因家族分析", llm_client=FakeLLMClient(), capabilities=None
|
||||
)
|
||||
assert result is None
|
||||
@@ -0,0 +1,255 @@
|
||||
"""test_bioflow_agent_s2.py —— S2 对话式编排闭环单测。
|
||||
|
||||
覆盖:card_result 回填协议(四类卡片)、参数抽取(白名单)、多轮记忆注入、
|
||||
session↔job 绑定上下文(active_job_id/chain)、job_status 主动查询。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent.parent / "src"))
|
||||
|
||||
from rag.qa_v1_1_73 import bioflow_agent as ba # noqa: E402
|
||||
from rag.qa_v1_1_73.capabilities.manifests import CapabilityManifest # noqa: E402
|
||||
from rag.qa_v1_1_73.tool_runner import ToolRunner # noqa: E402
|
||||
|
||||
|
||||
def _manifest(name="bioflow-gene-family", url="http://x/mcp/gene-family/"):
|
||||
return CapabilityManifest(name=name, provider="mcp", config={"mcp": {"url": url, "tool": name}})
|
||||
|
||||
|
||||
def _tool_manifests():
|
||||
"""覆盖四类卡片用到的工具 + job_status。"""
|
||||
names = ("create_gene_family_job", "create_rnaseq_job", "confirm_stage",
|
||||
"register_job_files", "fetch_source_file", "job_status")
|
||||
out = {}
|
||||
for t in names:
|
||||
out[t] = CapabilityManifest(
|
||||
name="bioflow-gene-family", provider="mcp",
|
||||
config={"mcp": {"url": "http://x/mcp/gene-family/", "tool": t}},
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
class FakeDispatcher:
|
||||
def __init__(self, results=None):
|
||||
self.results = results or {}
|
||||
self.calls = []
|
||||
|
||||
def execute(self, manifest, arguments):
|
||||
tool = manifest.config.get("mcp", {}).get("tool")
|
||||
self.calls.append((tool, arguments))
|
||||
return self.results.get(tool, '{"ok": true, "data": {"id": "job-1"}}')
|
||||
|
||||
|
||||
class FakeLLM:
|
||||
def __init__(self, complete_text="已收到,已为您处理。", json_payload=None):
|
||||
self.complete_text = complete_text
|
||||
self.json_payload = json_payload or {}
|
||||
|
||||
def complete(self, system, user, **kw):
|
||||
return self.complete_text
|
||||
|
||||
def complete_json(self, system, user, **kw):
|
||||
return self.json_payload
|
||||
|
||||
def complete_raw_messages(self, system, messages, tools=None, **kw):
|
||||
return {"content": "answer", "tool_calls": [], "finish_reason": "stop"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 参数抽取
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestExtractParams:
|
||||
def test_extract_known_and_unknown(self):
|
||||
llm = FakeLLM(json_payload={"params": {"threads": 8, "padj": 0.01}, "unknown": ["bogus"]})
|
||||
form_defaults = {"threads": {"default": 4}, "padj": {"default": 0.05}}
|
||||
known, unknown = ba._extract_params(llm, "用 8 核,padj 用 0.01", form_defaults)
|
||||
assert known == {"threads": 8, "padj": 0.01}
|
||||
assert unknown == ["bogus"]
|
||||
|
||||
def test_non_whitelist_param_filtered(self):
|
||||
llm = FakeLLM(json_payload={"params": {"evil": 1, "threads": 4}, "unknown": []})
|
||||
form_defaults = {"threads": {"default": 4}}
|
||||
known, unknown = ba._extract_params(llm, "x", form_defaults)
|
||||
assert known == {"threads": 4}
|
||||
|
||||
def test_bad_json_returns_empty(self):
|
||||
llm = FakeLLM(json_payload="not-a-dict")
|
||||
form_defaults = {"threads": {}}
|
||||
known, unknown = ba._extract_params(llm, "x", form_defaults)
|
||||
assert known == {} and unknown == []
|
||||
|
||||
def test_empty_whitelist(self):
|
||||
llm = FakeLLM()
|
||||
known, unknown = ba._extract_params(llm, "x", {})
|
||||
assert known == {} and unknown == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# card_result 回填协议
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestCardResult:
|
||||
def test_path_selection_creates_job_and_extracts_job_id(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
llm = FakeLLM()
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "path_selection",
|
||||
"payload": {"chain": "gene_family", "stages": ["identify", "tree"], "title": "水稻SPL"}},
|
||||
llm, {"session_id": "s1"}, _tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
assert out["bioflow"]["active_job_id"] == "job-1"
|
||||
assert out["bioflow"]["chain"] == "gene_family"
|
||||
assert dispatcher.calls[0][0] == "create_gene_family_job"
|
||||
assert dispatcher.calls[0][1]["stages"] == ["identify", "tree"]
|
||||
assert dispatcher.calls[0][1]["session_id"] == "s1"
|
||||
|
||||
def test_path_selection_missing_stages_returns_none(self):
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=FakeDispatcher()):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "path_selection", "payload": {"chain": "gene_family"}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is None
|
||||
|
||||
def test_form_confirm_calls_confirm_stage(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "form_confirm",
|
||||
"payload": {"job_id": "job-9", "stage": "identify", "params": {"evalue": 1e-5}}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
assert dispatcher.calls[0][0] == "confirm_stage"
|
||||
assert dispatcher.calls[0][1] == {"job_id": "job-9", "stage": "identify", "params": {"evalue": 1e-5}}
|
||||
|
||||
def test_form_confirm_uses_active_job_id_fallback(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
ba._handle_card_result(
|
||||
{"card": "form_confirm", "payload": {"stage": "identify", "params": {}}},
|
||||
FakeLLM(), {"active_job_id": "job-7"}, _tool_manifests(),
|
||||
)
|
||||
assert dispatcher.calls[0][1]["job_id"] == "job-7"
|
||||
|
||||
def test_file_upload_splits_remote_and_local(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "file_upload",
|
||||
"payload": {"job_id": "j1", "paths": ["https://x/a.fa", "/data/b.fa"]}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
tools = [c[0] for c in dispatcher.calls]
|
||||
assert "fetch_source_file" in tools and "register_job_files" in tools
|
||||
|
||||
def test_archive_choice_paraphrases(self):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "archive_choice", "payload": {"job_id": "j1", "archive": True}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
assert out["answer"] == "已收到,已为您处理。"
|
||||
|
||||
def test_unknown_card_returns_none(self):
|
||||
out = ba._handle_card_result({"card": "bogus", "payload": {}}, FakeLLM(), {}, _tool_manifests())
|
||||
assert out is None
|
||||
|
||||
def test_path_selection_reuse_failure_no_stages_returns_none(self):
|
||||
"""bug ②:复用失败且无显式 stages → 返回 None(不空链建任务)。"""
|
||||
dispatcher = FakeDispatcher({"reuse_archive": '{"ok": false, "code": "ARCHIVE_NOT_FOUND"}'})
|
||||
tm = _tool_manifests()
|
||||
tm["reuse_archive"] = CapabilityManifest(
|
||||
name="bioflow-gene-family", provider="mcp",
|
||||
config={"mcp": {"url": "http://x", "tool": "reuse_archive"}},
|
||||
)
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "path_selection",
|
||||
"payload": {"chain": "gene_family", "reuse_archive": "不存在的存档"}},
|
||||
FakeLLM(), {}, tm,
|
||||
)
|
||||
assert out is None
|
||||
# 只执行了 reuse_archive,未调用 create_gene_family_job
|
||||
assert [c[0] for c in dispatcher.calls] == ["reuse_archive"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 前门路由(card_result / bioflow 上下文)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestRouteCardResult:
|
||||
def test_try_bioflow_agent_card_result_short_circuits(self):
|
||||
registry_manifest = _manifest()
|
||||
with mock.patch.dict("os.environ", {ba.ENV_FUSE: "1"}):
|
||||
with mock.patch.object(ba, "_registered_capabilities", return_value=[registry_manifest]):
|
||||
with mock.patch.object(ba, "_build_tools", return_value=([{"type": "function", "function": {"name": "confirm_stage", "description": "", "parameters": {"type": "object", "properties": {}}}}], {"confirm_stage": registry_manifest})):
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=FakeDispatcher()):
|
||||
out = ba.try_bioflow_agent(
|
||||
"确认", llm_client=FakeLLM(),
|
||||
card_result={"card": "form_confirm", "payload": {"job_id": "j1", "stage": "s", "params": {}}},
|
||||
)
|
||||
assert out is not None
|
||||
assert out["bioflow"]["handled"] is True
|
||||
|
||||
def test_try_bioflow_agent_fuse_off_ignores_card_result(self):
|
||||
with mock.patch.dict("os.environ", {}, clear=False):
|
||||
ba.os.environ.pop(ba.ENV_FUSE, None)
|
||||
out = ba.try_bioflow_agent(
|
||||
"确认", llm_client=FakeLLM(),
|
||||
card_result={"card": "form_confirm", "payload": {}},
|
||||
)
|
||||
assert out is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 多轮记忆注入(ToolRunner.run history)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestMultiTurnHistory:
|
||||
def test_history_injected_into_messages(self):
|
||||
class RecordingLLM:
|
||||
def __init__(self):
|
||||
self.captured_messages = None
|
||||
def complete_raw_messages(self, system, messages, tools=None, **kw):
|
||||
self.captured_messages = list(messages)
|
||||
return {"content": "ok", "tool_calls": [], "finish_reason": "stop"}
|
||||
|
||||
from rag.qa_v1_1_73.tool_runner import make_llm_complete
|
||||
|
||||
llm = RecordingLLM()
|
||||
runner = ToolRunner(make_llm_complete(llm), max_rounds=1)
|
||||
history = [{"role": "user", "content": "上一轮问题"}, {"role": "assistant", "content": "上一轮答案"}]
|
||||
runner.run("sys", "当前问题", [], {}, history=history)
|
||||
msgs = llm.captured_messages
|
||||
assert msgs[0] == {"role": "user", "content": "上一轮问题"}
|
||||
assert msgs[1] == {"role": "assistant", "content": "上一轮答案"}
|
||||
assert msgs[-1] == {"role": "user", "content": "当前问题"}
|
||||
|
||||
def test_history_filters_bad_roles(self):
|
||||
class RecordingLLM:
|
||||
def __init__(self):
|
||||
self.captured_messages = None
|
||||
def complete_raw_messages(self, system, messages, tools=None, **kw):
|
||||
self.captured_messages = list(messages)
|
||||
return {"content": "ok", "tool_calls": [], "finish_reason": "stop"}
|
||||
|
||||
from rag.qa_v1_1_73.tool_runner import make_llm_complete
|
||||
|
||||
llm = RecordingLLM()
|
||||
runner = ToolRunner(make_llm_complete(llm), max_rounds=1)
|
||||
history = [{"role": "system", "content": "evil"}, {"role": "user", "content": ""}, {"role": "user", "content": "ok"}]
|
||||
runner.run("sys", "q", [], {}, history=history)
|
||||
roles = [m["role"] for m in llm.captured_messages]
|
||||
assert "system" not in roles
|
||||
assert llm.captured_messages[-1] == {"role": "user", "content": "q"}
|
||||
@@ -0,0 +1,511 @@
|
||||
"""test_bioflow_agent_s2_supplement.py —— S2 覆盖清单补强测试。
|
||||
|
||||
补强 dev 的 test_bioflow_agent_s2.py 未覆盖的分支:
|
||||
- _chain_to_capability / _extract_job_id / _execute_tool / _paraphrase / _query_job_status
|
||||
- card 四类 handler 缺必需字段回退 None、files 键、archive=False、非 dict payload
|
||||
- _route_card_result:chain 过滤 / 未注册能力→None / tools 空→None
|
||||
- NL 路径注入 bioflow 上下文(active_job_id/chain)
|
||||
- system prompt 硬规则在位(追问 ≤3 / 不可代确认 / 工具名不外泄)
|
||||
- ToolRunner.run(history) 边界(non-dict 条目 / history=None / 空内容剔除)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent.parent / "src"))
|
||||
|
||||
from rag.qa_v1_1_73 import bioflow_agent as ba # noqa: E402
|
||||
from rag.qa_v1_1_73.capabilities.manifests import CapabilityManifest # noqa: E402
|
||||
from rag.qa_v1_1_73.tool_runner import ToolRunner, make_llm_complete # noqa: E402
|
||||
|
||||
|
||||
def _manifest(name="bioflow-gene-family", url="http://x/mcp/gene-family/"):
|
||||
return CapabilityManifest(name=name, provider="mcp", config={"mcp": {"url": url, "tool": name}})
|
||||
|
||||
|
||||
def _tool_manifests():
|
||||
names = ("create_gene_family_job", "create_rnaseq_job", "confirm_stage",
|
||||
"register_job_files", "fetch_source_file", "job_status",
|
||||
"rerun_stage", "reuse_archive", "archive_job")
|
||||
out = {}
|
||||
for t in names:
|
||||
out[t] = CapabilityManifest(
|
||||
name="bioflow-gene-family", provider="mcp",
|
||||
config={"mcp": {"url": "http://x/mcp/gene-family/", "tool": t}},
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
class FakeLLM:
|
||||
def __init__(self, complete_text="已收到。", json_payload=None, complete_error=None):
|
||||
self.complete_text = complete_text
|
||||
self.json_payload = json_payload or {}
|
||||
self.complete_error = complete_error
|
||||
|
||||
def complete(self, system, user, **kw):
|
||||
if self.complete_error:
|
||||
raise self.complete_error
|
||||
return self.complete_text
|
||||
|
||||
def complete_json(self, system, user, **kw):
|
||||
return self.json_payload
|
||||
|
||||
def complete_raw_messages(self, system, messages, tools=None, **kw):
|
||||
return {"content": "answer", "tool_calls": [], "finish_reason": "stop"}
|
||||
|
||||
|
||||
class FakeDispatcher:
|
||||
def __init__(self, results=None):
|
||||
self.results = results or {}
|
||||
self.calls = []
|
||||
|
||||
def execute(self, manifest, arguments):
|
||||
tool = manifest.config.get("mcp", {}).get("tool")
|
||||
self.calls.append((tool, arguments))
|
||||
return self.results.get(tool, '{"ok": true, "data": {"id": "job-1"}}')
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 辅助函数
|
||||
# ===========================================================================
|
||||
|
||||
class TestChainToCapability:
|
||||
def test_known(self):
|
||||
assert ba._chain_to_capability("gene_family") == "bioflow-gene-family"
|
||||
assert ba._chain_to_capability("rnaseq") == "bioflow-rnaseq"
|
||||
|
||||
def test_unknown_and_empty(self):
|
||||
assert ba._chain_to_capability("wgs") is None
|
||||
assert ba._chain_to_capability("") is None
|
||||
assert ba._chain_to_capability(None) is None
|
||||
|
||||
def test_case_and_whitespace(self):
|
||||
assert ba._chain_to_capability(" Gene_Family ") == "bioflow-gene-family"
|
||||
|
||||
|
||||
class TestExtractJobId:
|
||||
def test_data_id(self):
|
||||
assert ba._extract_job_id('{"ok": true, "data": {"id": "job-9"}}') == "job-9"
|
||||
|
||||
def test_data_job_id(self):
|
||||
assert ba._extract_job_id('{"ok": true, "data": {"job_id": "job-8"}}') == "job-8"
|
||||
|
||||
def test_flat_id(self):
|
||||
assert ba._extract_job_id('{"ok": true, "id": "job-7"}') == "job-7"
|
||||
|
||||
def test_non_json_and_no_id(self):
|
||||
assert ba._extract_job_id("not json") is None
|
||||
assert ba._extract_job_id('{"ok": false}') is None
|
||||
assert ba._extract_job_id("") is None
|
||||
|
||||
|
||||
class TestExecuteTool:
|
||||
def test_unknown_tool(self):
|
||||
out = ba._execute_tool(_tool_manifests(), "ghost", {})
|
||||
assert "未注册" in out
|
||||
|
||||
def test_tool_exception_becomes_text(self):
|
||||
class Boom:
|
||||
def execute(self, manifest, arguments):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=Boom()):
|
||||
out = ba._execute_tool(_tool_manifests(), "job_status", {"job_id": "j"})
|
||||
assert "工具执行失败" in out
|
||||
|
||||
|
||||
class TestParaphrase:
|
||||
def test_llm_success(self):
|
||||
out = ba._paraphrase(FakeLLM(complete_text="已转述。"), "原文")
|
||||
assert out == "已转述。"
|
||||
|
||||
def test_llm_exception_falls_back_to_raw(self):
|
||||
out = ba._paraphrase(FakeLLM(complete_error=RuntimeError("boom")), "原文")
|
||||
assert out == "原文"
|
||||
|
||||
def test_llm_empty_falls_back_to_raw(self):
|
||||
out = ba._paraphrase(FakeLLM(complete_text=" "), "原文")
|
||||
assert out == "原文"
|
||||
|
||||
|
||||
class TestQueryJobStatus:
|
||||
def test_calls_job_status_and_paraphrases(self):
|
||||
with mock.patch.object(ba, "_execute_tool", return_value='{"ok": true, "data": {"status": "running"}}'):
|
||||
out = ba._query_job_status("j1", _tool_manifests(), FakeLLM(complete_text="任务进行中。"))
|
||||
assert out == "任务进行中。"
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# card 四类 handler 边界(缺必需字段 / files 键 / archive=False / 非 dict payload)
|
||||
# ===========================================================================
|
||||
|
||||
class TestCardHandlersEdge:
|
||||
def test_form_confirm_missing_job_and_stage_returns_none(self):
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=FakeDispatcher()):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "form_confirm", "payload": {"params": {}}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is None
|
||||
|
||||
def test_file_upload_missing_job_returns_none(self):
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=FakeDispatcher()):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "file_upload", "payload": {"paths": ["/data/a.fa"]}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is None
|
||||
|
||||
def test_file_upload_missing_files_returns_none(self):
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=FakeDispatcher()):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "file_upload", "payload": {"job_id": "j1"}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is None
|
||||
|
||||
def test_file_upload_uses_files_key(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "file_upload",
|
||||
"payload": {"job_id": "j1", "files": ["https://x/a.fa"]}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
assert dispatcher.calls[0][0] == "fetch_source_file"
|
||||
|
||||
def test_file_upload_string_files_coerced(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "file_upload",
|
||||
"payload": {"job_id": "j1", "files": "https://x/a.fa"}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
assert dispatcher.calls[0][0] == "fetch_source_file"
|
||||
|
||||
def test_archive_choice_missing_job_returns_none(self):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "archive_choice", "payload": {"archive": True}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is None
|
||||
|
||||
def test_archive_choice_skip(self):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "archive_choice", "payload": {"job_id": "j1", "archive": False}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
assert out["bioflow"]["audit"]["tool_calls"] == 0 # 存档不调工具(走 REST)
|
||||
|
||||
def test_path_selection_rnaseq_chain(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "path_selection",
|
||||
"payload": {"chain": "rnaseq", "stages": ["rnaseq_qc"]}},
|
||||
FakeLLM(), {"session_id": "s1"}, _tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
assert dispatcher.calls[0][0] == "create_rnaseq_job"
|
||||
|
||||
def test_card_non_dict_payload_treated_empty(self):
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=FakeDispatcher()):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "form_confirm", "payload": "not-a-dict"},
|
||||
FakeLLM(), {"active_job_id": "j1"}, _tool_manifests(),
|
||||
)
|
||||
assert out is None # payload 非 dict → 无 stage/job_id → None
|
||||
|
||||
def test_card_result_with_non_dict_is_ignored(self):
|
||||
# try_bioflow_agent 对非 dict 的 card_result 不做 card 路由,走 NL 意图
|
||||
with mock.patch.dict("os.environ", {ba.ENV_FUSE: "1"}):
|
||||
with mock.patch.object(ba, "_candidate_capabilities", return_value=[]):
|
||||
out = ba.try_bioflow_agent("确认", llm_client=FakeLLM(), card_result="bogus")
|
||||
assert out is None
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# _route_card_result 边界
|
||||
# ===========================================================================
|
||||
|
||||
class TestRouteCardResultEdge:
|
||||
def test_chain_filters_to_capability(self):
|
||||
gf = _manifest("bioflow-gene-family")
|
||||
rn = _manifest("bioflow-rnaseq", "http://x/mcp/rnaseq/")
|
||||
captured = {}
|
||||
|
||||
def fake_build(manifests):
|
||||
captured["manifests"] = [m.name for m in manifests]
|
||||
return ([], {})
|
||||
|
||||
with mock.patch.object(ba, "_registered_capabilities", return_value=[gf, rn]):
|
||||
with mock.patch.object(ba, "_build_tools", side_effect=fake_build):
|
||||
ba._route_card_result(
|
||||
{"card": "form_confirm", "payload": {"chain": "rnaseq"}},
|
||||
FakeLLM(), None, None, None,
|
||||
)
|
||||
assert captured["manifests"] == ["bioflow-rnaseq"] # chain 过滤到 rnaseq 能力
|
||||
|
||||
def test_chain_to_unregistered_capability_returns_none(self):
|
||||
gf = _manifest("bioflow-gene-family")
|
||||
with mock.patch.object(ba, "_registered_capabilities", return_value=[gf]):
|
||||
with mock.patch.object(ba, "_build_tools", return_value=([{"x": 1}], {})):
|
||||
out = ba._route_card_result(
|
||||
{"card": "form_confirm", "payload": {"chain": "rnaseq"}},
|
||||
FakeLLM(), None, None, None,
|
||||
)
|
||||
assert out is None # rnaseq 未注册 → manifests 空 → None
|
||||
|
||||
def test_tools_empty_returns_none(self):
|
||||
gf = _manifest("bioflow-gene-family")
|
||||
with mock.patch.object(ba, "_registered_capabilities", return_value=[gf]):
|
||||
with mock.patch.object(ba, "_build_tools", return_value=([], {})):
|
||||
out = ba._route_card_result(
|
||||
{"card": "form_confirm", "payload": {}},
|
||||
FakeLLM(), None, None, None,
|
||||
)
|
||||
assert out is None
|
||||
|
||||
def test_progress_callback_invoked(self):
|
||||
gf = _manifest("bioflow-gene-family")
|
||||
calls = []
|
||||
with mock.patch.object(ba, "_registered_capabilities", return_value=[gf]):
|
||||
with mock.patch.object(ba, "_build_tools", return_value=([{"x": 1}], {})):
|
||||
with mock.patch.object(ba, "_handle_card_result", return_value=None):
|
||||
ba._route_card_result(
|
||||
{"card": "form_confirm", "payload": {}},
|
||||
FakeLLM(), None, None, lambda p: calls.append(p),
|
||||
)
|
||||
assert calls and calls[0]["stage"] == "bioflow_card"
|
||||
assert calls[0]["card"] == "form_confirm"
|
||||
|
||||
def test_progress_callback_exception_swallowed(self):
|
||||
gf = _manifest("bioflow-gene-family")
|
||||
|
||||
def boom(p):
|
||||
raise RuntimeError("callback boom")
|
||||
|
||||
with mock.patch.object(ba, "_registered_capabilities", return_value=[gf]):
|
||||
with mock.patch.object(ba, "_build_tools", return_value=([{"x": 1}], {})):
|
||||
with mock.patch.object(ba, "_handle_card_result", return_value={"ok": True}):
|
||||
out = ba._route_card_result(
|
||||
{"card": "form_confirm", "payload": {}},
|
||||
FakeLLM(), None, None, boom,
|
||||
)
|
||||
assert out == {"ok": True} # 回调异常被吞,不影响主流程
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# NL 路径:bioflow 上下文注入 / 卡片与 NL 语义等价
|
||||
# ===========================================================================
|
||||
|
||||
class TestNLPathBioflowContext:
|
||||
def test_nl_path_injects_bioflow_context(self, tmp_path):
|
||||
from rag.qa_v1_1_73.capabilities.registry import CapabilityRegistry
|
||||
from rag.qa_v1_1_73.capabilities.manifests import parse_capabilities_text
|
||||
yaml = """
|
||||
capabilities:
|
||||
- name: bioflow-gene-family
|
||||
provider: mcp
|
||||
permission: read_only
|
||||
session_default: true
|
||||
nl_triggers: [基因家族]
|
||||
mcp:
|
||||
url: http://127.0.0.1:8098/mcp/gene-family/
|
||||
"""
|
||||
(tmp_path / "c.yaml").write_text(yaml, encoding="utf-8")
|
||||
registry = CapabilityRegistry(tmp_path / "c.yaml")
|
||||
registry.load()
|
||||
fake_tools = [{"type": "function", "function": {"name": "plan_gene_family",
|
||||
"description": "x",
|
||||
"parameters": {"type": "object", "properties": {}}}}]
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
with mock.patch.dict("os.environ", {ba.ENV_FUSE: "1"}):
|
||||
with mock.patch.object(ba, "_build_tools", return_value=(fake_tools, {})):
|
||||
with mock.patch.object(ba, "ToolRunner") as FakeRunner:
|
||||
FakeRunner.return_value.run.return_value = ("已答复", {"rounds": 1})
|
||||
out = ba.try_bioflow_agent(
|
||||
"帮我做基因家族分析", llm_client=FakeLLM(),
|
||||
bioflow={"session_id": "s1", "active_job_id": "job-9", "chain": "gene_family"},
|
||||
)
|
||||
assert out is not None
|
||||
assert out["bioflow"]["active_job_id"] == "job-9"
|
||||
assert out["bioflow"]["chain"] == "gene_family"
|
||||
|
||||
def test_card_and_nl_share_tool_chain(self):
|
||||
"""卡片操作与纯 NL 服务端语义等价:卡片 handler 只调 MCP 工具,不直连控制面。"""
|
||||
# 卡片 handler 用到的工具名 = MCP 工具(job_status/confirm_stage/create_* 等)
|
||||
tm = _tool_manifests()
|
||||
# 四类卡片触发的工具都在 tool_manifests(即 MCP 工具表)里
|
||||
for tool in ("create_gene_family_job", "create_rnaseq_job", "confirm_stage",
|
||||
"register_job_files", "fetch_source_file", "job_status"):
|
||||
assert tool in tm
|
||||
|
||||
|
||||
class TestExtractParamsEdge:
|
||||
def test_llm_exception_returns_empty(self):
|
||||
class RaisingLLM(FakeLLM):
|
||||
def complete_json(self, system, user, **kw):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
known, unknown = ba._extract_params(RaisingLLM(), "x", {"threads": {}})
|
||||
assert known == {} and unknown == []
|
||||
|
||||
def test_unknown_as_string_coerced_to_list(self):
|
||||
llm = FakeLLM(json_payload={"params": {"threads": 8}, "unknown": "bogus_param"})
|
||||
known, unknown = ba._extract_params(llm, "x", {"threads": {}})
|
||||
assert known == {"threads": 8}
|
||||
assert unknown == ["bogus_param"]
|
||||
|
||||
def test_path_selection_with_params(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
ba._handle_card_result(
|
||||
{"card": "path_selection",
|
||||
"payload": {"chain": "gene_family", "stages": ["identify"],
|
||||
"params": {"prefix": "Ft.SPL"}}},
|
||||
FakeLLM(), {"session_id": "s1"}, _tool_manifests(),
|
||||
)
|
||||
assert dispatcher.calls[0][1]["params"] == {"prefix": "Ft.SPL"}
|
||||
|
||||
def test_path_selection_create_tool_missing_returns_none(self):
|
||||
tm = _tool_manifests()
|
||||
tm.pop("create_gene_family_job")
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=FakeDispatcher()):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "path_selection",
|
||||
"payload": {"chain": "gene_family", "stages": ["identify"]}},
|
||||
FakeLLM(), {"session_id": "s1"}, tm,
|
||||
)
|
||||
assert out is None
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# system prompt 硬规则在位
|
||||
# ===========================================================================
|
||||
|
||||
class TestSystemPromptRules:
|
||||
def test_param_retry_rule(self):
|
||||
assert "最多 3 次" in ba._SYSTEM_BASE
|
||||
|
||||
def test_human_boundary_rule(self):
|
||||
assert "不可跳过" in ba._SYSTEM_BASE
|
||||
assert "不可代用户确认" in ba._SYSTEM_BASE
|
||||
|
||||
def test_tool_name_no_leak_rule(self):
|
||||
assert "绝不向用户泄露工具名" in ba._SYSTEM_BASE
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# ToolRunner.run(history) 边界
|
||||
# ===========================================================================
|
||||
|
||||
class TestHistoryEdge:
|
||||
def _run_with_history(self, history):
|
||||
captured = {}
|
||||
|
||||
class RecordingLLM:
|
||||
def complete_raw_messages(self, system, messages, tools=None, **kw):
|
||||
captured["messages"] = list(messages)
|
||||
return {"content": "ok", "tool_calls": [], "finish_reason": "stop"}
|
||||
|
||||
runner = ToolRunner(make_llm_complete(RecordingLLM()), max_rounds=1)
|
||||
runner.run("sys", "当前", [], {}, history=history)
|
||||
return captured["messages"]
|
||||
|
||||
def test_none_history(self):
|
||||
msgs = self._run_with_history(None)
|
||||
assert msgs == [{"role": "user", "content": "当前"}]
|
||||
|
||||
def test_non_dict_items_skipped(self):
|
||||
msgs = self._run_with_history(["not-a-dict", {"role": "user", "content": "ok"}])
|
||||
assert msgs == [{"role": "user", "content": "ok"}, {"role": "user", "content": "当前"}]
|
||||
|
||||
def test_empty_content_filtered(self):
|
||||
msgs = self._run_with_history([{"role": "user", "content": " "}])
|
||||
assert msgs == [{"role": "user", "content": "当前"}]
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# S3:card_result rerun / reuse_archive / archive_job 语义
|
||||
# ===========================================================================
|
||||
|
||||
class TestS3CardSemantics:
|
||||
def test_form_confirm_rerun_true_uses_rerun_stage(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "form_confirm",
|
||||
"payload": {"job_id": "j1", "stage": "identify",
|
||||
"params": {"evalue": 1e-5}, "rerun": True}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
assert dispatcher.calls[0][0] == "rerun_stage"
|
||||
assert out["bioflow"]["audit"]["tool_names"] == ["rerun_stage"]
|
||||
|
||||
def test_form_confirm_rerun_false_uses_confirm_stage(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
ba._handle_card_result(
|
||||
{"card": "form_confirm",
|
||||
"payload": {"job_id": "j1", "stage": "identify", "params": {}}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert dispatcher.calls[0][0] == "confirm_stage"
|
||||
|
||||
def test_path_selection_reuse_archive(self):
|
||||
dispatcher = FakeDispatcher({
|
||||
"reuse_archive": '{"ok": true, "data": {"id": "arch-1", '
|
||||
'"stage_chain": ["identify", "phylogeny"], '
|
||||
'"params": {"identify": {"prefix": "Ft"}}}}',
|
||||
})
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "path_selection",
|
||||
"payload": {"chain": "gene_family", "reuse_archive": "SPL 鉴定"}},
|
||||
FakeLLM(), {"session_id": "s1"}, _tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
tools = [c[0] for c in dispatcher.calls]
|
||||
assert tools[0] == "reuse_archive"
|
||||
assert tools[1] == "create_gene_family_job"
|
||||
create_args = dispatcher.calls[1][1]
|
||||
assert create_args["archive_id"] == "arch-1"
|
||||
assert create_args["stages"] == ["identify", "phylogeny"]
|
||||
assert create_args["params"] == {"identify": {"prefix": "Ft"}}
|
||||
assert out["bioflow"]["audit"]["tool_names"] == ["reuse_archive", "create_gene_family_job"]
|
||||
|
||||
def test_archive_choice_archive_true_calls_archive_job(self):
|
||||
dispatcher = FakeDispatcher({"archive_job": '{"ok": true, "data": {"id": "a1"}}'})
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "archive_choice",
|
||||
"payload": {"job_id": "j1", "archive": True,
|
||||
"name": "SPL 存档", "description": "d"}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
assert dispatcher.calls[0][0] == "archive_job"
|
||||
assert dispatcher.calls[0][1] == {"job_id": "j1", "name": "SPL 存档", "description": "d"}
|
||||
assert out["bioflow"]["audit"]["tool_names"] == ["archive_job"]
|
||||
|
||||
def test_archive_choice_archive_false_no_tool_call(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "archive_choice",
|
||||
"payload": {"job_id": "j1", "archive": False}},
|
||||
FakeLLM(), {}, _tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
assert dispatcher.calls == [] # 跳过存档不发工具请求
|
||||
assert out["bioflow"]["audit"]["tool_calls"] == 0
|
||||
@@ -0,0 +1,206 @@
|
||||
"""test_bioflow_agent_s4_supplement.py —— S4 覆盖清单补强测试。
|
||||
|
||||
覆盖:
|
||||
- _card_extension_action 六种 action 分发到对应扩展工具
|
||||
- 未知 action / 缺 ext_id/action 回退 None
|
||||
- confirm 携带 nl_triggers/parameters/runtime_image;EXT_STATUS 降级话术(确定性)
|
||||
- run 携带 params/files
|
||||
- 真实 configs/capabilities.yaml 三条目(含 bioflow-extension)可解析
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent.parent / "src"))
|
||||
|
||||
from rag.qa_v1_1_73 import bioflow_agent as ba # noqa: E402
|
||||
from rag.qa_v1_1_73.capabilities.manifests import CapabilityManifest, parse_capabilities_file # noqa: E402
|
||||
|
||||
EXT_TOOLS = (
|
||||
"create_extension", "extension_status", "build_extension", "trial_run_extension",
|
||||
"validate_extension", "confirm_extension", "publish_extension", "run_extension",
|
||||
"repair_extension",
|
||||
)
|
||||
|
||||
|
||||
def _ext_tool_manifests():
|
||||
out = {}
|
||||
for t in EXT_TOOLS:
|
||||
out[t] = CapabilityManifest(
|
||||
name="bioflow-extension", provider="mcp",
|
||||
config={"mcp": {"url": "http://127.0.0.1:8098/mcp/extension/", "tool": t}},
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
class FakeLLM:
|
||||
def complete(self, system, user, **kw):
|
||||
return "已收到。"
|
||||
|
||||
def complete_json(self, system, user, **kw):
|
||||
return {}
|
||||
|
||||
def complete_raw_messages(self, system, messages, tools=None, **kw):
|
||||
return {"content": "answer", "tool_calls": [], "finish_reason": "stop"}
|
||||
|
||||
|
||||
class FakeDispatcher:
|
||||
def __init__(self, results=None):
|
||||
self.results = results or {}
|
||||
self.calls = []
|
||||
|
||||
def execute(self, manifest, arguments):
|
||||
tool = manifest.config.get("mcp", {}).get("tool")
|
||||
self.calls.append((tool, arguments))
|
||||
return self.results.get(tool, '{"ok": true, "data": {"ext_id": "e1"}}')
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# _card_extension_action 分发
|
||||
# ===========================================================================
|
||||
|
||||
class TestExtensionAction:
|
||||
@pytest.mark.parametrize("action,expected_tool", [
|
||||
("build", "build_extension"),
|
||||
("trial_run", "trial_run_extension"),
|
||||
("validate", "validate_extension"),
|
||||
("confirm", "confirm_extension"),
|
||||
("publish", "publish_extension"),
|
||||
("run", "run_extension"),
|
||||
("repair", "repair_extension"),
|
||||
])
|
||||
def test_action_dispatches_to_tool(self, action, expected_tool):
|
||||
dispatcher = FakeDispatcher()
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "extension_action", "payload": {"ext_id": "e1", "action": action}},
|
||||
FakeLLM(), {}, _ext_tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
assert dispatcher.calls[0][0] == expected_tool
|
||||
assert dispatcher.calls[0][1] == {"ext_id": "e1"}
|
||||
assert out["bioflow"]["audit"]["tool_names"] == [expected_tool]
|
||||
|
||||
def test_unknown_action_returns_none(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "extension_action", "payload": {"ext_id": "e1", "action": "bogus"}},
|
||||
FakeLLM(), {}, _ext_tool_manifests(),
|
||||
)
|
||||
assert out is None
|
||||
|
||||
def test_missing_ext_id_or_action_returns_none(self):
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=FakeDispatcher()):
|
||||
assert ba._handle_card_result(
|
||||
{"card": "extension_action", "payload": {"action": "build"}},
|
||||
FakeLLM(), {}, _ext_tool_manifests(),
|
||||
) is None
|
||||
assert ba._handle_card_result(
|
||||
{"card": "extension_action", "payload": {"ext_id": "e1"}},
|
||||
FakeLLM(), {}, _ext_tool_manifests(),
|
||||
) is None
|
||||
|
||||
def test_confirm_passes_triggers_params_image(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
payload = {"ext_id": "e1", "action": "confirm",
|
||||
"nl_triggers": ["扩展", "触发词"], "parameters": [{"name": "x"}],
|
||||
"runtime_image": "img"}
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
ba._handle_card_result({"card": "extension_action", "payload": payload},
|
||||
FakeLLM(), {}, _ext_tool_manifests())
|
||||
args = dispatcher.calls[0][1]
|
||||
assert args == {"ext_id": "e1", "nl_triggers": ["扩展", "触发词"],
|
||||
"parameters": [{"name": "x"}], "runtime_image": "img"}
|
||||
|
||||
def test_confirm_status_insufficient_uses_deterministic_script(self):
|
||||
# confirm 返回 EXT_STATUS(状态不足)→ 确定性降级话术(不报 500、不经 LLM)
|
||||
dispatcher = FakeDispatcher({
|
||||
"confirm_extension": '{"ok": false, "code": "EXT_STATUS", "message": "状态不可确认"}',
|
||||
})
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "extension_action", "payload": {"ext_id": "e1", "action": "confirm"}},
|
||||
FakeLLM(), {}, _ext_tool_manifests(),
|
||||
)
|
||||
assert out is not None
|
||||
assert "还不能确认" in out["answer"]
|
||||
assert "先完成" in out["answer"] # 确定性话术(非 LLM 转述)
|
||||
|
||||
def test_run_passes_params_and_files(self):
|
||||
dispatcher = FakeDispatcher()
|
||||
payload = {"ext_id": "e1", "action": "run",
|
||||
"params": {"cpu": 8}, "files": ["/data/a.fa"]}
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
ba._handle_card_result({"card": "extension_action", "payload": payload},
|
||||
FakeLLM(), {}, _ext_tool_manifests())
|
||||
assert dispatcher.calls[0][1] == {"ext_id": "e1", "params": {"cpu": 8}, "files": ["/data/a.fa"]}
|
||||
|
||||
def test_extension_tool_not_in_manifests_returns_none(self):
|
||||
# tool_manifests 里没有扩展工具 → 回退 None(能力未安装/未发现工具)
|
||||
empty = CapabilityManifest(name="bioflow-gene-family", provider="mcp",
|
||||
config={"mcp": {"url": "http://x", "tool": "job_status"}})
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=FakeDispatcher()):
|
||||
out = ba._handle_card_result(
|
||||
{"card": "extension_action", "payload": {"ext_id": "e1", "action": "build"}},
|
||||
FakeLLM(), {}, {"job_status": empty},
|
||||
)
|
||||
assert out is None
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 真实 capabilities.yaml 三条目可解析
|
||||
# ===========================================================================
|
||||
|
||||
class TestCapabilitiesThreeEntries:
|
||||
def _real_yaml(self) -> Path:
|
||||
return Path(__file__).resolve().parents[5] / "agent" / "qa" / "configs" / "capabilities.yaml"
|
||||
|
||||
def test_three_bioflow_entries(self):
|
||||
manifests, warnings = parse_capabilities_file(self._real_yaml())
|
||||
by_name = {m.name: m for m in manifests}
|
||||
assert {"bioflow-gene-family", "bioflow-rnaseq", "bioflow-extension"} <= set(by_name)
|
||||
|
||||
ext = by_name["bioflow-extension"]
|
||||
assert ext.provider == "mcp"
|
||||
assert ext.permission == "read_only"
|
||||
assert ext.nl_triggers
|
||||
assert ext.config["mcp"]["url"].endswith("/mcp/extension/")
|
||||
assert ext.config["prompt_pack"]["file"] == "bioflow_extension_agent_rules.md"
|
||||
|
||||
def test_bioflow_extension_in_capability_names(self):
|
||||
assert "bioflow-extension" in ba.BIOFLOW_CAPABILITY_NAMES
|
||||
assert ba.CARD_EXTENSION_ACTION == "extension_action"
|
||||
assert ba.CARD_EXTENSION_ACTION in ba.CARD_TYPES
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# S5:_route_card_result 对 extension_action 只发现 extension 能力(一次 list_tools)
|
||||
# ===========================================================================
|
||||
|
||||
class TestRouteExtensionOptimization:
|
||||
def test_extension_action_discovers_only_extension_capability(self):
|
||||
gf = CapabilityManifest(name="bioflow-gene-family", provider="mcp",
|
||||
config={"mcp": {"url": "http://x/gene-family/"}})
|
||||
rn = CapabilityManifest(name="bioflow-rnaseq", provider="mcp",
|
||||
config={"mcp": {"url": "http://x/rnaseq/"}})
|
||||
ext = CapabilityManifest(name="bioflow-extension", provider="mcp",
|
||||
config={"mcp": {"url": "http://x/extension/"}})
|
||||
captured = {}
|
||||
|
||||
def fake_build(manifests):
|
||||
captured["names"] = [m.name for m in manifests]
|
||||
return ([], {})
|
||||
|
||||
with mock.patch.object(ba, "_registered_capabilities", return_value=[gf, rn, ext]):
|
||||
with mock.patch.object(ba, "_build_tools", side_effect=fake_build):
|
||||
ba._route_card_result(
|
||||
{"card": "extension_action", "payload": {"ext_id": "e1", "action": "build"}},
|
||||
FakeLLM(), None, None, None,
|
||||
)
|
||||
# 只把 bioflow-extension manifest 交给 _build_tools(避免对三能力各 list_tools 一次)
|
||||
assert captured["names"] == ["bioflow-extension"]
|
||||
@@ -0,0 +1,116 @@
|
||||
"""test_bioflow_agent_s6_config.py —— S6 多配置回归(QA 单元层)。
|
||||
|
||||
模拟四种部署配置,断言候选能力集合 + 意图预筛结果正确(真实 E2E/金标 20 题依赖部署,
|
||||
此处以 monkeypatch/env 在单元层覆盖漏判安全语义)。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent.parent / "src"))
|
||||
|
||||
from rag.qa_v1_1_73 import bioflow_agent as ba # noqa: E402
|
||||
from rag.qa_v1_1_73.capabilities.registry import CapabilityRegistry # noqa: E402
|
||||
|
||||
THREE_YAML = """
|
||||
capabilities:
|
||||
- name: bioflow-gene-family
|
||||
provider: mcp
|
||||
permission: read_only
|
||||
session_default: true
|
||||
nl_triggers: [基因家族, 亚家族]
|
||||
mcp: {url: "http://x/gene-family/"}
|
||||
- name: bioflow-rnaseq
|
||||
provider: mcp
|
||||
permission: read_only
|
||||
session_default: true
|
||||
nl_triggers: [转录组, 差异表达]
|
||||
mcp: {url: "http://x/rnaseq/"}
|
||||
- name: bioflow-extension
|
||||
provider: mcp
|
||||
permission: read_only
|
||||
session_default: true
|
||||
nl_triggers: [自定义扩展]
|
||||
mcp: {url: "http://x/extension/"}
|
||||
"""
|
||||
|
||||
|
||||
def _registry(tmp_path, yaml_text=THREE_YAML):
|
||||
path = tmp_path / "caps.yaml"
|
||||
path.write_text(yaml_text, encoding="utf-8")
|
||||
reg = CapabilityRegistry(path)
|
||||
reg.load()
|
||||
return reg
|
||||
|
||||
|
||||
class FakeLLM:
|
||||
def complete_raw_messages(self, *a, **kw):
|
||||
return {"content": "x", "tool_calls": [], "finish_reason": "stop"}
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 配置 1:全关(QA_V1_ENABLE_BIOFLOW_AGENT=0 + 三能力全卸载)
|
||||
# ===========================================================================
|
||||
|
||||
class TestConfigAllOff:
|
||||
def test_fuse_off_returns_none(self, tmp_path):
|
||||
reg = _registry(tmp_path)
|
||||
with mock.patch.object(ba, "get_registry", return_value=reg):
|
||||
with mock.patch.dict(ba.os.environ, {}, clear=False):
|
||||
ba.os.environ.pop(ba.ENV_FUSE, None)
|
||||
out = ba.try_bioflow_agent("帮我做基因家族分析", llm_client=FakeLLM())
|
||||
assert out is None # 走原 pipeline
|
||||
|
||||
def test_all_uninstalled_no_candidate(self, tmp_path):
|
||||
reg = _registry(tmp_path, "capabilities: []\n") # 三能力全卸载
|
||||
with mock.patch.object(ba, "get_registry", return_value=reg):
|
||||
with mock.patch.dict(ba.os.environ, {ba.ENV_FUSE: "1"}):
|
||||
assert ba._candidate_capabilities("帮我做基因家族分析", None) == []
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 配置 2:只开基因家族
|
||||
# ===========================================================================
|
||||
|
||||
class TestConfigGeneFamilyOnly:
|
||||
def test_gene_family_hit_rnaseq_miss(self, tmp_path):
|
||||
reg = _registry(tmp_path)
|
||||
caps = {"bioflow-gene-family": True, "bioflow-rnaseq": False, "bioflow-extension": False}
|
||||
with mock.patch.object(ba, "get_registry", return_value=reg):
|
||||
with mock.patch.dict(ba.os.environ, {ba.ENV_FUSE: "1"}):
|
||||
names = [m.name for m in ba._candidate_capabilities("帮我做水稻 SPL 基因家族分析", caps)]
|
||||
assert names == ["bioflow-gene-family"]
|
||||
assert ba._candidate_capabilities("帮我做转录组差异表达分析", caps) == []
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 配置 3:只开转录组
|
||||
# ===========================================================================
|
||||
|
||||
class TestConfigRnaseqOnly:
|
||||
def test_rnaseq_hit_gene_family_miss(self, tmp_path):
|
||||
reg = _registry(tmp_path)
|
||||
caps = {"bioflow-gene-family": False, "bioflow-rnaseq": True, "bioflow-extension": False}
|
||||
with mock.patch.object(ba, "get_registry", return_value=reg):
|
||||
with mock.patch.dict(ba.os.environ, {ba.ENV_FUSE: "1"}):
|
||||
names = [m.name for m in ba._candidate_capabilities("转录组差异表达分析", caps)]
|
||||
assert names == ["bioflow-rnaseq"]
|
||||
assert ba._candidate_capabilities("帮我做水稻 SPL 基因家族分析", caps) == []
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 配置 4:全开但非生信提问
|
||||
# ===========================================================================
|
||||
|
||||
class TestConfigAllOnNonBioflow:
|
||||
def test_non_bioflow_question_no_candidate(self, tmp_path):
|
||||
reg = _registry(tmp_path)
|
||||
caps = {"bioflow-gene-family": True, "bioflow-rnaseq": True, "bioflow-extension": True}
|
||||
with mock.patch.object(ba, "get_registry", return_value=reg):
|
||||
with mock.patch.dict(ba.os.environ, {ba.ENV_FUSE: "1"}):
|
||||
assert ba._candidate_capabilities("光合作用的电子传递链是什么", caps) == []
|
||||
assert ba._candidate_capabilities("什么是转录因子", caps) == []
|
||||
@@ -0,0 +1,496 @@
|
||||
"""test_bioflow_agent_supplement.py —— S1 覆盖清单补强测试。
|
||||
|
||||
补强 dev 的 test_bioflow_agent.py 未覆盖的分支(对照 S1 覆盖清单):
|
||||
- 熔断开关值解析(1/true/yes/on 与 0/false/off/空/缺失)
|
||||
- nl_triggers 同时命中 → 两个候选(交 LLM 依 description 选)
|
||||
- session_default 回退(True/False)/ registry 访问异常 / registry.get 异常
|
||||
- _build_tools:MCP 工具发现(空 url 跳过 / 发现异常 / 共享连接按能力名 / 无名跳过 / bad schema 兜底)
|
||||
- _build_system:prompt_pack 注入 / 去重 / 加载异常兜底
|
||||
- try_bioflow_agent:tools 空→None、answer 空→None、progress_callback 触发与吞异常
|
||||
- ToolRunner max_rounds=5 / bioflow_agent.MAX_ROUNDS 常量
|
||||
- make_llm_complete 包装(正常 / 异常)
|
||||
- 真实 configs/capabilities.yaml 两条目可解析
|
||||
- bioflow_agent 不 import pipeline/answer_composer(红线静态守卫)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent.parent / "src"))
|
||||
|
||||
from rag.qa_v1_1_73 import bioflow_agent as ba # noqa: E402
|
||||
from rag.qa_v1_1_73.capabilities.manifests import ( # noqa: E402
|
||||
CapabilityManifest,
|
||||
parse_capabilities_file,
|
||||
)
|
||||
from rag.qa_v1_1_73.capabilities.registry import CapabilityRegistry # noqa: E402
|
||||
|
||||
BIOFLOW_YAML = """
|
||||
capabilities:
|
||||
- name: bioflow-gene-family
|
||||
provider: mcp
|
||||
permission: read_only
|
||||
description: 基因家族分析
|
||||
session_default: true
|
||||
nl_triggers: [基因家族, 家族成员, 亚家族, 建树]
|
||||
mcp:
|
||||
url: http://127.0.0.1:8098/mcp/gene-family/
|
||||
prompt_pack:
|
||||
file: bioflow_gene_family_agent_rules.md
|
||||
- name: bioflow-rnaseq
|
||||
provider: mcp
|
||||
permission: read_only
|
||||
description: 转录组分析
|
||||
session_default: true
|
||||
nl_triggers: [转录组, RNA-seq, 差异表达, 富集]
|
||||
mcp:
|
||||
url: http://127.0.0.1:8098/mcp/rnaseq/
|
||||
prompt_pack:
|
||||
file: bioflow_rnaseq_agent_rules.md
|
||||
"""
|
||||
|
||||
|
||||
def _make_registry(tmp_path: Path, yaml_text: str = BIOFLOW_YAML) -> CapabilityRegistry:
|
||||
path = tmp_path / "capabilities.yaml"
|
||||
path.write_text(yaml_text, encoding="utf-8")
|
||||
registry = CapabilityRegistry(path)
|
||||
registry.load()
|
||||
return registry
|
||||
|
||||
|
||||
def _gf_manifest(**overrides) -> CapabilityManifest:
|
||||
cfg = {
|
||||
"mcp": {"url": "http://127.0.0.1:8098/mcp/gene-family/"},
|
||||
"prompt_pack": {"file": "bioflow_gene_family_agent_rules.md"},
|
||||
}
|
||||
cfg.update(overrides.pop("config", {}))
|
||||
return CapabilityManifest(
|
||||
name="bioflow-gene-family", provider="mcp", permission="read_only",
|
||||
description="基因家族分析", config=cfg, nl_triggers=["基因家族"], **overrides,
|
||||
)
|
||||
|
||||
|
||||
class FakeLLMClient:
|
||||
def __init__(self, content: str = "已为您规划分析路径。") -> None:
|
||||
self.content = content
|
||||
|
||||
def complete_raw_messages(self, system, messages, tools=None, **kwargs):
|
||||
return {"content": self.content, "tool_calls": [], "finish_reason": "stop"}
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 门 1:熔断开关值解析
|
||||
# ===========================================================================
|
||||
|
||||
class TestEnvFuse:
|
||||
@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on", " On ", "YES"])
|
||||
def test_truthy(self, value):
|
||||
with mock.patch.dict(ba.os.environ, {ba.ENV_FUSE: value}):
|
||||
assert ba._env_enabled() is True
|
||||
|
||||
@pytest.mark.parametrize("value", ["0", "false", "off", "", "no", " "])
|
||||
def test_falsy(self, value):
|
||||
with mock.patch.dict(ba.os.environ, {ba.ENV_FUSE: value}):
|
||||
assert ba._env_enabled() is False
|
||||
|
||||
def test_missing_defaults_off(self):
|
||||
with mock.patch.dict(ba.os.environ, {}, clear=False):
|
||||
ba.os.environ.pop(ba.ENV_FUSE, None)
|
||||
assert ba._env_enabled() is False
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 门 3:nl_triggers 意图预筛
|
||||
# ===========================================================================
|
||||
|
||||
class TestIntentMatchEdge:
|
||||
def test_empty_triggers_never_match(self):
|
||||
assert ba._intent_match("做基因家族分析", []) is False
|
||||
assert ba._intent_match("做基因家族分析", [""]) is False
|
||||
|
||||
def test_case_insensitive_substring(self):
|
||||
assert ba._intent_match("RNA-SEQ 分析", ["rna-seq"]) is True
|
||||
assert ba._intent_match("ka/ks 选择压力", ["Ka/Ks"]) is True
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 门 2 + 门 3:候选能力(registry 异常 / session_default 回退 / 同时命中)
|
||||
# ===========================================================================
|
||||
|
||||
class TestCandidateCapabilities:
|
||||
def test_registry_access_exception_returns_empty(self):
|
||||
with mock.patch.object(ba, "get_registry", side_effect=RuntimeError("boom")):
|
||||
assert ba._candidate_capabilities("做基因家族", None) == []
|
||||
|
||||
def test_registry_get_exception_skips(self):
|
||||
class BadRegistry:
|
||||
def get(self, name):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
with mock.patch.object(ba, "get_registry", return_value=BadRegistry()):
|
||||
assert ba._candidate_capabilities("做基因家族", None) == []
|
||||
|
||||
def test_both_triggers_match_returns_both(self, tmp_path):
|
||||
registry = _make_registry(tmp_path)
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
caps = ba._candidate_capabilities("做基因家族和转录组差异表达分析", None)
|
||||
assert [m.name for m in caps] == ["bioflow-gene-family", "bioflow-rnaseq"]
|
||||
|
||||
def test_session_default_false_skips(self, tmp_path):
|
||||
yaml = BIOFLOW_YAML.replace("session_default: true", "session_default: false")
|
||||
registry = _make_registry(tmp_path, yaml)
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
# 无会话开关时回退 session_default=false → 不候选
|
||||
assert ba._candidate_capabilities("做基因家族分析", None) == []
|
||||
# 显式会话开关 true → 候选
|
||||
caps = ba._candidate_capabilities("做基因家族分析", {"bioflow-gene-family": True})
|
||||
assert [m.name for m in caps] == ["bioflow-gene-family"]
|
||||
|
||||
def test_session_switch_true_overrides_default_false(self, tmp_path):
|
||||
yaml = BIOFLOW_YAML.replace("session_default: true", "session_default: false", 1)
|
||||
registry = _make_registry(tmp_path, yaml)
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
caps = ba._candidate_capabilities("做基因家族分析", {"bioflow-gene-family": True})
|
||||
assert [m.name for m in caps] == ["bioflow-gene-family"]
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# _build_tools:MCP 工具发现
|
||||
# ===========================================================================
|
||||
|
||||
class _FakeDispatcher:
|
||||
def __init__(self, tools_by_name=None, exc=None):
|
||||
self.tools_by_name = tools_by_name or {}
|
||||
self.exc = exc
|
||||
|
||||
def list_mcp_tools(self, manifest):
|
||||
if self.exc:
|
||||
raise self.exc
|
||||
return self.tools_by_name.get(manifest.name, [])
|
||||
|
||||
|
||||
class TestBuildTools:
|
||||
def test_empty_url_skipped(self):
|
||||
manifest = CapabilityManifest(
|
||||
name="x", provider="mcp", config={"mcp": {"url": ""}}
|
||||
)
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=_FakeDispatcher()):
|
||||
tools, tool_manifests = ba._build_tools([manifest])
|
||||
assert tools == []
|
||||
assert tool_manifests == {}
|
||||
|
||||
def test_discovery_exception_skipped(self):
|
||||
manifest = _gf_manifest()
|
||||
with mock.patch.object(
|
||||
ba, "get_dispatcher",
|
||||
return_value=_FakeDispatcher(exc=RuntimeError("mcp down")),
|
||||
):
|
||||
tools, tool_manifests = ba._build_tools([manifest])
|
||||
assert tools == []
|
||||
assert tool_manifests == {}
|
||||
|
||||
def test_builds_schema_and_shared_connection(self):
|
||||
dispatcher = _FakeDispatcher({
|
||||
"bioflow-gene-family": [
|
||||
{"name": "plan_gene_family", "description": "规划路径",
|
||||
"input_schema": {"type": "object", "properties": {"message": {"type": "string"}}}},
|
||||
]
|
||||
})
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
tools, tool_manifests = ba._build_tools([_gf_manifest()])
|
||||
assert len(tools) == 1
|
||||
assert tools[0]["function"]["name"] == "plan_gene_family"
|
||||
assert tools[0]["function"]["description"] == "规划路径"
|
||||
# 共享连接:manifest.name 用能力名,mcp.tool 指向真实 server 工具名
|
||||
tm = tool_manifests["plan_gene_family"]
|
||||
assert tm.name == "bioflow-gene-family"
|
||||
assert tm.config["mcp"]["tool"] == "plan_gene_family"
|
||||
assert tm.config["mcp"]["url"].endswith("/mcp/gene-family/")
|
||||
|
||||
def test_skips_nameless_and_bad_schema(self):
|
||||
dispatcher = _FakeDispatcher({
|
||||
"bioflow-gene-family": [
|
||||
{"name": "", "description": "无名", "input_schema": {}},
|
||||
{"name": "job_status", "description": "查状态", "input_schema": "not-a-dict"},
|
||||
]
|
||||
})
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
tools, tool_manifests = ba._build_tools([_gf_manifest()])
|
||||
assert [t["function"]["name"] for t in tools] == ["job_status"]
|
||||
assert tools[0]["function"]["parameters"] == {"type": "object", "properties": {}}
|
||||
|
||||
def test_shared_tool_deduped_across_capabilities(self):
|
||||
"""两能力同时命中:共享工具去重(只保留首个能力),专有工具各自保留。
|
||||
|
||||
覆盖 bug ① 修复:job_status 等 12 个共享工具两端点语义一致,按名去重后
|
||||
LLM tools 列表无重名,且 job_status 路由到首个能力(gene-family)。
|
||||
"""
|
||||
dispatcher = _FakeDispatcher({
|
||||
"bioflow-gene-family": [
|
||||
{"name": "plan_gene_family", "description": "规划基因家族", "input_schema": {}},
|
||||
{"name": "job_status", "description": "查状态", "input_schema": {}},
|
||||
{"name": "confirm_stage", "description": "确认参数", "input_schema": {}},
|
||||
],
|
||||
"bioflow-rnaseq": [
|
||||
{"name": "plan_rnaseq", "description": "规划转录组", "input_schema": {}},
|
||||
{"name": "job_status", "description": "查状态", "input_schema": {}},
|
||||
{"name": "confirm_stage", "description": "确认参数", "input_schema": {}},
|
||||
],
|
||||
})
|
||||
rn = CapabilityManifest(
|
||||
name="bioflow-rnaseq", provider="mcp", permission="read_only",
|
||||
description="转录组分析", nl_triggers=["转录组"],
|
||||
config={"mcp": {"url": "http://127.0.0.1:8098/mcp/rnaseq/"}},
|
||||
)
|
||||
with mock.patch.object(ba, "get_dispatcher", return_value=dispatcher):
|
||||
tools, tool_manifests = ba._build_tools([_gf_manifest(), rn])
|
||||
|
||||
names = [t["function"]["name"] for t in tools]
|
||||
# 专有工具各保留 + 共享工具只保留一份,且无重名
|
||||
assert names == ["plan_gene_family", "job_status", "confirm_stage", "plan_rnaseq"]
|
||||
assert len(names) == len(set(names))
|
||||
|
||||
# 共享工具路由到首个能力(gene-family)
|
||||
assert tool_manifests["job_status"].name == "bioflow-gene-family"
|
||||
assert tool_manifests["job_status"].config["mcp"]["url"].endswith("/mcp/gene-family/")
|
||||
assert tool_manifests["confirm_stage"].name == "bioflow-gene-family"
|
||||
# 专有工具各自路由到自身能力
|
||||
assert tool_manifests["plan_rnaseq"].name == "bioflow-rnaseq"
|
||||
assert tool_manifests["plan_rnaseq"].config["mcp"]["url"].endswith("/mcp/rnaseq/")
|
||||
assert tool_manifests["plan_gene_family"].name == "bioflow-gene-family"
|
||||
|
||||
# ===========================================================================
|
||||
# _build_system:prompt_pack 注入
|
||||
# ===========================================================================
|
||||
|
||||
class TestBuildSystem:
|
||||
def test_appends_prompt_pack(self):
|
||||
manifest = _gf_manifest()
|
||||
with mock.patch.object(ba.PromptPackProvider, "load", return_value="基因家族规则正文"):
|
||||
system = ba._build_system([manifest])
|
||||
assert ba._SYSTEM_BASE in system
|
||||
assert "基因家族规则正文" in system
|
||||
|
||||
def test_dedup_same_file_once(self):
|
||||
m1 = _gf_manifest()
|
||||
m2 = _gf_manifest() # 同 file
|
||||
calls = []
|
||||
with mock.patch.object(ba.PromptPackProvider, "load", side_effect=lambda m, root: calls.append(m) or "RULES"):
|
||||
system = ba._build_system([m1, m2])
|
||||
assert len(calls) == 1 # 去重后只加载一次
|
||||
assert system.count("RULES") == 1
|
||||
|
||||
def test_empty_file_skipped(self):
|
||||
manifest = CapabilityManifest(name="x", provider="mcp", config={"prompt_pack": {"file": ""}})
|
||||
system = ba._build_system([manifest])
|
||||
assert system == ba._SYSTEM_BASE
|
||||
|
||||
def test_load_exception_falls_back(self):
|
||||
manifest = _gf_manifest()
|
||||
with mock.patch.object(ba.PromptPackProvider, "load", side_effect=RuntimeError("boom")):
|
||||
system = ba._build_system([manifest])
|
||||
assert ba._SYSTEM_BASE in system # 不炸,只用基础提示
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# try_bioflow_agent:tools 空 / answer 空 / progress_callback
|
||||
# ===========================================================================
|
||||
|
||||
class TestFullFlowEdge:
|
||||
def test_tools_empty_returns_none(self, tmp_path):
|
||||
registry = _make_registry(tmp_path)
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
with mock.patch.dict(ba.os.environ, {ba.ENV_FUSE: "1"}):
|
||||
with mock.patch.object(ba, "_build_tools", return_value=([], {})):
|
||||
result = ba.try_bioflow_agent(
|
||||
"帮我做基因家族分析", llm_client=FakeLLMClient(), capabilities=None
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_answer_empty_returns_none(self, tmp_path):
|
||||
registry = _make_registry(tmp_path)
|
||||
fake_tools = [{"type": "function", "function": {"name": "plan_gene_family",
|
||||
"description": "x",
|
||||
"parameters": {"type": "object", "properties": {}}}}]
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
with mock.patch.dict(ba.os.environ, {ba.ENV_FUSE: "1"}):
|
||||
with mock.patch.object(ba, "_build_tools", return_value=(fake_tools, {})):
|
||||
with mock.patch.object(ba, "ToolRunner") as FakeRunner:
|
||||
FakeRunner.return_value.run.return_value = ("", {"rounds": 1})
|
||||
result = ba.try_bioflow_agent(
|
||||
"帮我做基因家族分析", llm_client=FakeLLMClient(), capabilities=None
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_progress_callback_invoked(self, tmp_path):
|
||||
registry = _make_registry(tmp_path)
|
||||
fake_tools = [{"type": "function", "function": {"name": "plan_gene_family",
|
||||
"description": "x",
|
||||
"parameters": {"type": "object", "properties": {}}}}]
|
||||
calls = []
|
||||
|
||||
def cb(payload):
|
||||
calls.append(payload)
|
||||
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
with mock.patch.dict(ba.os.environ, {ba.ENV_FUSE: "1"}):
|
||||
with mock.patch.object(ba, "_build_tools", return_value=(fake_tools, {})):
|
||||
with mock.patch.object(ba, "ToolRunner") as FakeRunner:
|
||||
FakeRunner.return_value.run.return_value = ("已答复", {"rounds": 1})
|
||||
result = ba.try_bioflow_agent(
|
||||
"帮我做基因家族分析", llm_client=FakeLLMClient(),
|
||||
capabilities=None, progress_callback=cb,
|
||||
)
|
||||
assert result is not None
|
||||
assert calls and calls[0]["stage"] == "bioflow_agent_started"
|
||||
assert calls[0]["capabilities"] == ["bioflow-gene-family"]
|
||||
assert calls[0]["tool_count"] == 1
|
||||
|
||||
def test_progress_callback_exception_swallowed(self, tmp_path):
|
||||
registry = _make_registry(tmp_path)
|
||||
fake_tools = [{"type": "function", "function": {"name": "plan_gene_family",
|
||||
"description": "x",
|
||||
"parameters": {"type": "object", "properties": {}}}}]
|
||||
|
||||
def cb(payload):
|
||||
raise RuntimeError("callback boom")
|
||||
|
||||
with mock.patch.object(ba, "get_registry", return_value=registry):
|
||||
with mock.patch.dict(ba.os.environ, {ba.ENV_FUSE: "1"}):
|
||||
with mock.patch.object(ba, "_build_tools", return_value=(fake_tools, {})):
|
||||
with mock.patch.object(ba, "ToolRunner") as FakeRunner:
|
||||
FakeRunner.return_value.run.return_value = ("已答复", {"rounds": 1})
|
||||
result = ba.try_bioflow_agent(
|
||||
"帮我做基因家族分析", llm_client=FakeLLMClient(),
|
||||
capabilities=None, progress_callback=cb,
|
||||
)
|
||||
assert result is not None # 回调异常被吞,不影响主流程
|
||||
assert result["answer"] == "已答复"
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# ToolRunner max_rounds=5 / MAX_ROUNDS 常量
|
||||
# ===========================================================================
|
||||
|
||||
class TestMaxRounds:
|
||||
def test_max_rounds_constant_is_5(self):
|
||||
assert ba.MAX_ROUNDS == 5
|
||||
|
||||
def test_tool_runner_respects_max_rounds_5(self):
|
||||
from rag.qa_v1_1_73.tool_runner import ToolRunner
|
||||
|
||||
class LLM:
|
||||
def __call__(self, system, messages, tools):
|
||||
return "", [{"id": "c", "name": "echo", "arguments": "{}"}], None
|
||||
|
||||
class Dispatcher:
|
||||
def execute(self, manifest, arguments):
|
||||
return "ok"
|
||||
|
||||
manifest = CapabilityManifest(name="echo", provider="python")
|
||||
tools = [{"type": "function", "function": {"name": "echo", "description": "e",
|
||||
"parameters": {"type": "object", "properties": {}}}}]
|
||||
runner = ToolRunner(LLM(), dispatcher=Dispatcher(), max_rounds=5)
|
||||
content, audit = runner.run("sys", "user", tools, {"echo": manifest})
|
||||
assert audit["rounds"] == 5
|
||||
assert audit["tool_calls"] == 5
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# make_llm_complete 包装
|
||||
# ===========================================================================
|
||||
|
||||
class TestMakeLLMComplete:
|
||||
def test_normal(self):
|
||||
from rag.qa_v1_1_73.tool_runner import make_llm_complete
|
||||
|
||||
class Client:
|
||||
def complete_raw_messages(self, system, messages, tools=None, **kw):
|
||||
return {"content": "hi", "tool_calls": [{"name": "x"}], "finish_reason": "stop"}
|
||||
|
||||
complete = make_llm_complete(Client())
|
||||
content, tool_calls, error = complete("sys", [{"role": "user", "content": "q"}], [])
|
||||
assert content == "hi"
|
||||
assert tool_calls == [{"name": "x"}]
|
||||
assert error is None
|
||||
|
||||
def test_exception_becomes_error(self):
|
||||
from rag.qa_v1_1_73.tool_runner import make_llm_complete
|
||||
|
||||
class Client:
|
||||
def complete_raw_messages(self, *a, **kw):
|
||||
raise RuntimeError("llm down")
|
||||
|
||||
complete = make_llm_complete(Client())
|
||||
content, tool_calls, error = complete("sys", [], [])
|
||||
assert content == ""
|
||||
assert tool_calls == []
|
||||
assert "RuntimeError" in error
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 真实 configs/capabilities.yaml 两条目可解析
|
||||
# ===========================================================================
|
||||
|
||||
class TestRealCapabilitiesYaml:
|
||||
def _real_yaml(self) -> Path:
|
||||
repo_root = Path(__file__).resolve().parents[5]
|
||||
return repo_root / "agent" / "qa" / "configs" / "capabilities.yaml"
|
||||
|
||||
def test_both_bioflow_entries_parse(self):
|
||||
manifests, warnings = parse_capabilities_file(self._real_yaml())
|
||||
by_name = {m.name: m for m in manifests}
|
||||
assert "bioflow-gene-family" in by_name
|
||||
assert "bioflow-rnaseq" in by_name
|
||||
|
||||
gf = by_name["bioflow-gene-family"]
|
||||
assert gf.provider == "mcp"
|
||||
assert gf.permission == "read_only"
|
||||
assert gf.nl_triggers
|
||||
assert gf.config["mcp"]["url"].endswith("/mcp/gene-family/")
|
||||
assert gf.config["prompt_pack"]["file"] == "bioflow_gene_family_agent_rules.md"
|
||||
|
||||
rn = by_name["bioflow-rnaseq"]
|
||||
assert rn.permission == "read_only"
|
||||
assert rn.nl_triggers
|
||||
assert rn.config["mcp"]["url"].endswith("/mcp/rnaseq/")
|
||||
assert rn.config["prompt_pack"]["file"] == "bioflow_rnaseq_agent_rules.md"
|
||||
|
||||
def test_nl_triggers_cover_expected_words(self):
|
||||
manifests, _ = parse_capabilities_file(self._real_yaml())
|
||||
gf = next(m for m in manifests if m.name == "bioflow-gene-family")
|
||||
rn = next(m for m in manifests if m.name == "bioflow-rnaseq")
|
||||
# 关键触发词在位(意图预筛依赖)
|
||||
assert "基因家族" in gf.nl_triggers
|
||||
assert "转录组" in rn.nl_triggers
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 红线静态守卫:bioflow_agent 不依赖 pipeline / answer_composer
|
||||
# ===========================================================================
|
||||
|
||||
class TestRedLineNoPipelineDep:
|
||||
def test_no_pipeline_or_answer_composer_import(self):
|
||||
import ast
|
||||
|
||||
src = Path(ba.__file__).read_text(encoding="utf-8")
|
||||
tree = ast.parse(src)
|
||||
imported = []
|
||||
for node in tree.body:
|
||||
if isinstance(node, ast.Import):
|
||||
imported += [a.name for a in node.names]
|
||||
elif isinstance(node, ast.ImportFrom):
|
||||
imported.append(node.module or "")
|
||||
assert not any("pipeline" in m or "answer_composer" in m for m in imported)
|
||||
|
||||
def test_never_raises(self, tmp_path):
|
||||
# 任何异常都返回 None/正常 payload,绝不外抛
|
||||
with mock.patch.object(ba, "get_registry", side_effect=RuntimeError("boom")):
|
||||
with mock.patch.dict(ba.os.environ, {ba.ENV_FUSE: "1"}):
|
||||
result = ba.try_bioflow_agent("做基因家族", llm_client=FakeLLMClient())
|
||||
assert result is None
|
||||
@@ -0,0 +1,335 @@
|
||||
"""test_score_gate.py —— 相关性下限单测(缺陷 #28 回改 v2,2026-09-11)。
|
||||
|
||||
运行(venv 内,源码树):
|
||||
cd agent/qa/src/rag/qa_v1_1_73
|
||||
python -m pytest test_score_gate.py -v
|
||||
|
||||
覆盖:
|
||||
- search(retrieval_only)模式:无 LLM 调用、无 rerank 调用;
|
||||
vector 域 0.45 预筛滤掉 0.32 垃圾、保留 0.63 正常;
|
||||
sources 的 score 字段取 vector_score(final_score=vector_score 域)。
|
||||
- chat 模式:final 域 0.25 后置过滤语义不变(5 高+25 低 -> 5、30 低 -> 0);
|
||||
Qdrant score_threshold 不传(None,保召回)。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# 本文件位于 agent/qa/src/rag/qa_v1_1_73/,parents[2] = agent/qa/src
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
|
||||
from rag.qa_v1_1_73 import score_gate # noqa: E402
|
||||
from rag.qa_v1_1_73.schemas import EvidenceItem # noqa: E402
|
||||
from rag.qa_v1_1_73.pipeline import QAV1Pipeline # noqa: E402
|
||||
|
||||
|
||||
def _item(
|
||||
idx: int, *, final: float = 0.0, vector: float = 0.0, graph: bool = False, key: str = ""
|
||||
) -> EvidenceItem:
|
||||
return EvidenceItem(
|
||||
evidence_id=f"E{idx}",
|
||||
kb="V16",
|
||||
source=f"src{idx}",
|
||||
relation="has",
|
||||
target=f"tgt{idx}",
|
||||
evidence=f"evidence body {idx}",
|
||||
relation_id=(key or f"rid-{idx}"),
|
||||
retrieval_methods=(["neo4j_bounded_expand"] if graph else ["qdrant_vector"]),
|
||||
vector_score=vector,
|
||||
final_score=final,
|
||||
)
|
||||
|
||||
|
||||
def _select(items, *, top_k=30, graph_min=0):
|
||||
selected, audit = QAV1Pipeline._select_graph_balanced(
|
||||
items, top_k_for_llm=top_k, graph_min_evidence=graph_min
|
||||
)
|
||||
return selected, audit
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# chat 模式 final 域 0.25(r4-28 语义不变)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_chat_all_below_source_threshold_selects_zero():
|
||||
"""30 条 final 全部低于 0.25 -> 选入 0。"""
|
||||
items = [_item(i, final=0.20) for i in range(30)]
|
||||
selected, audit = _select(items)
|
||||
assert selected == []
|
||||
assert audit["below_min_source_score"] == 30
|
||||
assert audit["min_source_score"] == score_gate.min_source_score()
|
||||
|
||||
|
||||
def test_chat_five_high_twentyfive_low_selects_only_five():
|
||||
"""5 高(final>=0.25)+ 25 低 -> 仅选入 5(保 rerank 序)。"""
|
||||
items = [_item(i, final=0.90 - i * 0.01) for i in range(5)]
|
||||
items += [_item(5 + i, final=0.20) for i in range(25)]
|
||||
selected, audit = _select(items)
|
||||
assert len(selected) == 5
|
||||
assert [item.evidence_id for item in selected] == [f"E{i}" for i in range(5)]
|
||||
assert audit["below_min_source_score"] == 25
|
||||
|
||||
|
||||
def test_chat_below_threshold_graph_candidates_excluded():
|
||||
"""低分 graph 候选不占用 graph 配额。"""
|
||||
items = [
|
||||
_item(1, final=0.30, graph=True),
|
||||
_item(2, final=0.21, graph=True),
|
||||
_item(3, final=0.95, graph=False),
|
||||
_item(4, final=0.19, graph=False),
|
||||
]
|
||||
selected, audit = _select(items, top_k=30, graph_min=1)
|
||||
assert [item.evidence_id for item in selected] == ["E1", "E3"]
|
||||
assert audit["graph_selected"] == 1
|
||||
assert audit["below_min_source_score"] == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# search / retrieval-only 模式:免 LLM、免 rerank、vector 域 0.45 预筛
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _StubLLM:
|
||||
available = True # 可用也必须不调用(回改核心断言)
|
||||
|
||||
def complete_json(self, *args, **kwargs):
|
||||
raise AssertionError("search 模式不得调用 LLM")
|
||||
|
||||
def complete(self, *args, **kwargs):
|
||||
raise AssertionError("search 模式不得调用 LLM")
|
||||
|
||||
|
||||
class _StubDirectRetriever:
|
||||
"""捕获 retrieve 入参,供断言 Qdrant score_threshold 传递。"""
|
||||
|
||||
def __init__(self, items):
|
||||
self.items = items
|
||||
self.last_kwargs = {}
|
||||
|
||||
def retrieve(self, question, queries, hypothesis, **kwargs):
|
||||
self.last_kwargs = kwargs
|
||||
return list(self.items), {
|
||||
"qdrant": [],
|
||||
"seed_extraction": {},
|
||||
"neo4j": [],
|
||||
"dedup": {"qdrant_items": len(self.items)},
|
||||
"options": dict(kwargs),
|
||||
}
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
|
||||
def _make_pipeline(graph_schema):
|
||||
from rag.qa_v1_1_73.config import QAConfig
|
||||
|
||||
cfg = QAConfig.from_env()
|
||||
if graph_schema == "legacy":
|
||||
object.__setattr__(cfg, "graph_schema", "legacy")
|
||||
import threading
|
||||
|
||||
p = QAV1Pipeline.__new__(QAV1Pipeline)
|
||||
p.config = cfg
|
||||
p._reranker_lock = threading.Lock()
|
||||
p.llm = _StubLLM()
|
||||
p.provisional = _RaisingComposer("provisional")
|
||||
p.hypothesis_extractor = _RaisingComposer("hypothesis")
|
||||
p.query_builder = _StubQueryBuilder()
|
||||
p.reranker = _StubReranker()
|
||||
p.agentic_controller = None
|
||||
p.question_decomposer = None
|
||||
p.subquestion_orchestrator = None
|
||||
p._retriever_stub = _StubDirectRetriever(ITEMS)
|
||||
p.direct_retriever_factory = lambda: p._retriever_stub
|
||||
p.audit = _NullAudit()
|
||||
p.domain = _StubDomain()
|
||||
return p
|
||||
|
||||
|
||||
class _RaisingComposer:
|
||||
def __init__(self, name):
|
||||
self.name = name
|
||||
|
||||
def generate(self, *args, **kwargs):
|
||||
raise AssertionError(f"{self.name} 不得被调用")
|
||||
|
||||
def extract(self, *args, **kwargs):
|
||||
raise AssertionError(f"{self.name} 不得被调用")
|
||||
|
||||
|
||||
class _StubQueryBuilder:
|
||||
def build(self, question, provisional_answer, hypothesis):
|
||||
return {"hypothesis_lane": [], "direct": [], "concept": []}
|
||||
|
||||
|
||||
class _StubReranker:
|
||||
def rerank(self, question, items):
|
||||
raise AssertionError("search 模式不得调用 rerank")
|
||||
|
||||
|
||||
class _NullAudit:
|
||||
def new_query_id(self):
|
||||
return "qid"
|
||||
|
||||
def write(self, query_id, payload):
|
||||
return None
|
||||
|
||||
|
||||
class _StubDomain:
|
||||
def normalize_hypothesis(self, hypothesis):
|
||||
return {"entities": []}
|
||||
|
||||
def direct_entities_from_question(self, question):
|
||||
return []
|
||||
|
||||
|
||||
# 集成(answer())路径:search 免 rerank,final_score = vector_score(同域),
|
||||
# 按 TL 实测向量域构造:垃圾 0.32(< 0.45 滤掉)、正常 0.63(保留)。
|
||||
JUNK = _item(1, vector=0.32)
|
||||
GOOD = _item(2, vector=0.63)
|
||||
ITEMS = [JUNK, GOOD]
|
||||
|
||||
|
||||
def test_search_mode_no_llm_no_rerank_keeps_good_only():
|
||||
p = _make_pipeline("legacy")
|
||||
out = p.answer(
|
||||
"景甜是谁",
|
||||
retrieval_only=True,
|
||||
native_chain_enabled=False,
|
||||
progress_callback=None,
|
||||
)
|
||||
assert out["sources"] != [] # vector 0.63 保留
|
||||
titles = [src["source"] for src in out["sources"]]
|
||||
assert "src2" in titles and "src1" not in titles # 0.32 垃圾被 0.45 预筛滤掉
|
||||
assert p._retriever_stub.last_kwargs["qdrant_score_threshold"] == 0.45
|
||||
assert out["retrieval_audit"]["options"]["qdrant_score_threshold"] == 0.45
|
||||
assert out["query_plan"]["agentic_plan"]["plan_source"] == "deterministic_retrieval_only_plan"
|
||||
assert out["degraded_execution"] is False
|
||||
# 跳过 rerank 后:final_score = vector_score(TL 回改语义,与 Qdrant 阈值
|
||||
# 同域);rerank_score 未精排(0.0)。BFF normalize 的 score 取 final_score。
|
||||
for src in out["sources"]:
|
||||
assert src["rerank_score"] == 0.0
|
||||
assert src["final_score"] == src["vector_score"]
|
||||
|
||||
|
||||
def test_search_mode_use_neo4j_forced_off():
|
||||
"""回改 v3:search 模式 retriever 收到的检索选项 use_neo4j=False
|
||||
(r3 语义补完——图证据种子检索在 13M 边库上 150s+,检索模式只走 Qdrant)。
|
||||
"""
|
||||
p = _make_pipeline("legacy")
|
||||
out = p.answer(
|
||||
"arginine biosynthesis",
|
||||
retrieval_only=True,
|
||||
native_chain_enabled=False,
|
||||
progress_callback=None,
|
||||
)
|
||||
kwargs = p._retriever_stub.last_kwargs
|
||||
assert kwargs["use_neo4j"] is False
|
||||
assert kwargs["use_graph_expansion"] is False # 依赖 use_neo4j 一并关
|
||||
assert kwargs["use_qdrant"] is True
|
||||
# 即使 ablation 显式开图检索,检索模式也不得重开(性能红线)
|
||||
p2 = _make_pipeline("legacy")
|
||||
out2 = p2.answer(
|
||||
"arginine biosynthesis",
|
||||
retrieval_only=True,
|
||||
native_chain_enabled=False,
|
||||
ablation={"retrieval": {"use_neo4j": True, "use_graph_expansion": True}},
|
||||
progress_callback=None,
|
||||
)
|
||||
kwargs2 = p2._retriever_stub.last_kwargs
|
||||
assert kwargs2["use_neo4j"] is False
|
||||
assert kwargs2["use_graph_expansion"] is False
|
||||
assert out2["success"] is True
|
||||
|
||||
|
||||
def test_chat_mode_passes_no_threshold():
|
||||
p = _make_pipeline("legacy")
|
||||
object.__setattr__(p.config, "enable_bge_rerank", False)
|
||||
out = p.answer(
|
||||
"景甜是谁",
|
||||
retrieval_only=False,
|
||||
native_chain_enabled=False,
|
||||
progress_callback=None,
|
||||
)
|
||||
assert p._retriever_stub.last_kwargs["qdrant_score_threshold"] is None
|
||||
|
||||
|
||||
def test_direct_flow_retriever_threshold_passthrough():
|
||||
"""真实 retriever 的 qdrant_score_threshold 参数链路(不触发 IO)。
|
||||
|
||||
用对象伪造构造 DirectFlowRetriever 实例,直接调用 qdrant_search:
|
||||
- 传阈值时透传给 qdrant.query_points(同域向量阈值);
|
||||
- 不传阈值时 score_threshold=None(chat 保召回)。
|
||||
"""
|
||||
import rag.qa_v1_1_73.direct_flow_retriever as dfr_mod
|
||||
|
||||
class _FakeQdrant:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def query_points(self, **kwargs):
|
||||
self.calls.append(kwargs)
|
||||
from qdrant_client.http import models
|
||||
|
||||
return models.QueryResponse(points=[])
|
||||
|
||||
class _MinimalRetriever(dfr_mod.DirectFlowRetriever):
|
||||
"""绕过 __init__(不建连接),只测 qdrant_search 参数透传。"""
|
||||
|
||||
def __init__(self):
|
||||
self.qdrant = _FakeQdrant()
|
||||
self.collection = "plant_srt_pc2025_v3_20260817_normalized"
|
||||
self.embed_url = ""
|
||||
self.embed_model = ""
|
||||
|
||||
def embed(self, text):
|
||||
return [0.0] * 1024
|
||||
|
||||
r = _MinimalRetriever()
|
||||
r.qdrant_search("景甜是谁", 10, score_threshold=0.45)
|
||||
assert r.qdrant.calls[-1]["score_threshold"] == 0.45
|
||||
assert r.qdrant.calls[-1]["collection_name"] == r.collection
|
||||
|
||||
r.qdrant_search("arginine", 10)
|
||||
assert r.qdrant.calls[-1]["score_threshold"] is None
|
||||
|
||||
|
||||
def test_item_from_payload_derives_citation_from_paper_id():
|
||||
"""回改 v4 根因单测:Qdrant payload 无 pmid/pmcid/doi 时从 paper_id
|
||||
派生 DOI——chat 模式 public_citation 过滤不得把有正文的 qdrant 条目
|
||||
全滤掉(此前 rerank 输入全是无 evidence 的 neo4j 图断言,bge 全
|
||||
0.02-0.047,final 门 0.25 全拒 sources=[])。"""
|
||||
import rag.qa_v1_1_73.direct_flow_retriever as dfr_mod
|
||||
|
||||
r = dfr_mod.DirectFlowRetriever.__new__(dfr_mod.DirectFlowRetriever)
|
||||
r.paper_metadata = {}
|
||||
# 内部 DOI 形态
|
||||
it = r.item_from_payload(
|
||||
{"paper_id": "10.3390_f12050650", "source": "s", "relation": "r",
|
||||
"target": "t", "evidence": "e"},
|
||||
"qdrant_vector", "q",
|
||||
)
|
||||
assert it.doi == "10.3390/f12050650"
|
||||
assert it.retrieval_methods == ["qdrant_vector"]
|
||||
# doi: 前缀形态
|
||||
it2 = r.item_from_payload(
|
||||
{"paper_id": "doi:10.1002/ps.6526", "source": "s", "relation": "r",
|
||||
"target": "t", "evidence": "e"},
|
||||
"qdrant_vector", "q",
|
||||
)
|
||||
assert it2.doi == "10.1002/ps.6526"
|
||||
# 派生出的 DOI 能通过 public_citation 校验(过滤不再误杀)
|
||||
from rag.qa_v1_1_73.citation_contract import public_citation
|
||||
cite = public_citation(evidence_id="", pmid="", pmcid="", doi=it.doi)
|
||||
assert cite == "[DOI: 10.3390/f12050650]"
|
||||
|
||||
|
||||
def test_env_overrides(monkeypatch):
|
||||
monkeypatch.setenv("QA_MIN_VECTOR_SCORE", "0.5")
|
||||
assert score_gate._env_value("QA_MIN_VECTOR_SCORE", 0.45) == 0.5
|
||||
monkeypatch.setenv("QA_MIN_VECTOR_SCORE", "abc")
|
||||
assert score_gate._env_value("QA_MIN_VECTOR_SCORE", 0.45) == 0.45
|
||||
monkeypatch.setenv("QA_MIN_SOURCE_SCORE", "1.5")
|
||||
assert score_gate._env_value("QA_MIN_SOURCE_SCORE", 0.25) == 0.25
|
||||
monkeypatch.delenv("QA_MIN_VECTOR_SCORE", raising=False)
|
||||
monkeypatch.delenv("QA_MIN_SOURCE_SCORE", raising=False)
|
||||
@@ -51,9 +51,13 @@ class ToolRunner:
|
||||
user: str,
|
||||
tools: List[Dict[str, Any]],
|
||||
manifests: Dict[str, CapabilityManifest],
|
||||
history: Optional[List[Dict[str, str]]] = None,
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
"""执行工具调用循环。返回 (content, audit)。
|
||||
|
||||
history:多轮记忆(S2),List[{role: user|assistant, content}],旧→新,
|
||||
不含当前这轮提问;注入为 messages 前置轮次(作为上下文,不参与工具结果回填)。
|
||||
|
||||
audit: {"rounds": n, "tool_calls": k, "used_json_fallback": bool,
|
||||
"tool_names": [...], "errors": [...]}
|
||||
"""
|
||||
@@ -62,9 +66,19 @@ class ToolRunner:
|
||||
"tool_calls": 0,
|
||||
"used_json_fallback": False,
|
||||
"tool_names": [],
|
||||
"tool_results": {},
|
||||
"errors": [],
|
||||
}
|
||||
messages: List[Dict[str, Any]] = [{"role": "user", "content": user}]
|
||||
messages: List[Dict[str, Any]] = []
|
||||
for h in history or []:
|
||||
if not isinstance(h, dict):
|
||||
continue
|
||||
role = str(h.get("role") or "").strip()
|
||||
content = str(h.get("content") or "").strip()
|
||||
if role not in {"user", "assistant"} or not content:
|
||||
continue
|
||||
messages.append({"role": role, "content": content})
|
||||
messages.append({"role": "user", "content": user})
|
||||
system_text = system
|
||||
for round_index in range(1, self.max_rounds + 1):
|
||||
audit["rounds"] = round_index
|
||||
@@ -82,6 +96,7 @@ class ToolRunner:
|
||||
break
|
||||
assistant_message: Dict[str, Any] = {"role": "assistant", "content": content or ""}
|
||||
tool_call_blocks: List[Dict[str, Any]] = []
|
||||
messages.append(assistant_message)
|
||||
for call in tool_calls:
|
||||
name = str(call.get("name") or call.get("function", {}).get("name") or "")
|
||||
raw_args = call.get("arguments") or call.get("function", {}).get("arguments") or {}
|
||||
@@ -112,9 +127,9 @@ class ToolRunner:
|
||||
"function": {"name": name or "unknown", "arguments": json.dumps(arguments, ensure_ascii=False)},
|
||||
}
|
||||
)
|
||||
audit["tool_results"][name] = result_text
|
||||
messages.append({"role": "tool", "tool_call_id": call_id, "content": result_text})
|
||||
assistant_message["tool_calls"] = tool_call_blocks
|
||||
messages.append(assistant_message)
|
||||
continue
|
||||
# 最后一轮:拿最终文本(无工具或轮次用尽后由调用方用 messages 再 complete 一次)
|
||||
return self._final_text(system_text, messages, tools), audit
|
||||
|
||||
@@ -0,0 +1,470 @@
|
||||
#!/usr/bin/env python3
|
||||
"""run_extension.py —— 扩展分析容器通用入口(宿主挂载 /opt/bioflow-entry,镜像零改动)。
|
||||
|
||||
执行形态(契约 §3):
|
||||
python -B /opt/bioflow-entry/run_extension.py --ext-dir /job/ext \
|
||||
--input-dir /job/input --output-dir /job/output [--param name=value ...]
|
||||
|
||||
职责:
|
||||
1. 读 <ext-dir>/extension.json(command_template 唯一事实源,契约 §2 schema)。
|
||||
2. render_command_template:移植 175 executor 模板渲染纯函数(白名单占位符
|
||||
12 个 + {param_*};{output_dir} 必含)。
|
||||
3. subprocess 执行(无 shell:shlex 切分;stdout/stderr 透传;超时进程组
|
||||
SIGTERM → 3s → SIGKILL,与 176 服务侧 terminate→kill 约定同源)。
|
||||
4. 输出契约校验:输出目录非空,否则退出码 3。
|
||||
|
||||
退出码约定(容器内零第三方依赖,仅 python 标准库):
|
||||
0 成功
|
||||
1.. 命令自身退出码透传(1/2/7/137…)
|
||||
2 extension.json 缺失/非法、模板渲染失败、白名单外占位符、缺 {output_dir}
|
||||
3 输出契约校验失败(输出目录为空)
|
||||
124 执行超时(进程组已终止)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shlex
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
EXIT_OK = 0
|
||||
EXIT_SETUP = 2
|
||||
EXIT_OUTPUT_CONTRACT = 3
|
||||
EXIT_TIMEOUT = 124
|
||||
|
||||
DEFAULT_TIMEOUT_SEC = 36000.0 # 与 176 单阶段 10h 超时对齐
|
||||
|
||||
# 占位符白名单(spec_v2 PLACEHOLDERS 移植;扩展 {param_<name>})
|
||||
PLACEHOLDERS = {
|
||||
"code_dir",
|
||||
"input_dir",
|
||||
"input_file",
|
||||
"input_hmm",
|
||||
"input_fasta",
|
||||
"input_query_fasta",
|
||||
"input_reference_fasta",
|
||||
"input_pep_fasta",
|
||||
"input_cds_fasta",
|
||||
"input_id_file",
|
||||
"output_dir",
|
||||
"work_dir",
|
||||
}
|
||||
|
||||
_ID_SUFFIXES = {".id", ".ids", ".geneid", ".geneids", ".txt", ".tsv", ".csv"}
|
||||
_FASTA_SUFFIXES = {".fa", ".fasta", ".faa", ".fsa", ".pep"}
|
||||
_GFF_SUFFIXES = {".gff", ".gff3"}
|
||||
|
||||
_ACTIVE_PROC: subprocess.Popen | None = None
|
||||
|
||||
|
||||
class ExtensionRenderError(RuntimeError):
|
||||
"""模板渲染失败(占位符缺失/参数缺值)。"""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 模板渲染纯函数(executor.py 移植裁剪;不 import gfa_core,零第三方依赖)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _safe_parameter_name(value: Any) -> str:
|
||||
text = re.sub(r"[^A-Za-z0-9_]+", "_", str(value or "").strip())
|
||||
text = re.sub(r"_+", "_", text).strip("_")
|
||||
if not text:
|
||||
return ""
|
||||
if text[0].isdigit():
|
||||
text = "p_" + text
|
||||
if text.startswith("param_"):
|
||||
text = text[6:]
|
||||
return text
|
||||
|
||||
|
||||
def _shell_parameter_value(value: Any) -> str:
|
||||
if isinstance(value, bool):
|
||||
return "true" if value else "false"
|
||||
if value is None:
|
||||
return ""
|
||||
text = str(value)
|
||||
if re.fullmatch(r"[A-Za-z0-9_./:+,@%-]+", text):
|
||||
return text
|
||||
return shlex.quote(text)
|
||||
|
||||
|
||||
def _first_input_by_suffix(paths: list[Path], suffixes: set[str]) -> str:
|
||||
for path in paths:
|
||||
name = path.name.lower()
|
||||
if path.suffix.lower() in suffixes or any(name.endswith(s) for s in suffixes):
|
||||
return path.as_posix()
|
||||
return ""
|
||||
|
||||
|
||||
def _input_files_by_suffix(paths: list[Path], suffixes: set[str]) -> list[Path]:
|
||||
matched = []
|
||||
for path in paths:
|
||||
name = path.name.lower()
|
||||
if path.suffix.lower() in suffixes or any(name.endswith(s) for s in suffixes):
|
||||
matched.append(path)
|
||||
return matched
|
||||
|
||||
|
||||
def _first_input_by_name(paths: list[Path], tokens: set[str]) -> str:
|
||||
for path in paths:
|
||||
name = path.name.lower()
|
||||
if any(token in name for token in tokens):
|
||||
return path.as_posix()
|
||||
return ""
|
||||
|
||||
|
||||
def _largest_file(paths: list[Path]) -> str:
|
||||
if not paths:
|
||||
return ""
|
||||
return max(paths, key=lambda p: p.stat().st_size if p.exists() else -1).as_posix()
|
||||
|
||||
|
||||
def _smallest_file(paths: list[Path]) -> str:
|
||||
if not paths:
|
||||
return ""
|
||||
return min(paths, key=lambda p: p.stat().st_size if p.exists() else 10**18).as_posix()
|
||||
|
||||
|
||||
def _role_of(extension: dict[str, Any], placeholder: str) -> str:
|
||||
"""input_roles 双形态兼容:契约 §2 平铺 {"input_file": "fastq"},175 旧版
|
||||
{"input_file": {"role": "generic_input_file"}}。仅用于语义占位符消歧。"""
|
||||
roles = extension.get("input_roles")
|
||||
if not isinstance(roles, dict):
|
||||
return ""
|
||||
spec = roles.get(placeholder)
|
||||
if isinstance(spec, dict):
|
||||
return str(spec.get("role") or "").lower()
|
||||
return str(spec or "").lower()
|
||||
|
||||
|
||||
def _used_placeholders(template: str) -> set[str]:
|
||||
return set(re.findall(r"\{([A-Za-z_][A-Za-z0-9_]*)\}", template or ""))
|
||||
|
||||
|
||||
def _validate_template(template: str) -> str | None:
|
||||
"""模板白名单校验(与 spec_v2 normalize 同纪律);None = 通过。"""
|
||||
if not template.strip():
|
||||
return "extension.json 缺少 command_template"
|
||||
if "{output_dir}" not in template:
|
||||
return "command_template 必须包含 {output_dir}(输出契约硬约束)"
|
||||
bad = sorted(
|
||||
name for name in _used_placeholders(template)
|
||||
if name not in PLACEHOLDERS and not name.startswith("param_")
|
||||
)
|
||||
if bad:
|
||||
return "command_template 含白名单外占位符: " + ", ".join(bad)
|
||||
return None
|
||||
|
||||
|
||||
def build_render_context(
|
||||
extension: dict[str, Any],
|
||||
*,
|
||||
run_dir: Path,
|
||||
input_dir: Path,
|
||||
input_files: list[Path],
|
||||
parameters: dict[str, Any] | None = None,
|
||||
output_dir: Path | None = None,
|
||||
code_dir: Path | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""渲染上下文:12 个白名单占位符(语义占位符按输入文件后缀/角色推导)+ param_*。"""
|
||||
output_dir = output_dir or (run_dir / "outputs")
|
||||
code_dir = code_dir or (run_dir / "code")
|
||||
input_file_role = _role_of(extension, "input_file")
|
||||
id_file = _first_input_by_suffix(input_files, _ID_SUFFIXES)
|
||||
annotation_file = (
|
||||
_first_input_by_suffix(input_files, _GFF_SUFFIXES)
|
||||
if ("annotation_gff" in input_file_role or "gff" in input_file_role)
|
||||
else ""
|
||||
)
|
||||
cis_element_file = (
|
||||
_first_input_by_name(input_files, {"plantcare", "cis_element", "cis-elements", "cis.elements"})
|
||||
if input_file_role == "generic_input_file"
|
||||
else ""
|
||||
)
|
||||
input_file = (
|
||||
annotation_file or cis_element_file or id_file
|
||||
or (input_files[0].as_posix() if input_files else input_dir.as_posix())
|
||||
)
|
||||
input_hmm = _first_input_by_suffix(input_files, {".hmm"}) or input_file
|
||||
fasta_paths = _input_files_by_suffix(input_files, _FASTA_SUFFIXES)
|
||||
input_fasta = fasta_paths[0].as_posix() if fasta_paths else input_file
|
||||
query_fasta = _largest_file(fasta_paths) or input_fasta
|
||||
reference_fasta = (
|
||||
_smallest_file([p for p in fasta_paths if p.as_posix() != query_fasta])
|
||||
if len(fasta_paths) > 1 else input_fasta
|
||||
)
|
||||
pep_fasta = _first_input_by_name(fasta_paths, {"pep", "protein", "faa"}) or query_fasta
|
||||
cds_fasta = _first_input_by_name(fasta_paths, {"cds"}) or input_fasta
|
||||
render_context: dict[str, Any] = {
|
||||
"code_dir": code_dir.as_posix(),
|
||||
"input_dir": input_dir.as_posix(),
|
||||
"input_file": input_file,
|
||||
"input_hmm": input_hmm,
|
||||
"input_fasta": input_fasta,
|
||||
"input_query_fasta": query_fasta,
|
||||
"input_reference_fasta": reference_fasta,
|
||||
"input_pep_fasta": pep_fasta,
|
||||
"input_cds_fasta": cds_fasta,
|
||||
"input_id_file": id_file or input_file,
|
||||
"output_dir": output_dir.as_posix(),
|
||||
"work_dir": run_dir.as_posix(),
|
||||
}
|
||||
for name, value in (parameters or {}).items():
|
||||
render_context["param_" + _safe_parameter_name(name)] = _shell_parameter_value(value)
|
||||
return render_context
|
||||
|
||||
|
||||
def render_command_template(
|
||||
extension: dict[str, Any],
|
||||
*,
|
||||
template: str,
|
||||
run_dir: Path,
|
||||
input_dir: Path,
|
||||
input_files: list[Path],
|
||||
parameters: dict[str, Any] | None = None,
|
||||
output_dir: Path | None = None,
|
||||
code_dir: Path | None = None,
|
||||
) -> str:
|
||||
"""模板渲染纯函数(executor.py 同签名移植)。占位符缺值 → ExtensionRenderError。"""
|
||||
context = build_render_context(
|
||||
extension,
|
||||
run_dir=run_dir,
|
||||
input_dir=input_dir,
|
||||
input_files=input_files,
|
||||
parameters=parameters,
|
||||
output_dir=output_dir,
|
||||
code_dir=code_dir,
|
||||
)
|
||||
try:
|
||||
return template.format(**context)
|
||||
except KeyError as exc:
|
||||
missing = str(exc).strip("'\"")
|
||||
raise ExtensionRenderError(f"Missing command template value: {missing}") from exc
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 执行与输出契约
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _output_contract_ok(output_dir: Path) -> bool:
|
||||
"""输出契约:输出目录存在且非空(至少一个条目)。"""
|
||||
return output_dir.is_dir() and any(output_dir.iterdir())
|
||||
|
||||
|
||||
def _find_code_dir(ext_dir: Path) -> Path | None:
|
||||
"""代码目录:converted_code 优先(契约 §1 布局),raw_code 兜底。"""
|
||||
for name in ("converted_code", "raw_code"):
|
||||
candidate = ext_dir / name
|
||||
if candidate.is_dir():
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def _copy_code_tree(src: Path, dst: Path) -> set[str]:
|
||||
"""把代码树复制进可写目录;返回原始代码相对路径集合(收割时排除)。"""
|
||||
import shutil
|
||||
rels: set[str] = set()
|
||||
if not src.is_dir():
|
||||
return rels
|
||||
for item in src.rglob("*"):
|
||||
rel = item.relative_to(src).as_posix()
|
||||
dest = dst / rel
|
||||
if item.is_dir():
|
||||
dest.mkdir(parents=True, exist_ok=True)
|
||||
else:
|
||||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(item, dest)
|
||||
rels.add(rel)
|
||||
return rels
|
||||
|
||||
|
||||
def _harvest_outputs(src: Path, dst: Path, exclude: set[str]) -> None:
|
||||
"""把 src 内产物(排除原始代码路径)合并进 dst,保留相对路径。"""
|
||||
import shutil
|
||||
for item in src.rglob("*"):
|
||||
rel = item.relative_to(src).as_posix()
|
||||
if rel in exclude:
|
||||
continue
|
||||
# 排除路径是原始代码文件所在目录时也跳过其祖先目录本身
|
||||
if any(rel == e or rel.startswith(e + "/") for e in exclude):
|
||||
continue
|
||||
dest = dst / rel
|
||||
if item.is_dir():
|
||||
dest.mkdir(parents=True, exist_ok=True)
|
||||
else:
|
||||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(item, dest)
|
||||
|
||||
|
||||
def _kill_group(pid: int, sig: int) -> None:
|
||||
try:
|
||||
os.killpg(pid, sig)
|
||||
except (ProcessLookupError, PermissionError):
|
||||
pass
|
||||
|
||||
|
||||
def _sigterm_handler(signum: int, frame: Any) -> None:
|
||||
"""docker stop 的 SIGTERM:转发到子进程组后以 124 退出(日志可辨)。"""
|
||||
proc = _ACTIVE_PROC
|
||||
if proc is not None and proc.poll() is None:
|
||||
_kill_group(proc.pid, signal.SIGTERM)
|
||||
sys.exit(EXIT_TIMEOUT)
|
||||
|
||||
|
||||
def _run_command(command: list[str], cwd: Path, env: dict[str, str], timeout_sec: float) -> int:
|
||||
"""执行命令(stdout/stderr 直接透传容器流;超时进程组 SIGTERM→3s→SIGKILL)。
|
||||
|
||||
可执行文件缺失(镜像缺工具)→ 可读退出码 2(不抛 traceback)。"""
|
||||
global _ACTIVE_PROC
|
||||
try:
|
||||
proc = subprocess.Popen(command, cwd=str(cwd), env=env, start_new_session=True)
|
||||
except FileNotFoundError:
|
||||
print(f"[run_extension] 可执行文件缺失(镜像 PATH 中不存在): {command[0]}",
|
||||
file=sys.stderr)
|
||||
return EXIT_SETUP
|
||||
_ACTIVE_PROC = proc
|
||||
try:
|
||||
return proc.wait(timeout=timeout_sec)
|
||||
except subprocess.TimeoutExpired:
|
||||
_kill_group(proc.pid, signal.SIGTERM)
|
||||
try:
|
||||
proc.wait(timeout=3)
|
||||
except subprocess.TimeoutExpired:
|
||||
_kill_group(proc.pid, signal.SIGKILL)
|
||||
proc.wait()
|
||||
return EXIT_TIMEOUT
|
||||
finally:
|
||||
_ACTIVE_PROC = None
|
||||
|
||||
|
||||
def _parse_params(items: list[str]) -> dict[str, Any]:
|
||||
parameters: dict[str, Any] = {}
|
||||
for item in items:
|
||||
name, sep, value = item.partition("=")
|
||||
if not sep or not name.strip():
|
||||
raise ExtensionRenderError(f"invalid --param (需要 name=value): {item}")
|
||||
parameters[name.strip()] = value
|
||||
return parameters
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="run_extension",
|
||||
description="扩展分析容器通用入口(契约 EXTENSION_CONTRACT §3)",
|
||||
)
|
||||
parser.add_argument("--ext-dir", default="/job/ext", help="扩展目录(挂载 :ro)")
|
||||
parser.add_argument("--input-dir", default="/job/input", help="输入目录")
|
||||
parser.add_argument("--output-dir", default="/job/output", help="输出目录(契约校验对象)")
|
||||
parser.add_argument("--work-dir", default="", help="工作目录(默认=输出目录)")
|
||||
parser.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT_SEC, help="执行超时(秒)")
|
||||
parser.add_argument("--param", action="append", default=[], help="参数 name=value(可重复)")
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
try:
|
||||
signal.signal(signal.SIGTERM, _sigterm_handler)
|
||||
except ValueError: # 非主线程(测试)不支持 handler
|
||||
pass
|
||||
|
||||
ext_dir = Path(args.ext_dir)
|
||||
ext_json_path = ext_dir / "extension.json"
|
||||
if not ext_json_path.is_file():
|
||||
print(f"[run_extension] extension.json 缺失: {ext_json_path}", file=sys.stderr)
|
||||
return EXIT_SETUP
|
||||
try:
|
||||
extension = json.loads(ext_json_path.read_text(encoding="utf-8"))
|
||||
except ValueError as exc:
|
||||
print(f"[run_extension] extension.json 解析失败: {exc}", file=sys.stderr)
|
||||
return EXIT_SETUP
|
||||
if not isinstance(extension, dict):
|
||||
print("[run_extension] extension.json 必须是 JSON 对象", file=sys.stderr)
|
||||
return EXIT_SETUP
|
||||
|
||||
template = str(extension.get("command_template") or "")
|
||||
error = _validate_template(template)
|
||||
if error:
|
||||
print(f"[run_extension] {error}", file=sys.stderr)
|
||||
return EXIT_SETUP
|
||||
|
||||
code_dir = _find_code_dir(ext_dir)
|
||||
if code_dir is None:
|
||||
print("[run_extension] 代码目录缺失(converted_code/raw_code 均不存在)", file=sys.stderr)
|
||||
return EXIT_SETUP
|
||||
|
||||
input_dir = Path(args.input_dir)
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
# 2026-09-15 修复:代码挂载只读,但 18 用例语义是「cd 代码目录后脚本写
|
||||
# outputs/ 相对目录」。等价做法 = 把代码复制进可写 run_dir 执行,
|
||||
# 结束后把 run_dir 产物收割进 output_dir。显式 --work-dir 可覆盖。
|
||||
work_dir = Path(args.work_dir) if args.work_dir else (output_dir / "_work")
|
||||
work_dir.mkdir(parents=True, exist_ok=True)
|
||||
code_rel_paths = _copy_code_tree(code_dir, work_dir)
|
||||
# 模板中 {code_dir} 指向可写副本(而非只读挂载)
|
||||
effective_code_dir = work_dir
|
||||
|
||||
try:
|
||||
parameters = _parse_params(args.param)
|
||||
except ExtensionRenderError as exc:
|
||||
print(f"[run_extension] {exc}", file=sys.stderr)
|
||||
return EXIT_SETUP
|
||||
|
||||
input_files = sorted(p for p in input_dir.rglob("*") if p.is_file())
|
||||
try:
|
||||
command_text = render_command_template(
|
||||
extension,
|
||||
template=template,
|
||||
run_dir=work_dir,
|
||||
input_dir=input_dir,
|
||||
input_files=input_files,
|
||||
parameters=parameters,
|
||||
output_dir=output_dir,
|
||||
code_dir=effective_code_dir,
|
||||
)
|
||||
except ExtensionRenderError as exc:
|
||||
print(f"[run_extension] 模板渲染失败: {exc}", file=sys.stderr)
|
||||
return EXIT_SETUP
|
||||
|
||||
try:
|
||||
command = shlex.split(command_text)
|
||||
except ValueError as exc:
|
||||
print(f"[run_extension] 命令切分失败: {exc}", file=sys.stderr)
|
||||
return EXIT_SETUP
|
||||
if not command:
|
||||
print("[run_extension] 渲染后命令为空", file=sys.stderr)
|
||||
return EXIT_SETUP
|
||||
|
||||
print("[run_extension] command: " + command_text)
|
||||
sys.stdout.flush()
|
||||
|
||||
env = os.environ.copy()
|
||||
env.setdefault("PYTHONUNBUFFERED", "1")
|
||||
code = _run_command(command, work_dir, env, args.timeout)
|
||||
if code == EXIT_TIMEOUT:
|
||||
print(f"[run_extension] 执行超时({args.timeout}s),进程组已终止 (exit {EXIT_TIMEOUT})", file=sys.stderr)
|
||||
return EXIT_TIMEOUT
|
||||
if code != EXIT_OK:
|
||||
print(f"[run_extension] command exited with code {code}", file=sys.stderr)
|
||||
return code
|
||||
# 收割:work_dir 产物(排除原始代码文件)合并进 output_dir,然后删除 work 目录
|
||||
try:
|
||||
_harvest_outputs(work_dir, output_dir, code_rel_paths)
|
||||
import shutil
|
||||
shutil.rmtree(work_dir, ignore_errors=True)
|
||||
except OSError as exc:
|
||||
print(f"[run_extension] 产物收割失败: {exc}", file=sys.stderr)
|
||||
return EXIT_SETUP
|
||||
if not _output_contract_ok(output_dir):
|
||||
print(f"[run_extension] 输出契约校验失败:输出目录为空 {output_dir} (exit {EXIT_OUTPUT_CONTRACT})",
|
||||
file=sys.stderr)
|
||||
return EXIT_OUTPUT_CONTRACT
|
||||
return EXIT_OK
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,218 @@
|
||||
"""run_extension 容器入口单测 —— 纯 stdlib unittest(容器内无 pytest)。
|
||||
|
||||
覆盖:模板渲染白名单/输出契约(退出码 3)/参数注入/退出码透传/超时/extension.json
|
||||
异常。真 subprocess 执行(容器入口语义,宿主机 python3 直跑,不需要 docker)。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
ENTRY = REPO_ROOT / "bioflow_kernel" / "cli" / "run_extension.py"
|
||||
|
||||
if not ENTRY.is_file():
|
||||
raise SystemExit(f"入口脚本缺失: {ENTRY}")
|
||||
|
||||
|
||||
def _write_extension(ext_dir: Path, template: str, **extra) -> None:
|
||||
payload = {
|
||||
"schema_version": "extension_stage_v1",
|
||||
"ext_id": "ext_test",
|
||||
"display_name": "测试扩展",
|
||||
"status": "published",
|
||||
"command_template": template,
|
||||
"original_command": "",
|
||||
"runtime_image": "bioflow-gf",
|
||||
"requirements": {"pip": [], "apt": []},
|
||||
"parameters": [],
|
||||
"input_roles": {"input_file": "fastq"},
|
||||
"nl_triggers": [],
|
||||
"analysis_type": "table_filter",
|
||||
"golden": None,
|
||||
"repair_count": 0,
|
||||
}
|
||||
payload.update(extra)
|
||||
(ext_dir / "extension.json").write_text(json.dumps(payload), encoding="utf-8")
|
||||
|
||||
|
||||
def _run_entry(ext_dir: Path, input_dir: Path, output_dir: Path,
|
||||
params: list[str] | None = None, timeout: float = 120.0,
|
||||
extra_args: list[str] | None = None) -> subprocess.CompletedProcess:
|
||||
cmd = [sys.executable, "-B", str(ENTRY),
|
||||
"--ext-dir", str(ext_dir),
|
||||
"--input-dir", str(input_dir),
|
||||
"--output-dir", str(output_dir)]
|
||||
cmd += extra_args or []
|
||||
for item in params or []:
|
||||
cmd += ["--param", item]
|
||||
# 测试宿主无裸 python(生产镜像内由 /opt/conda/bin 提供):注入 shim 验证执行语义
|
||||
import tempfile
|
||||
|
||||
shim_dir = tempfile.mkdtemp(prefix="ext-shim-")
|
||||
shim = Path(shim_dir) / "python"
|
||||
shim.write_text(f"#!/bin/sh\nexec {sys.executable} \"$@\"\n", encoding="utf-8")
|
||||
shim.chmod(0o755)
|
||||
env = dict(__import__("os").environ)
|
||||
env["PATH"] = shim_dir + ":" + env.get("PATH", "")
|
||||
return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=env)
|
||||
|
||||
|
||||
class RenderTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
import importlib.util
|
||||
|
||||
spec = importlib.util.spec_from_file_location("run_extension", ENTRY)
|
||||
self.mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(self.mod)
|
||||
|
||||
def _ctx(self, template: str, **extra):
|
||||
ext = {"command_template": template, "input_roles": {"input_file": "fastq"}}
|
||||
ext.update(extra)
|
||||
return self.mod.build_render_context(
|
||||
ext,
|
||||
run_dir=Path("/job/output"),
|
||||
input_dir=Path("/job/input"),
|
||||
input_files=[Path("/job/input/sample.fastq")],
|
||||
output_dir=Path("/job/output"),
|
||||
code_dir=Path("/job/ext/converted_code"),
|
||||
)
|
||||
|
||||
def test_basic_context(self):
|
||||
ctx = self._ctx("python {code_dir}/x.py --fastq {input_file} --out {output_dir}")
|
||||
self.assertEqual(ctx["code_dir"], "/job/ext/converted_code")
|
||||
self.assertEqual(ctx["input_file"], "/job/input/sample.fastq")
|
||||
self.assertEqual(ctx["output_dir"], "/job/output")
|
||||
self.assertEqual(ctx["work_dir"], "/job/output")
|
||||
|
||||
def test_input_hmm_suffix(self):
|
||||
ctx = self._ctx("hmmsearch {input_hmm} > {output_dir}/out.txt")
|
||||
# 无 .hmm 输入时回退 input_file(与 executor.py 同语义)
|
||||
self.assertEqual(ctx["input_hmm"], "/job/input/sample.fastq")
|
||||
|
||||
def test_template_render_unknown_placeholder_rejected(self):
|
||||
mod = self.mod
|
||||
with self.assertRaises(mod.ExtensionRenderError):
|
||||
mod.render_command_template(
|
||||
{"command_template": "run {not_a_placeholder}"},
|
||||
template="run {not_a_placeholder}",
|
||||
run_dir=Path("/job/output"),
|
||||
input_dir=Path("/job/input"),
|
||||
input_files=[Path("/job/input/a.fastq")],
|
||||
output_dir=Path("/job/output"),
|
||||
code_dir=Path("/job/ext/converted_code"),
|
||||
)
|
||||
|
||||
def test_validate_template_output_dir_required(self):
|
||||
self.assertIn("output_dir", self.mod._validate_template("python {code_dir}/x.py") or "")
|
||||
self.assertIsNone(self.mod._validate_template("python {code_dir}/x.py > {output_dir}/r.txt"))
|
||||
|
||||
def test_validate_template_whitelist(self):
|
||||
err = self.mod._validate_template("run {unknown_placeholder} > {output_dir}/x")
|
||||
self.assertIn("unknown_placeholder", err or "")
|
||||
self.assertIsNone(self.mod._validate_template("x {param_alpha} {param_beta} {work_dir} {output_dir}"))
|
||||
|
||||
def test_shell_parameter_value_quoting(self):
|
||||
self.assertEqual(self.mod._shell_parameter_value("a-b_c/d.e"), "a-b_c/d.e")
|
||||
self.assertEqual(self.mod._shell_parameter_value("a b"), "'a b'")
|
||||
|
||||
def test_safe_parameter_name(self):
|
||||
self.assertEqual(self.mod._safe_parameter_name("my-param"), "my_param")
|
||||
self.assertEqual(self.mod._safe_parameter_name("2fast"), "p_2fast")
|
||||
|
||||
|
||||
class EntryTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
import tempfile
|
||||
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
self.root = Path(self._tmp.name)
|
||||
self.ext = self.root / "ext"
|
||||
self.code = self.ext / "converted_code"
|
||||
self.input_dir = self.root / "input"
|
||||
self.output = self.root / "output"
|
||||
for d in (self.code, self.input_dir):
|
||||
d.mkdir(parents=True)
|
||||
self.script = self.code / "hello.py"
|
||||
self.script.write_text(
|
||||
textwrap.dedent("""\
|
||||
import pathlib, sys
|
||||
print("hello from ext", flush=True)
|
||||
pathlib.Path(sys.argv[1]).mkdir(parents=True, exist_ok=True)
|
||||
pathlib.Path(sys.argv[1], "result.tsv").write_text("a\\tb\\n")
|
||||
"""), encoding="utf-8")
|
||||
(self.input_dir / "sample.fastq").write_text(">s1\\nACGT\\n", encoding="utf-8")
|
||||
|
||||
def tearDown(self):
|
||||
self._tmp.cleanup()
|
||||
|
||||
def _template(self) -> str:
|
||||
return ("python {code_dir}/hello.py {output_dir} --fastq {input_file} "
|
||||
"--min {param_min} --label {param_label}")
|
||||
|
||||
def test_success_and_output_contract(self):
|
||||
_write_extension(self.ext, self._template())
|
||||
result = _run_entry(self.ext, self.input_dir, self.output,
|
||||
params=["min=30", "label=run_1"])
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
self.assertIn("hello from ext", result.stdout)
|
||||
self.assertTrue((self.output / "result.tsv").is_file())
|
||||
|
||||
def test_output_contract_empty_exit_3(self):
|
||||
(self.script).write_text("print('no output written')\n", encoding="utf-8")
|
||||
_write_extension(self.ext, "python {code_dir}/hello.py --out {output_dir}")
|
||||
result = _run_entry(self.ext, self.input_dir, self.output)
|
||||
self.assertEqual(result.returncode, 3)
|
||||
self.assertIn("输出契约", result.stderr)
|
||||
|
||||
def test_missing_extension_json_exit_2(self):
|
||||
result = _run_entry(self.ext, self.input_dir, self.output)
|
||||
self.assertEqual(result.returncode, 2)
|
||||
|
||||
def test_invalid_template_exit_2(self):
|
||||
_write_extension(self.ext, "python {code_dir}/hello.py --unknown {nope} --out {output_dir}")
|
||||
result = _run_entry(self.ext, self.input_dir, self.output)
|
||||
self.assertEqual(result.returncode, 2)
|
||||
self.assertIn("白名单外", result.stderr)
|
||||
|
||||
def test_exit_code_passthrough(self):
|
||||
(self.code / "fail.sh").write_text("#!/bin/sh\nexit 7\n", encoding="utf-8")
|
||||
(self.code / "fail.sh").chmod(0o755)
|
||||
_write_extension(self.ext, "sh {code_dir}/fail.sh > {output_dir}/o.txt")
|
||||
result = _run_entry(self.ext, self.input_dir, self.output)
|
||||
self.assertEqual(result.returncode, 7)
|
||||
self.assertIn("code 7", result.stderr)
|
||||
|
||||
def test_template_missing_param_exit_2(self):
|
||||
_write_extension(self.ext, "python {code_dir}/hello.py --x {param_missing} --out {output_dir}")
|
||||
result = _run_entry(self.ext, self.input_dir, self.output)
|
||||
self.assertEqual(result.returncode, 2)
|
||||
self.assertIn("模板渲染失败", result.stderr)
|
||||
|
||||
def test_param_with_spaces_single_arg(self):
|
||||
# 参数含空格 → shlex.quote 保证独立 argv(等价 kernel list-form 无 shell 注入)
|
||||
(self.code / "args.py").write_text(
|
||||
"import sys; print(repr(sys.argv[1]))\n", encoding="utf-8")
|
||||
_write_extension(self.ext, "python {code_dir}/args.py {param_label} --out {output_dir}")
|
||||
result = _run_entry(self.ext, self.input_dir, self.output,
|
||||
params=["label=hello world"])
|
||||
self.assertEqual(result.returncode, 3, result.stderr) # 未写输出 → 契约校验
|
||||
self.assertIn("'hello world'", result.stdout)
|
||||
|
||||
def test_timeout_returns_124(self):
|
||||
(self.code / "sleep.py").write_text(
|
||||
"import time\ntime.sleep(600)\n", encoding="utf-8")
|
||||
_write_extension(self.ext, "python {code_dir}/sleep.py --out {output_dir}")
|
||||
result = _run_entry(self.ext, self.input_dir, self.output,
|
||||
timeout=10.0, extra_args=["--timeout", "2"])
|
||||
self.assertEqual(result.returncode, 124)
|
||||
self.assertIn("超时", result.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
@@ -0,0 +1,85 @@
|
||||
---
|
||||
name: bioflow-analysis
|
||||
description: >-
|
||||
PlantAI 生信分析编排 Skill。当用户提供原始组学数据(蛋白 FASTA/基因组/FASTQ/表达矩阵等)
|
||||
并提出基因家族、转录组、变异或自定义分析需求时使用。负责从原始数据到交付报告的
|
||||
全流程编排:选链路、登记资产、就绪检查、参数确认、执行、诊断、版本化修改、交付清单。
|
||||
触发词见 nl_triggers;仅在明确涉及生信分析时触发。
|
||||
nl_triggers:
|
||||
- 基因家族
|
||||
- 家族成员
|
||||
- 亚家族
|
||||
- 系统发育
|
||||
- 建树
|
||||
- 进化树
|
||||
- motif
|
||||
- 结构域
|
||||
- 共线性
|
||||
- Ka/Ks
|
||||
- 转录组
|
||||
- RNA-seq
|
||||
- RNAseq
|
||||
- rnaseq
|
||||
- 差异表达
|
||||
- 表达定量
|
||||
- 富集
|
||||
- WGCNA
|
||||
- GSEA
|
||||
- fastq
|
||||
- 实验设计
|
||||
- 生信分析
|
||||
- 变异检测
|
||||
- GWAS
|
||||
---
|
||||
|
||||
# bioflow-analysis 能力声明
|
||||
|
||||
PlantAI 生信分析 Skill,面向「原始数据 → 交付报告」的分析编排。所有分析经
|
||||
`bioflowctl`(唯一事实源)与 11 个分析控制 MCP 工具完成,服务端状态经 SQLite 持久化。
|
||||
|
||||
## 11 工具白名单
|
||||
|
||||
Skill 只允许调用以下 11 个分析控制工具(其余为内部实现,不外泄):
|
||||
|
||||
| 工具 | 职责 | 对应 CLI |
|
||||
|------|------|----------|
|
||||
| get_task_context | 任务绑定(session/job/plan_revision/run_id/当前阶段) | — |
|
||||
| inspect_asset | 资产/文件探测(格式/预览/列结构) | file inspect |
|
||||
| validate_plan | 读单一参数 schema 校验参数 | (补齐) |
|
||||
| check_readiness | 四象限就绪检查(环境/参考/磁盘/预算) | check readiness |
|
||||
| get_run_evidence | 运行证据(日志切片/失败原因/manifest 校验值) | — |
|
||||
| propose_change | 参数变更影响范围 | change propose |
|
||||
| apply_change | 幂等应用变更 | change apply |
|
||||
| execute_nodes | 只重跑受影响节点 | run node-execute |
|
||||
| control_execution | 执行所有权(stop/resume) | run stop/resume |
|
||||
| validate_delivery | 交付清单完整性校验 | delivery validate |
|
||||
| publish_delivery | 发布交付清单 | delivery publish |
|
||||
|
||||
## 对话剧本
|
||||
|
||||
1. **原始数据**:识别用户提供的原始数据(蛋白 FASTA / 基因组 / FASTQ / 表达矩阵 / 设计表)。
|
||||
2. **选链路**:`inspect_asset` 探测数据格式;据此选择分析链路(基因家族 / 转录组 / …)。
|
||||
3. **登记**:登记数据资产(角色化 `asset register`),保证「下游直入」可用。
|
||||
4. **就绪**:`check_readiness` 四象限检查;缺项返回结构化错误阻断,不启动。
|
||||
5. **参数确认**:`validate_plan` 校验参数;未知参数/越界拒绝。
|
||||
6. **执行**:`execute_nodes` 只重跑受影响节点;`control_execution` 管理执行所有权。
|
||||
7. **诊断**:失败时 `get_run_evidence` 取日志切片/manifest 校验值,按 8 类错误分类
|
||||
(input/resource/environment/transient/adapter/artifact/reference/scientific_precondition)。
|
||||
8. **版本化**:`propose_change` + `apply_change` 版本化修改,历史不覆盖、幂等。
|
||||
9. **交付**:`validate_delivery` + `publish_delivery` 校验并发布交付清单。
|
||||
|
||||
## 人工边界(requires_human)
|
||||
|
||||
- `subfamily_confirm`:亚家族分类确认,须人工确认,不可代用户。
|
||||
- `rnaseq_de`:差异表达实验设计确认,须人工确认,不可代用户。
|
||||
上述阶段在 pending_actions 中 `requires_human=true`,Skill 只呈现表单默认值,不代确认。
|
||||
|
||||
## 验收清单
|
||||
|
||||
- [ ] 原始数据格式已探测(inspect_asset)
|
||||
- [ ] 就绪检查通过(check_readiness ready=true)
|
||||
- [ ] 参数经单一 schema 校验(validate_plan)
|
||||
- [ ] 执行所有权唯一(control_execution 不重复接管/不重叠)
|
||||
- [ ] 失败已诊断(get_run_evidence + 8 类错误分类)
|
||||
- [ ] 变更版本化(propose_change/apply_change 幂等、历史不覆盖)
|
||||
- [ ] 交付清单完整并已发布(validate_delivery ok → publish_delivery approved)
|
||||
@@ -35,6 +35,10 @@ class Settings(BaseSettings):
|
||||
|
||||
# ---- 阶段→镜像映射(A1 交付物,仓库内 bioflow_kernel/stage_image_map.json)----
|
||||
image_map_path: str = "bioflow_kernel/stage_image_map.json"
|
||||
# 内核树:运行时只读挂载进执行容器(-v kernel:/bioflow-kernel:ro + -w)
|
||||
# 除 bioflow-gf 外其余镜像未内置内核,挂载后全部阶段镜像统一可用
|
||||
kernel_dir: str = "bioflow_kernel"
|
||||
container_kernel_dir: str = "/bioflow-kernel"
|
||||
|
||||
# ---- 数据与任务目录 ----
|
||||
data_root: str = "bioflow_data" # 生信数据根(hmm/rag_db 等,只读挂载进容器)
|
||||
@@ -54,7 +58,28 @@ class Settings(BaseSettings):
|
||||
default_memory: str = "32g"
|
||||
log_tail_lines: int = 80 # manifest 接口日志尾行数(≤80)
|
||||
max_upload_gb: float = 100.0 # 上传大小上限(GB;超限 413)
|
||||
upload_ext_allow: str = ".fasta,.fa,.faa,.fna,.pep,.cds,.gz,.gff,.gff3,.gtf,.txt,.tsv,.csv,.xlsx,.xls,.bed,.vcf,.vcf.gz,.bam,.cram,.sam,.hmm,.nwk,.newick,.treefile,.json,.yaml,.yml,.pdf"
|
||||
upload_ext_allow: str = ".fasta,.fa,.faa,.fna,.pep,.cds,.gz,.gff,.gff3,.gtf,.txt,.tsv,.csv,.xlsx,.xls,.bed,.vcf,.vcf.gz,.bam,.cram,.sam,.hmm,.nwk,.newick,.treefile,.json,.yaml,.yml,.pdf,.zip,.py,.r,.sh"
|
||||
|
||||
# ---- 服务器路径登记(FR-BIO-007;用户数据可在任意位置)----
|
||||
max_register_files_per_job: int = 200 # 目录展开单任务文件数上限(防误填大目录)
|
||||
# 敏感路径黑名单:命中前缀即 403(白名单根校验已按用户要求移除)
|
||||
sensitive_path_prefixes: tuple[str, ...] = (
|
||||
"/etc", "/root", "/proc", "/sys", "/dev", "/run", "/boot",
|
||||
"/bin", "/sbin", "/lib", "/lib64", "/usr/bin", "/usr/sbin",
|
||||
"/usr/lib", "/usr/lib64", "/usr/local/bin", "/var/lib", "/var/log",
|
||||
"/opt/neo4j", "/var/lib/docker", "/srv/plantai/install",
|
||||
)
|
||||
# 敏感文件名(正则,re.match 整名);.env 等含密钥配置、密钥文件不可登记
|
||||
sensitive_filename_patterns: tuple[str, ...] = (
|
||||
r"\.env.*", r".*\.key$", r".*\.pem$", r"id_rsa.*", r"id_ed25519.*",
|
||||
r".*\.p12$", r"\.gitconfig.*",
|
||||
)
|
||||
|
||||
# ---- 用户自定义扩展分析(并行链路;存储根 + AI 转换 LLM 键 + 修复上限)----
|
||||
extensions_dir: str = "bioflow_data/extensions" # 扩展存储根({ext_id}/extension.json...)
|
||||
extension_max_repair_rounds: int = 20 # AI 修复循环上限(契约 §6 可配)
|
||||
# 扩展转换 LLM 键:DOCGEN_LLM_* 直接复用(docgen 直连模式),
|
||||
# BIOFLOW_EXTENSION_* 为可选别名(extension_llm.llm_config 兜底读取)。
|
||||
|
||||
# ---- 规划(planner)----
|
||||
llm_plan_timeout: float = 60.0 # 规划 LLM 单次调用超时(秒)
|
||||
@@ -104,6 +129,10 @@ class Settings(BaseSettings):
|
||||
def data_root_abs(self) -> str:
|
||||
return _resolve_root(self.data_root)
|
||||
|
||||
@property
|
||||
def kernel_dir_abs(self) -> str:
|
||||
return _resolve_root(self.kernel_dir)
|
||||
|
||||
@property
|
||||
def jobs_dir_abs(self) -> str:
|
||||
return _resolve_root(self.jobs_dir)
|
||||
@@ -120,6 +149,10 @@ class Settings(BaseSettings):
|
||||
def flow_docs_dir_abs(self) -> str:
|
||||
return _resolve_root(self.flow_docs_dir)
|
||||
|
||||
@property
|
||||
def extensions_dir_abs(self) -> str:
|
||||
return _resolve_root(self.extensions_dir)
|
||||
|
||||
@property
|
||||
def upload_ext_set(self) -> set[str]:
|
||||
return {ext.strip().lower() for ext in self.upload_ext_allow.split(",") if ext.strip()}
|
||||
|
||||
@@ -161,6 +161,22 @@ async def init_context(db_path: str | None = None) -> Context:
|
||||
database = await init_db()
|
||||
flow_store = init_flow_store(get_settings().flow_docs_dir_abs)
|
||||
context = build_context(db=database, flow_store=flow_store)
|
||||
# TL 接线(2026-09-12):扩展执行器接 Dev-B extension_exec。
|
||||
# 失败(如 docker 不可用/配置缺失)静默降级——trial/run 返回 503 而非 500。
|
||||
try:
|
||||
from .services.bioflow import extensions as ext_mod
|
||||
# 2026-09-15 修复:裸 ExtensionExecConfig() 的 entry_script_dir 为空串,
|
||||
# build_command 用相对路径 run_extension.py 在 176 cwd 下找不到入口。
|
||||
# 改走 _default_exec_config()(经 _resolve_extension_entry_dir 解析部署树路径)。
|
||||
from .services.bioflow.extension_exec import _default_exec_config, ExtensionExecutor as RealExecutor
|
||||
cfg = _default_exec_config()
|
||||
real = RealExecutor(cfg)
|
||||
store = ext_mod.get_store()
|
||||
if store is not None:
|
||||
ext_mod.inject_real_executor(real, store)
|
||||
logger.info("extension executor wired (image suffix=%s)", cfg.docker_tag_suffix)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("extension executor wiring skipped: %s", exc)
|
||||
return context
|
||||
|
||||
|
||||
@@ -189,6 +205,17 @@ def _record_from_job(job: dict[str, Any], stages: list[dict[str, Any]]) -> Any:
|
||||
auto_continue=bool(job.get("auto_continue")),
|
||||
archive_id=job.get("archive_id"),
|
||||
)
|
||||
# 登记路径恢复(server_registered.json;供执行容器自动挂载)
|
||||
try:
|
||||
import json as _json
|
||||
reg_file = record.job_dir / "server_registered.json"
|
||||
if reg_file.is_file():
|
||||
record.registered_server_paths = (
|
||||
_json.loads(reg_file.read_text(encoding="utf-8"))
|
||||
.get("registered_server_paths") or []
|
||||
)
|
||||
except Exception: # noqa: BLE001 - 恢复失败不阻断(挂载缺失→阶段报错可查)
|
||||
pass
|
||||
for s in stages:
|
||||
record._stage_states[s["stage"]] = {
|
||||
"status": s["status"],
|
||||
|
||||
+39
-10
@@ -2,13 +2,21 @@
|
||||
|
||||
- 独立部署树 bioflow_service/(FR-BIO-028:独立进程 8098 + 独立 SQLite)。
|
||||
- lifespan:init_db(幂等迁移)→ init_context(kernel/catalog/planner/runner)
|
||||
→ recover(服务重启恢复 in-flight 续跑,FR-BIO-031)→ shutdown 清理。
|
||||
→ recover(服务重启恢复 in-flight 续跑,FR-BIO-031)→ MCP session manager
|
||||
(anyio task group,streamable-http 处理请求必需)→ shutdown 清理。
|
||||
- 端点:health(/api/bioflow/health,前端探活 FR-BIO-025)。
|
||||
- 不依赖 BFF/生产服务(8090/8091 等绝不探测/连接)。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# 把 bioflow_service/.env 注入 os.environ(llm.py 等模块直读 DOCGEN_LLM_*,
|
||||
# 不走 pydantic-settings 的 env_file,避免「.env 配了但 os.environ 读不到」的配置断层)。
|
||||
load_dotenv(Path(__file__).resolve().parent.parent / ".env")
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
@@ -17,11 +25,18 @@ from fastapi.middleware.cors import CORSMiddleware
|
||||
from .config import get_settings
|
||||
from .context import context, init_context, shutdown_context
|
||||
from .db import close_db, init_db
|
||||
from .routers import bioarchives, bioflow, biojobs
|
||||
from .mcp_server import extension_mcp, gene_family_mcp, rnaseq_mcp
|
||||
from .routers import bioarchives, bioflow, biojobs, extension
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 三个 MCP streamable-http app(调用即创建 session_manager,供 lifespan 进入 run())。
|
||||
# 端点地址:/mcp/gene-family、/mcp/rnaseq、/mcp/extension(streamable_http_path="/" + 挂载前缀)。
|
||||
_gene_family_mcp_app = gene_family_mcp.streamable_http_app(streamable_http_path="/")
|
||||
_rnaseq_mcp_app = rnaseq_mcp.streamable_http_app(streamable_http_path="/")
|
||||
_extension_mcp_app = extension_mcp.streamable_http_app(streamable_http_path="/")
|
||||
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
settings = get_settings()
|
||||
@@ -60,20 +75,34 @@ def create_app() -> FastAPI:
|
||||
logger.info("main: 重启恢复 %d 个 in-flight 任务续跑", len(recovered))
|
||||
except Exception:
|
||||
logger.exception("main: 重启恢复扫描失败(任务不重跑,等待后续操作)")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await shutdown_context()
|
||||
try:
|
||||
await close_db()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# MCP session manager 的 task group 必须在服务运行期内保持活跃
|
||||
# (streamable-http 处理请求时要求已初始化)。Starlette 不会自动运行
|
||||
# 子 Mount 的 lifespan,故在此显式进入三个 MCP 的 run()。
|
||||
async with gene_family_mcp.session_manager.run():
|
||||
async with rnaseq_mcp.session_manager.run():
|
||||
async with extension_mcp.session_manager.run():
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await shutdown_context()
|
||||
try:
|
||||
await close_db()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
app.router.lifespan_context = _lifespan
|
||||
|
||||
app.include_router(bioflow.router)
|
||||
app.include_router(biojobs.router)
|
||||
app.include_router(bioarchives.router)
|
||||
app.include_router(extension.router)
|
||||
|
||||
# 生信 MCP 能力外挂(mcp 2.x:MCPServer.streamable_http_app)。
|
||||
# 三个独立 streamable-http 端点,各自装卸(REST 面不受影响)。
|
||||
app.mount("/mcp/gene-family", _gene_family_mcp_app)
|
||||
app.mount("/mcp/rnaseq", _rnaseq_mcp_app)
|
||||
app.mount("/mcp/extension", _extension_mcp_app)
|
||||
return app
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,900 @@
|
||||
"""app/mcp_server.py —— 生信 MCP server(两个 streamable-http 端点)。
|
||||
|
||||
- `/mcp/gene-family`:基因家族能力单元(plan/create/ask + 共享控制组 + 文件组)
|
||||
- `/mcp/rnaseq`:转录组能力单元(plan/create/generate_samples_yaml/ask + 共享控制组 + 文件组)
|
||||
|
||||
设计约束(实施方案 S0):
|
||||
- 工具只承诺「已受理」,执行进度/日志由 SSE 事件总线直推,不走 MCP。
|
||||
- 每个能力单元只接受本链阶段(create 跨链拒绝 CROSS_CHAIN)。
|
||||
- 人工确认边界(subfamily_confirm / rnaseq_de 实验设计确认)在 job_status 的
|
||||
pending_actions 里标 requires_human=true,工具不提供绕过。
|
||||
- 白名单沿用 runner._allowed_param_names(单一事实源)。
|
||||
- 工具名不外泄由 QA agent(S1)负责,本模块只负责工具语义。
|
||||
|
||||
mcp 2.x 注意:MCPServer(非 FastMCP,来自 mcp.server.mcpserver);挂载用
|
||||
`mcp.streamable_http_app()`。本文件位于 app/ 层(非 services/),可用 pydantic(mcp
|
||||
自身也依赖 pydantic 构造 schema)。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import re
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .config import get_settings
|
||||
from .context import load_record
|
||||
from .services.bioflow import extensions as ext_mod
|
||||
from .services.bioflow import store
|
||||
from .services.bioflow.archive import archive_from_job
|
||||
from .services.bioflow.catalog import GENE_FAMILY_CHAIN
|
||||
from .services.bioflow.errors import error_result
|
||||
from .services.bioflow.file_transform import validate_format
|
||||
from .services.bioflow.runner import RunnerError, ensure_samples_yaml, llm_sample_fallback
|
||||
from .services.bioflow.source_fetcher import fetch_source
|
||||
|
||||
# 薄适配:委托 bioflowctl CLI(唯一事实源),不重复业务逻辑
|
||||
from bioflowctl.commands import (
|
||||
cmd_archive,
|
||||
cmd_archive_reuse,
|
||||
cmd_file_inspect,
|
||||
cmd_file_transform,
|
||||
cmd_source_fetch,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _mcp_err(code: str, message: str, *, evidence_refs: list[str] | None = None) -> dict[str, Any]:
|
||||
"""G13:MCP 工具失败返回统一结构化错误对象(code 顶层 + error 对象)。"""
|
||||
return error_result(code, message, evidence_refs=evidence_refs)
|
||||
|
||||
# ---- 能力单元阶段白名单 ----
|
||||
|
||||
GENE_FAMILY_STAGES = set(GENE_FAMILY_CHAIN)
|
||||
|
||||
RNASEQ_STAGES = {
|
||||
"rnaseq_qc", "rnaseq_align", "rnaseq_quant", "rnaseq_sample_qc", "rnaseq_de",
|
||||
"rnaseq_wgcna", "rnaseq_go_enrich", "rnaseq_kegg_enrich", "rnaseq_enrich",
|
||||
"rnaseq_fgsea", "rnaseq_candidate_family", "rnaseq_report",
|
||||
"annotation", "expression",
|
||||
}
|
||||
|
||||
# 人工确认边界:这些阶段 waiting_confirm 时 requires_human=true,工具/QA 不可代确认
|
||||
HUMAN_CONFIRM_STAGES = {"subfamily_confirm", "rnaseq_de"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# context 访问(动态读模块全局,测试 monkeypatch 兼容)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _ctx():
|
||||
from . import context as context_mod
|
||||
|
||||
return context_mod.context
|
||||
|
||||
|
||||
def _job_detail(record: Any) -> dict[str, Any]:
|
||||
stages = []
|
||||
for stage in record.stage_chain:
|
||||
st = record._stage_states.get(stage, {})
|
||||
stages.append({
|
||||
"stage": stage,
|
||||
"status": st.get("status", "waiting"),
|
||||
"params": st.get("params") or {},
|
||||
"attempt": int(st.get("attempt") or 0),
|
||||
"manifest_path": st.get("manifest_path"),
|
||||
"log_tail": st.get("log_tail") or "",
|
||||
"error": st.get("error"),
|
||||
})
|
||||
return {
|
||||
"id": record.id,
|
||||
"session_id": record.session_id,
|
||||
"title": record.title,
|
||||
"status": record.status,
|
||||
"stages": stages,
|
||||
"params": record.params,
|
||||
"auto_continue": record.auto_continue,
|
||||
"archive_id": record.archive_id,
|
||||
"progress": record.progress(),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# pending_actions 派生(job_status 工具返回给 QA agent 的下一步动作)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def derive_pending_actions(record: Any, runner: Any) -> list[dict[str, Any]]:
|
||||
"""由任务/阶段状态派生待办动作(QA agent 据此引导用户,人工边界不可代确认)。
|
||||
|
||||
按阶段链顺序找第一个「卡住」的环节,返回动作清单(通常 0-1 条)。
|
||||
"""
|
||||
if record.status == "paused":
|
||||
return [{"action": "resume", "stage": None, "requires_human": False,
|
||||
"message": "任务已暂停,可恢复"}]
|
||||
if record.status == "done":
|
||||
return [{"action": "archive", "stage": None, "requires_human": False,
|
||||
"message": "任务已完成,可存档复用"}]
|
||||
if record.status == "archived":
|
||||
return [{"action": "none", "stage": None, "requires_human": False,
|
||||
"message": "任务已存档"}]
|
||||
|
||||
for stage in record.stage_chain:
|
||||
status = record._stage_states.get(stage, {}).get("status", "waiting")
|
||||
if status == "await_files":
|
||||
return [{"action": "provide_files", "stage": stage, "requires_human": False,
|
||||
"message": f"阶段 {stage} 等待输入文件"}]
|
||||
if status == "waiting_confirm":
|
||||
human = stage in HUMAN_CONFIRM_STAGES
|
||||
return [{"action": "confirm", "stage": stage, "requires_human": human,
|
||||
"form_defaults": runner._form_defaults(stage, record),
|
||||
"message": (f"阶段 {stage} 需人工确认(不可代确认)" if human
|
||||
else f"阶段 {stage} 等待参数确认")}]
|
||||
if status == "failed":
|
||||
return [{"action": "rerun_or_skip", "stage": stage, "requires_human": False,
|
||||
"message": f"阶段 {stage} 失败,可重跑或跳过"}]
|
||||
if status == "partial_success":
|
||||
# G4:partial_success 独立终态,不被当作 success/done
|
||||
return [{"action": "rerun_or_skip", "stage": stage, "requires_human": False,
|
||||
"message": f"阶段 {stage} 部分成功(产物不全),可重跑或跳过"}]
|
||||
# waiting/running/queued/done/skipped → 继续
|
||||
return []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 工具实现(命名 async 函数,带类型提示,供两个 MCPServer 注册)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def _plan(message: str, allowed_stages: set[str]) -> dict[str, Any]:
|
||||
from .services.bioflow.planner import PlanCandidate, PlannerError, PlanResult
|
||||
|
||||
ctx = _ctx()
|
||||
if ctx.planner is None:
|
||||
return {"ok": False, "code": "NOT_READY", "message": "规划器未就绪"}
|
||||
try:
|
||||
plan = await ctx.planner.plan(message)
|
||||
except PlannerError as exc:
|
||||
fallback = [s for s in ("identify", "rnaseq_qc") if s in allowed_stages and ctx.catalog.has(s)]
|
||||
plan = PlanResult(candidates=[PlanCandidate(purpose="内置默认流程", stages=fallback)],
|
||||
source="llm_only", message=exc.message)
|
||||
candidates = []
|
||||
for c in plan.candidates:
|
||||
stages = [s for s in c.stages if s in allowed_stages]
|
||||
if stages:
|
||||
candidates.append({"purpose": c.purpose, "stages": stages})
|
||||
return {
|
||||
"ok": True,
|
||||
"candidates": candidates,
|
||||
"source": plan.source,
|
||||
"reuse_archive": plan.reuse_archive_id,
|
||||
"reuse_name": plan.reuse_name,
|
||||
"fallback": plan.message or None,
|
||||
}
|
||||
|
||||
|
||||
async def _create_job(session_id: str, title: str, stages: list[str],
|
||||
params: dict[str, Any] | None, auto_continue: bool,
|
||||
archive_id: str | None, allowed_stages: set[str]) -> dict[str, Any]:
|
||||
ctx = _ctx()
|
||||
unknown = [s for s in stages if s not in allowed_stages]
|
||||
if unknown:
|
||||
return {"ok": False, "code": "CROSS_CHAIN",
|
||||
"message": f"阶段不属于该能力单元: {', '.join(unknown)}"}
|
||||
try:
|
||||
record = await ctx.runner.create_job(
|
||||
session_id, title or "生信任务", stages, params or {},
|
||||
auto_continue=auto_continue, archive_id=archive_id,
|
||||
)
|
||||
except RunnerError as exc:
|
||||
return {"ok": False, "code": exc.code, "message": exc.message}
|
||||
ctx.records[record.id] = record
|
||||
return {"ok": True, "data": _job_detail(record)}
|
||||
|
||||
|
||||
async def _job_status(job_id: str) -> dict[str, Any]:
|
||||
ctx = _ctx()
|
||||
record = await load_record(job_id)
|
||||
if record is None:
|
||||
return {"ok": False, "code": "NOT_FOUND", "message": "任务不存在"}
|
||||
detail = _job_detail(record)
|
||||
detail["pending_actions"] = derive_pending_actions(record, ctx.runner)
|
||||
return {"ok": True, "data": detail}
|
||||
|
||||
|
||||
async def _confirm_stage(job_id: str, stage: str, params: dict[str, Any] | None) -> dict[str, Any]:
|
||||
ctx = _ctx()
|
||||
record = await load_record(job_id)
|
||||
if record is None:
|
||||
return {"ok": False, "code": "NOT_FOUND", "message": "任务不存在"}
|
||||
if stage not in record.stage_chain:
|
||||
return {"ok": False, "code": "NOT_FOUND", "message": f"阶段不存在: {stage}"}
|
||||
try:
|
||||
await ctx.runner.confirm_stage(record, stage, params or {})
|
||||
except RunnerError as exc:
|
||||
return {"ok": False, "code": exc.code, "message": exc.message}
|
||||
return {"ok": True, "data": {"job_id": job_id, "stage": stage,
|
||||
"status": ctx.runner.stage_status(record, stage)}}
|
||||
|
||||
|
||||
async def _rerun_stage(job_id: str, stage: str, params: dict[str, Any] | None) -> dict[str, Any]:
|
||||
ctx = _ctx()
|
||||
record = await load_record(job_id)
|
||||
if record is None:
|
||||
return {"ok": False, "code": "NOT_FOUND", "message": "任务不存在"}
|
||||
if stage not in record.stage_chain:
|
||||
return {"ok": False, "code": "NOT_FOUND", "message": f"阶段不存在: {stage}"}
|
||||
try:
|
||||
await ctx.runner.rerun_stage(record, stage, params)
|
||||
except RunnerError as exc:
|
||||
return {"ok": False, "code": exc.code, "message": exc.message}
|
||||
return {"ok": True, "data": {"job_id": job_id, "stage": stage,
|
||||
"status": ctx.runner.stage_status(record, stage)}}
|
||||
|
||||
|
||||
async def _skip_stage(job_id: str, stage: str) -> dict[str, Any]:
|
||||
ctx = _ctx()
|
||||
record = await load_record(job_id)
|
||||
if record is None:
|
||||
return {"ok": False, "code": "NOT_FOUND", "message": "任务不存在"}
|
||||
if stage not in record.stage_chain:
|
||||
return {"ok": False, "code": "NOT_FOUND", "message": f"阶段不存在: {stage}"}
|
||||
try:
|
||||
await ctx.runner.skip_stage(record, stage)
|
||||
except RunnerError as exc:
|
||||
return {"ok": False, "code": exc.code, "message": exc.message}
|
||||
return {"ok": True, "data": {"job_id": job_id, "stage": stage,
|
||||
"status": ctx.runner.stage_status(record, stage)}}
|
||||
|
||||
|
||||
async def _pause_job(job_id: str) -> dict[str, Any]:
|
||||
ctx = _ctx()
|
||||
record = await load_record(job_id)
|
||||
if record is None:
|
||||
return {"ok": False, "code": "NOT_FOUND", "message": "任务不存在"}
|
||||
try:
|
||||
await ctx.runner.pause_job(record)
|
||||
except RunnerError as exc:
|
||||
return {"ok": False, "code": exc.code, "message": exc.message}
|
||||
return {"ok": True, "data": {"job_id": job_id, "status": record.status}}
|
||||
|
||||
|
||||
async def _resume_job(job_id: str) -> dict[str, Any]:
|
||||
ctx = _ctx()
|
||||
record = await load_record(job_id)
|
||||
if record is None:
|
||||
return {"ok": False, "code": "NOT_FOUND", "message": "任务不存在"}
|
||||
try:
|
||||
await ctx.runner.resume_job(record)
|
||||
except RunnerError as exc:
|
||||
return {"ok": False, "code": exc.code, "message": exc.message}
|
||||
return {"ok": True, "data": {"job_id": job_id, "status": record.status}}
|
||||
|
||||
|
||||
async def _archive_job(job_id: str, name: str, description: str,
|
||||
include_intermediate: bool) -> dict[str, Any]:
|
||||
"""任务存档(薄适配):委托 cmd_archive(唯一事实源),仅保留内存 record 校验 + 事件发布。"""
|
||||
ctx = _ctx()
|
||||
record = await load_record(job_id)
|
||||
if record is None:
|
||||
return _mcp_err("NOT_FOUND", "任务不存在")
|
||||
result = await cmd_archive(
|
||||
job_id, name, description,
|
||||
db=ctx.db,
|
||||
record=record,
|
||||
include_intermediate=bool(include_intermediate),
|
||||
archive_fn=archive_from_job,
|
||||
archives_dir=get_settings().archives_dir_abs,
|
||||
)
|
||||
if result.get("ok"):
|
||||
ctx.events.publish("job_archived", {"job_id": job_id, "archive_id": result["data"]["id"]})
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 扩展能力单元(第三个能力单元 bioflow-extension):复用 extensions.py 状态机
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_EXT_VIEW_KEYS = (
|
||||
"ext_id", "display_name", "status", "command_template", "original_command",
|
||||
"runtime_image", "requirements", "parameters", "input_roles", "nl_triggers",
|
||||
"analysis_type", "golden", "repair_count", "trial", "validation", "last_build",
|
||||
"created_at", "updated_at",
|
||||
)
|
||||
|
||||
|
||||
def _ext_view(ext: dict[str, Any]) -> dict[str, Any]:
|
||||
"""扩展精简视图(剔除 source_files 等超长字段,供 MCP 工具结果返回)。"""
|
||||
return {k: ext.get(k) for k in _EXT_VIEW_KEYS if k in ext}
|
||||
|
||||
|
||||
def _file_items(files: Any) -> list[tuple[str, bytes]]:
|
||||
"""把 MCP 入参 [{path, content|data}] 转为 [(rel_path, bytes)]。"""
|
||||
out: list[tuple[str, bytes]] = []
|
||||
if not isinstance(files, list):
|
||||
return out
|
||||
for f in files:
|
||||
if not isinstance(f, dict):
|
||||
continue
|
||||
path = str(f.get("path") or f.get("name") or "").strip()
|
||||
if not path:
|
||||
continue
|
||||
content = f.get("content")
|
||||
if content is None:
|
||||
content = f.get("data") or ""
|
||||
if isinstance(content, (bytes, bytearray)):
|
||||
blob = bytes(content)
|
||||
else:
|
||||
blob = str(content).encode("utf-8")
|
||||
out.append((path, blob))
|
||||
return out
|
||||
|
||||
|
||||
def _ext_err(exc: Exception) -> str:
|
||||
if isinstance(exc, ext_mod.ExtensionNotFoundError):
|
||||
return "NOT_FOUND"
|
||||
if isinstance(exc, ext_mod.ExtensionStatusError):
|
||||
return "EXT_STATUS"
|
||||
if isinstance(exc, ext_mod.ExtensionBuildError):
|
||||
return "EXT_BUILD"
|
||||
if isinstance(exc, ext_mod.ExtensionRunError):
|
||||
return "EXT_RUN"
|
||||
return "EXT_ERROR"
|
||||
|
||||
|
||||
async def _create_extension(display_name: str, command: str, runtime_image: str,
|
||||
code_files: Any, data_files: Any) -> dict[str, Any]:
|
||||
code = _file_items(code_files)
|
||||
if not code:
|
||||
return {"ok": False, "code": "VALIDATION", "message": "扩展代码包不能为空"}
|
||||
try:
|
||||
ext = ext_mod.get_store().create_extension(
|
||||
display_name=display_name, command=command, runtime_image=runtime_image or "",
|
||||
code_files=code, data_files=_file_items(data_files),
|
||||
)
|
||||
except ext_mod.ExtensionError as exc:
|
||||
return {"ok": False, "code": _ext_err(exc), "message": str(exc)}
|
||||
return {"ok": True, "data": _ext_view(ext)}
|
||||
|
||||
|
||||
async def _extension_status(ext_id: str) -> dict[str, Any]:
|
||||
try:
|
||||
ext = ext_mod.get_store().load_extension(ext_id)
|
||||
except ext_mod.ExtensionNotFoundError:
|
||||
return {"ok": False, "code": "NOT_FOUND", "message": "扩展不存在"}
|
||||
return {"ok": True, "data": _ext_view(ext) if ext else None}
|
||||
|
||||
|
||||
async def _build_extension(ext_id: str) -> dict[str, Any]:
|
||||
try:
|
||||
ext = await ext_mod.get_store().build_extension(ext_id)
|
||||
except ext_mod.ExtensionError as exc:
|
||||
return {"ok": False, "code": _ext_err(exc), "message": str(exc)}
|
||||
return {"ok": True, "data": _ext_view(ext)}
|
||||
|
||||
|
||||
async def _trial_run_extension(ext_id: str) -> dict[str, Any]:
|
||||
try:
|
||||
ext = await ext_mod.get_store().trial_run(ext_id)
|
||||
except ext_mod.ExtensionError as exc:
|
||||
return {"ok": False, "code": _ext_err(exc), "message": str(exc)}
|
||||
return {"ok": True, "data": _ext_view(ext)}
|
||||
|
||||
|
||||
async def _validate_extension(ext_id: str) -> dict[str, Any]:
|
||||
try:
|
||||
ext = ext_mod.get_store().validate(ext_id)
|
||||
except ext_mod.ExtensionError as exc:
|
||||
return {"ok": False, "code": _ext_err(exc), "message": str(exc)}
|
||||
return {"ok": True, "data": _ext_view(ext)}
|
||||
|
||||
|
||||
async def _confirm_extension(ext_id: str, nl_triggers: Any, parameters: Any,
|
||||
runtime_image: str) -> dict[str, Any]:
|
||||
triggers = [str(x).strip() for x in nl_triggers if str(x).strip()] if isinstance(nl_triggers, list) else None
|
||||
params = parameters if isinstance(parameters, list) else None
|
||||
try:
|
||||
ext = ext_mod.get_store().confirm(
|
||||
ext_id, nl_triggers=triggers, parameters=params,
|
||||
runtime_image=runtime_image or None,
|
||||
)
|
||||
except ext_mod.ExtensionError as exc:
|
||||
return {"ok": False, "code": _ext_err(exc), "message": str(exc)}
|
||||
return {"ok": True, "data": _ext_view(ext)}
|
||||
|
||||
|
||||
async def _publish_extension(ext_id: str) -> dict[str, Any]:
|
||||
try:
|
||||
ext = ext_mod.get_store().publish(ext_id)
|
||||
except ext_mod.ExtensionError as exc:
|
||||
return {"ok": False, "code": _ext_err(exc), "message": str(exc)}
|
||||
return {"ok": True, "data": _ext_view(ext)}
|
||||
|
||||
|
||||
async def _run_extension(ext_id: str, params: Any, files: Any) -> dict[str, Any]:
|
||||
try:
|
||||
result = await ext_mod.get_store().run(
|
||||
ext_id,
|
||||
params=params if isinstance(params, dict) else None,
|
||||
files=[str(f) for f in files] if isinstance(files, list) else None,
|
||||
)
|
||||
except ext_mod.ExtensionError as exc:
|
||||
return {"ok": False, "code": _ext_err(exc), "message": str(exc)}
|
||||
return {"ok": True, "data": result}
|
||||
|
||||
|
||||
async def _repair_extension(ext_id: str) -> dict[str, Any]:
|
||||
try:
|
||||
ext = await ext_mod.get_store().repair(ext_id)
|
||||
except ext_mod.ExtensionError as exc:
|
||||
return {"ok": False, "code": _ext_err(exc), "message": str(exc)}
|
||||
return {"ok": True, "data": _ext_view(ext)}
|
||||
|
||||
|
||||
# ---- 文件登记(复刻 biojobs._register_path 核心,无 FastAPI 依赖)----
|
||||
|
||||
def _blocked_reason(path: Path, settings: Any) -> str | None:
|
||||
s = str(path)
|
||||
for prefix in settings.sensitive_path_prefixes:
|
||||
if s == prefix or s.startswith(prefix + "/"):
|
||||
return f"敏感路径前缀 {prefix}"
|
||||
for pat in settings.sensitive_filename_patterns:
|
||||
if re.match(pat, path.name):
|
||||
return f"敏感文件名 {path.name}"
|
||||
return None
|
||||
|
||||
|
||||
async def _register_job_files(job_id: str, paths: list[str]) -> dict[str, Any]:
|
||||
"""登记服务器/局域网路径(单文件或目录展开)到任务 input。"""
|
||||
ctx = _ctx()
|
||||
settings = get_settings()
|
||||
record = await load_record(job_id)
|
||||
if record is None:
|
||||
return _mcp_err("NOT_FOUND", "任务不存在")
|
||||
|
||||
input_dir = record.input_dir
|
||||
input_dir.mkdir(parents=True, exist_ok=True)
|
||||
items: list[dict[str, Any]] = []
|
||||
registered: list[str] = []
|
||||
|
||||
for raw in paths:
|
||||
p = Path(raw).expanduser().resolve()
|
||||
if not p.exists():
|
||||
return _mcp_err("PATH_NOT_EXIST", f"路径不存在: {raw}")
|
||||
blocked = _blocked_reason(p, settings)
|
||||
if blocked:
|
||||
return _mcp_err("PATH_NOT_ALLOWED", f"路径不允许: {raw}({blocked})")
|
||||
files: list[Path] = []
|
||||
if p.is_file():
|
||||
files = [p]
|
||||
else:
|
||||
ext_set = settings.upload_ext_set
|
||||
for f in sorted(p.rglob("*")):
|
||||
if f.is_file() and (not ext_set or f.suffix.lower() in ext_set) and not _blocked_reason(f, settings):
|
||||
files.append(f)
|
||||
if not files:
|
||||
return _mcp_err("VALIDATION", f"目录内无可登记文件: {raw}")
|
||||
for f in files:
|
||||
link = input_dir / f.name
|
||||
try:
|
||||
if not link.exists():
|
||||
link.symlink_to(f)
|
||||
except OSError as exc: # noqa: BLE001
|
||||
logger.warning("mcp register symlink 失败: %s -> %s (%s)", f, link, exc)
|
||||
items.append({"file_id": f"{record.id}:{f.name}", "name": f.name,
|
||||
"size": f.stat().st_size, "path": str(f), "source": "server_path"})
|
||||
registered.append(str(f))
|
||||
|
||||
# rnaseq 系:samples.yaml 生成(描述文件优先 + 实况分析;规则不可归纳走 LLM 兜底)
|
||||
try:
|
||||
source_roots = [str(Path(x).parent) for x in paths if Path(x).exists()]
|
||||
extra_mounts = ensure_samples_yaml(input_dir, source_roots=[Path(r) for r in source_roots])
|
||||
if not extra_mounts:
|
||||
try:
|
||||
await llm_sample_fallback(input_dir, [Path(r) for r in source_roots])
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("samples.yaml LLM 兜底失败: %s", exc)
|
||||
if extra_mounts:
|
||||
record.registered_server_paths = list(
|
||||
dict.fromkeys((getattr(record, "registered_server_paths", None) or [])
|
||||
+ [str(Path(m)) for m in extra_mounts]))
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("samples.yaml 分析失败: %s", exc)
|
||||
|
||||
if registered:
|
||||
import json
|
||||
record.registered_server_paths = list(
|
||||
dict.fromkeys((getattr(record, "registered_server_paths", None) or []) + registered)
|
||||
)
|
||||
(record.job_dir / "server_registered.json").write_text(
|
||||
json.dumps({"registered_server_paths": record.registered_server_paths},
|
||||
ensure_ascii=False), encoding="utf-8")
|
||||
|
||||
# 自动续跑(await_inputs 任务重推;非等待态 no-op)
|
||||
await ctx.runner.retry_after_files(record)
|
||||
return {"ok": True, "data": {"items": items, "total": len(items)}}
|
||||
|
||||
|
||||
async def _ask_question(question: str) -> dict[str, Any]:
|
||||
ctx = _ctx()
|
||||
hits: list[dict[str, Any]] = []
|
||||
if ctx.flow_store is not None:
|
||||
try:
|
||||
hits = ctx.flow_store.search(question, top_k=8).get("items") or []
|
||||
except Exception: # noqa: BLE001
|
||||
hits = []
|
||||
citations = []
|
||||
for h in hits:
|
||||
citations.append({
|
||||
"source": h.get("source") or h.get("doc_kind") or "",
|
||||
"title_path": h.get("title_path") or "",
|
||||
"path": h.get("path") or "",
|
||||
"text": str(h.get("text") or "")[:1200],
|
||||
})
|
||||
return {"ok": True, "citations": citations, "degraded": not citations}
|
||||
|
||||
|
||||
async def _list_archives(page: int, page_size: int) -> dict[str, Any]:
|
||||
ctx = _ctx()
|
||||
items, total = await store.list_archives(ctx.db, page=page, page_size=page_size)
|
||||
return {"ok": True, "items": items, "total": total, "page": page, "page_size": page_size}
|
||||
|
||||
|
||||
async def _reuse_archive(name: str) -> dict[str, Any]:
|
||||
"""复用存档(薄适配):委托 cmd_archive_reuse(读各阶段实际生效参数)。"""
|
||||
return await cmd_archive_reuse(name, db=_ctx().db)
|
||||
|
||||
|
||||
async def _fetch_source_file(job_id: str | None, source: str,
|
||||
credentials: dict[str, Any] | None) -> dict[str, Any]:
|
||||
"""拉取源文件(薄适配):委托 cmd_source_fetch,仅保留内存 record 校验。"""
|
||||
if job_id:
|
||||
record = await load_record(job_id)
|
||||
if record is None:
|
||||
return _mcp_err("NOT_FOUND", "任务不存在")
|
||||
return await cmd_source_fetch(source, job_id=job_id, credentials=credentials,
|
||||
fetch_fn=fetch_source)
|
||||
|
||||
|
||||
async def _inspect_file(path: str) -> dict[str, Any]:
|
||||
"""探测文件(薄适配):委托 cmd_file_inspect。"""
|
||||
return cmd_file_inspect(path)
|
||||
|
||||
|
||||
async def _transform_file(src_path: str, target_format: str,
|
||||
column_map: dict[str, str] | None,
|
||||
dest_path: str | None) -> dict[str, Any]:
|
||||
"""转换文件(薄适配):委托 cmd_file_transform(注入 LLM 适配器)。"""
|
||||
from .services.bioflow.llm import complete_json
|
||||
|
||||
async def llm_adapter(user: str, system: str) -> Any:
|
||||
return await complete_json(system, user, temperature=0.2, max_tokens=2000)
|
||||
|
||||
return await cmd_file_transform(
|
||||
src_path, target_format, column_map=column_map, dest=dest_path, llm=llm_adapter,
|
||||
)
|
||||
|
||||
|
||||
async def _generate_samples_yaml(job_id: str, sample_description: str) -> dict[str, Any]:
|
||||
import yaml
|
||||
|
||||
from .services.bioflow.llm import complete_json
|
||||
|
||||
record = await load_record(job_id)
|
||||
if record is None:
|
||||
return {"ok": False, "code": "NOT_FOUND", "message": "任务不存在"}
|
||||
system = (
|
||||
"你是 RNA-seq 样本配置助手。根据用户的自然语言描述,输出 samples.yaml 结构 JSON:"
|
||||
'{"samples": {"样本名": {"r1": "fastq路径", "r2": "fastq路径", "condition": "...", ...}}}。'
|
||||
"只输出 JSON。"
|
||||
)
|
||||
try:
|
||||
raw = await complete_json(system, sample_description, temperature=0.2, max_tokens=2000)
|
||||
samples = raw.get("samples") if isinstance(raw, dict) else None
|
||||
if not isinstance(samples, dict) or not samples:
|
||||
return {"ok": False, "code": "EMPTY_SAMPLES", "message": "未能从描述解析出样本"}
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return {"ok": False, "code": "SAMPLES_GEN_FAILED", "message": f"样本配置生成失败: {exc}"}
|
||||
record.input_dir.mkdir(parents=True, exist_ok=True)
|
||||
out = record.input_dir / "samples.yaml"
|
||||
out.write_text(yaml.safe_dump({"samples": samples}, allow_unicode=True), encoding="utf-8")
|
||||
v = validate_format(out, "samples.yaml")
|
||||
if not v.get("ok"):
|
||||
out.unlink(missing_ok=True)
|
||||
return {"ok": False, "code": "VALIDATION_FAILED", "message": f"样本配置校验失败: {v.get('reason')}"}
|
||||
return {"ok": True, "data": {"path": str(out), "samples": samples, "validated": v}}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 能力单元工具(命名 async,带类型提示 + 中文描述)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def plan_gene_family(message: str) -> dict[str, Any]:
|
||||
"""根据用户自然语言需求规划基因家族分析候选路径(仅基因家族链阶段)。"""
|
||||
return await _plan(message, GENE_FAMILY_STAGES)
|
||||
|
||||
|
||||
async def create_gene_family_job(session_id: str, stages: list[str], title: str = "",
|
||||
params: dict[str, Any] | None = None,
|
||||
auto_continue: bool = False,
|
||||
archive_id: str | None = None) -> dict[str, Any]:
|
||||
"""按确认的阶段链创建基因家族任务(stages 只能是基因家族链阶段,跨链拒绝)。"""
|
||||
return await _create_job(session_id, title, stages, params, auto_continue, archive_id, GENE_FAMILY_STAGES)
|
||||
|
||||
|
||||
async def ask_gene_family_question(question: str) -> dict[str, Any]:
|
||||
"""回答基因家族流程/参数/生物学意义类问题(流程知识库检索,带引用)。"""
|
||||
return await _ask_question(question)
|
||||
|
||||
|
||||
async def plan_rnaseq(message: str) -> dict[str, Any]:
|
||||
"""根据用户自然语言需求规划转录组分析候选路径(仅转录组链阶段)。"""
|
||||
return await _plan(message, RNASEQ_STAGES)
|
||||
|
||||
|
||||
async def create_rnaseq_job(session_id: str, stages: list[str], title: str = "",
|
||||
params: dict[str, Any] | None = None,
|
||||
auto_continue: bool = False,
|
||||
archive_id: str | None = None) -> dict[str, Any]:
|
||||
"""按确认的阶段链创建转录组任务(stages 只能是转录组链阶段,跨链拒绝)。"""
|
||||
return await _create_job(session_id, title, stages, params, auto_continue, archive_id, RNASEQ_STAGES)
|
||||
|
||||
|
||||
async def generate_samples_yaml(job_id: str, sample_description: str) -> dict[str, Any]:
|
||||
"""把用户的自然语言样本描述(样本名/对照处理组/重复/单双端等)解析为 samples.yaml 写入任务 input。"""
|
||||
return await _generate_samples_yaml(job_id, sample_description)
|
||||
|
||||
|
||||
async def ask_rnaseq_question(question: str) -> dict[str, Any]:
|
||||
"""回答转录组流程/参数/实验设计类问题(流程知识库检索,带引用)。"""
|
||||
return await _ask_question(question)
|
||||
|
||||
|
||||
async def job_status(job_id: str) -> dict[str, Any]:
|
||||
"""查询生信任务状态:任务/各阶段状态 + 下一步待办动作(pending_actions)。卡住或需确认时用本工具查真相。"""
|
||||
return await _job_status(job_id)
|
||||
|
||||
|
||||
async def confirm_stage(job_id: str, stage: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
"""确认某阶段参数并执行。仅 waiting_confirm 状态可确认;参数名须在白名单内(未知参数被拒)。人工确认边界阶段(subfamily_confirm/rnaseq_de)不可代用户确认。"""
|
||||
return await _confirm_stage(job_id, stage, params)
|
||||
|
||||
|
||||
async def register_job_files(job_id: str, paths: list[str]) -> dict[str, Any]:
|
||||
"""把服务器/局域网路径(单文件或目录)登记为任务输入文件。目录按扩展名白名单展开。"""
|
||||
return await _register_job_files(job_id, paths)
|
||||
|
||||
|
||||
async def rerun_stage(job_id: str, stage: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
"""重跑失败阶段(带参数则先等确认)。仅 failed 状态可重跑。"""
|
||||
return await _rerun_stage(job_id, stage, params)
|
||||
|
||||
|
||||
async def skip_stage(job_id: str, stage: str) -> dict[str, Any]:
|
||||
"""跳过失败阶段并推进后续。仅 failed 状态可跳过。"""
|
||||
return await _skip_stage(job_id, stage)
|
||||
|
||||
|
||||
async def pause_job(job_id: str) -> dict[str, Any]:
|
||||
"""暂停运行中的任务(running 阶段会立即 kill 容器)。"""
|
||||
return await _pause_job(job_id)
|
||||
|
||||
|
||||
async def resume_job(job_id: str) -> dict[str, Any]:
|
||||
"""恢复已暂停的任务(从下一阶段继续,确认点不跳过)。"""
|
||||
return await _resume_job(job_id)
|
||||
|
||||
|
||||
async def list_archives(page: int = 1, page_size: int = 20) -> dict[str, Any]:
|
||||
"""列出历史存档任务(供复用参考)。"""
|
||||
return await _list_archives(page, page_size)
|
||||
|
||||
|
||||
async def reuse_archive(name: str) -> dict[str, Any]:
|
||||
"""按名称匹配存档,返回其阶段链+参数(供复用发起)。"""
|
||||
return await _reuse_archive(name)
|
||||
|
||||
|
||||
async def archive_job(job_id: str, name: str, description: str = "",
|
||||
include_intermediate: bool = False) -> dict[str, Any]:
|
||||
"""把已完成任务存档(打包结果/日志/事件 + 落档登记);重名拒绝(ARCHIVE_EXISTS)。"""
|
||||
return await _archive_job(job_id, name, description, include_intermediate)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 扩展能力单元工具(第三个能力单元 bioflow-extension)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def create_extension(display_name: str, command: str, runtime_image: str = "",
|
||||
code_files: list[dict[str, Any]] | None = None,
|
||||
data_files: list[dict[str, Any]] | None = None) -> dict[str, Any]:
|
||||
"""创建扩展草稿(draft):上传代码文件清单(每项 {path, content 文本}) + 运行命令 + 显示名。"""
|
||||
return await _create_extension(display_name, command, runtime_image, code_files, data_files)
|
||||
|
||||
|
||||
async def extension_status(ext_id: str) -> dict[str, Any]:
|
||||
"""查询扩展详情与状态(draft/build/trial/validated/confirming/published)。"""
|
||||
return await _extension_status(ext_id)
|
||||
|
||||
|
||||
async def build_extension(ext_id: str) -> dict[str, Any]:
|
||||
"""AI 转换扩展(原始命令→command_template+参数+dependencies);LLM 失败自动回退规则版。"""
|
||||
return await _build_extension(ext_id)
|
||||
|
||||
|
||||
async def trial_run_extension(ext_id: str) -> dict[str, Any]:
|
||||
"""试运行扩展(docker 执行;未接入执行器时返回 EXT_RUN 503 语义)。"""
|
||||
return await _trial_run_extension(ext_id)
|
||||
|
||||
|
||||
async def validate_extension(ext_id: str) -> dict[str, Any]:
|
||||
"""校验扩展(金标/试运行校验)。仅 built/trial_passed 可校验。"""
|
||||
return await _validate_extension(ext_id)
|
||||
|
||||
|
||||
async def confirm_extension(ext_id: str, nl_triggers: list[str] | None = None,
|
||||
parameters: list[dict[str, Any]] | None = None,
|
||||
runtime_image: str = "") -> dict[str, Any]:
|
||||
"""确认扩展参数/触发词(用户可编辑 nl_triggers);仅 validated/trial_passed 可确认。"""
|
||||
return await _confirm_extension(ext_id, nl_triggers, parameters, runtime_image)
|
||||
|
||||
|
||||
async def publish_extension(ext_id: str) -> dict[str, Any]:
|
||||
"""发布扩展(仅 published 状态可被自然语言触发运行)。"""
|
||||
return await _publish_extension(ext_id)
|
||||
|
||||
|
||||
async def run_extension(ext_id: str, params: dict[str, Any] | None = None,
|
||||
files: list[str] | None = None) -> dict[str, Any]:
|
||||
"""正式运行已发布扩展(同步返回 run_result)。"""
|
||||
return await _run_extension(ext_id, params, files)
|
||||
|
||||
|
||||
async def repair_extension(ext_id: str) -> dict[str, Any]:
|
||||
"""AI 修复扩展(试运行失败上下文→LLM 修复 command_template/依赖;LLM 失败回退规则版)。"""
|
||||
return await _repair_extension(ext_id)
|
||||
|
||||
|
||||
async def fetch_source_file(source: str, job_id: str | None = None,
|
||||
credentials: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
"""从 web(http/https 流式)、OSS(oss:// 或预签名 URL)或局域网路径拉取源文件到任务 input;下载前检查磁盘余量并记录 provenance。"""
|
||||
return await _fetch_source_file(job_id, source, credentials)
|
||||
|
||||
|
||||
async def inspect_file(path: str) -> dict[str, Any]:
|
||||
"""探测文件格式(扩展名仅参考):返回预览/列结构/大小/格式猜测/乱码检测,供判断是否需要预处理。"""
|
||||
return await _inspect_file(path)
|
||||
|
||||
|
||||
async def transform_file(src_path: str, target_format: str,
|
||||
column_map: dict[str, str] | None = None,
|
||||
dest_path: str | None = None) -> dict[str, Any]:
|
||||
"""把非标准文件(xlsx/csv/tsv/txt)转换为 fasta/gff3/tsv/samples.yaml;LLM 定列映射规则+脚本执行,结果过格式验证器。"""
|
||||
return await _transform_file(src_path, target_format, column_map, dest_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# G12:11 个分析控制 MCP 工具(异步薄封装委托 skill_tools,不写业务逻辑)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def get_task_context(job_id: str | None = None) -> dict[str, Any]:
|
||||
from .services.bioflow import skill_tools
|
||||
return skill_tools.get_task_context(job_id)
|
||||
|
||||
|
||||
async def inspect_asset(role: str = "pep.fa", path: str = "") -> dict[str, Any]:
|
||||
from .services.bioflow import skill_tools
|
||||
return skill_tools.inspect_asset(role, path)
|
||||
|
||||
|
||||
async def validate_plan(stages: list[str] | None = None) -> dict[str, Any]:
|
||||
from .services.bioflow import skill_tools
|
||||
return skill_tools.validate_plan(stages)
|
||||
|
||||
|
||||
async def check_readiness(assets_root: str = "", stages: list[str] | None = None) -> dict[str, Any]:
|
||||
from .services.bioflow import skill_tools
|
||||
return skill_tools.check_readiness(assets_root, stages)
|
||||
|
||||
|
||||
async def get_run_evidence(run_id: str = "") -> dict[str, Any]:
|
||||
from .services.bioflow import skill_tools
|
||||
return skill_tools.get_run_evidence(run_id)
|
||||
|
||||
|
||||
async def propose_change(job_id: str = "", stage: str = "",
|
||||
params: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
from .services.bioflow import skill_tools
|
||||
return skill_tools.propose_change(job_id, stage, params)
|
||||
|
||||
|
||||
async def apply_change(action_id: str = "", change: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
from .services.bioflow import skill_tools
|
||||
return skill_tools.apply_change(action_id, change)
|
||||
|
||||
|
||||
async def execute_nodes(job_id: str = "", nodes: list[str] | None = None) -> dict[str, Any]:
|
||||
from .services.bioflow import skill_tools
|
||||
return skill_tools.execute_nodes(job_id, nodes)
|
||||
|
||||
|
||||
async def control_execution(run_id: str = "", command: str = "") -> dict[str, Any]:
|
||||
from .services.bioflow import skill_tools
|
||||
return skill_tools.control_execution(run_id, command)
|
||||
|
||||
|
||||
async def validate_delivery(manifest: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
from .services.bioflow import skill_tools
|
||||
return skill_tools.validate_delivery(manifest)
|
||||
|
||||
|
||||
async def publish_delivery(manifest: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
from .services.bioflow import skill_tools
|
||||
return skill_tools.publish_delivery(manifest)
|
||||
|
||||
|
||||
SKILL_CONTROL_TOOLS = [
|
||||
get_task_context, inspect_asset, validate_plan,
|
||||
check_readiness, get_run_evidence, propose_change,
|
||||
apply_change, execute_nodes, control_execution,
|
||||
validate_delivery, publish_delivery,
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 构建三个 MCPServer(基因家族 / 转录组 / 扩展)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def build_mcp_servers() -> tuple[Any, Any, Any]:
|
||||
from mcp.server.mcpserver import MCPServer
|
||||
|
||||
gf = MCPServer(
|
||||
"bioflow-gene-family",
|
||||
title="PlantAI 基因家族分析",
|
||||
description="基因家族分析能力单元:路径规划/建任务/确认参数/查状态/文件获取与预处理。",
|
||||
)
|
||||
rn = MCPServer(
|
||||
"bioflow-rnaseq",
|
||||
title="PlantAI 转录组分析",
|
||||
description="转录组分析能力单元:路径规划/样本配置生成/建任务/确认参数/查状态/文件获取与预处理。",
|
||||
)
|
||||
ext = MCPServer(
|
||||
"bioflow-extension",
|
||||
title="PlantAI 用户自定义扩展分析",
|
||||
description="用户自定义扩展分析能力单元:创建/转换/试运行/校验/确认/发布/运行。",
|
||||
)
|
||||
|
||||
gf.tool(name="plan_gene_family", description="根据用户自然语言需求规划基因家族分析候选路径(仅基因家族链阶段)。返回 2-3 条候选路径供用户确认。")(plan_gene_family)
|
||||
gf.tool(name="create_gene_family_job", description="按确认的阶段链创建基因家族任务。stages 只能是基因家族链阶段(跨链拒绝)。")(create_gene_family_job)
|
||||
gf.tool(name="ask_gene_family_question", description="回答基因家族流程/参数/生物学意义类问题(流程知识库检索,带引用)。")(ask_gene_family_question)
|
||||
|
||||
rn.tool(name="plan_rnaseq", description="根据用户自然语言需求规划转录组分析候选路径(仅转录组链阶段)。返回候选路径供用户确认。")(plan_rnaseq)
|
||||
rn.tool(name="create_rnaseq_job", description="按确认的阶段链创建转录组任务。stages 只能是转录组链阶段(跨链拒绝)。")(create_rnaseq_job)
|
||||
rn.tool(name="generate_samples_yaml", description="把用户的自然语言样本描述解析为 samples.yaml 写入任务 input。")(generate_samples_yaml)
|
||||
rn.tool(name="ask_rnaseq_question", description="回答转录组流程/参数/实验设计类问题(流程知识库检索,带引用)。")(ask_rnaseq_question)
|
||||
|
||||
ext.tool(name="create_extension", description="创建扩展草稿(draft):上传代码文件清单(每项 {path, content 文本}) + 运行命令 + 显示名。")(create_extension)
|
||||
ext.tool(name="extension_status", description="查询扩展详情与状态(draft/build/trial/validated/confirming/published)。")(extension_status)
|
||||
ext.tool(name="build_extension", description="AI 转换扩展(原始命令→command_template+参数+依赖);LLM 失败自动回退规则版。")(build_extension)
|
||||
ext.tool(name="trial_run_extension", description="试运行扩展(docker 执行;未接入执行器时返回 EXT_RUN)。")(trial_run_extension)
|
||||
ext.tool(name="validate_extension", description="校验扩展(金标/试运行校验)。仅 built/trial_passed 可校验。")(validate_extension)
|
||||
ext.tool(name="confirm_extension", description="确认扩展参数/触发词(用户可编辑 nl_triggers);仅 validated/trial_passed 可确认。")(confirm_extension)
|
||||
ext.tool(name="publish_extension", description="发布扩展(仅 published 状态可被自然语言触发运行)。")(publish_extension)
|
||||
ext.tool(name="run_extension", description="正式运行已发布扩展(同步返回 run_result)。")(run_extension)
|
||||
ext.tool(name="repair_extension", description="AI 修复扩展(试运行失败上下文→LLM 修复;LLM 失败回退规则版)。")(repair_extension)
|
||||
|
||||
shared = [
|
||||
job_status, confirm_stage, register_job_files, rerun_stage, skip_stage,
|
||||
pause_job, resume_job, list_archives, reuse_archive, archive_job,
|
||||
fetch_source_file, inspect_file, transform_file,
|
||||
]
|
||||
for mcp in (gf, rn):
|
||||
for fn in shared:
|
||||
mcp.tool()(fn)
|
||||
# G12:11 个分析控制工具(Skill 编排,异步薄封装委托 skill_tools)
|
||||
for fn in SKILL_CONTROL_TOOLS:
|
||||
mcp.tool()(fn)
|
||||
return gf, rn, ext
|
||||
|
||||
|
||||
gene_family_mcp, rnaseq_mcp, extension_mcp = build_mcp_servers()
|
||||
@@ -0,0 +1,11 @@
|
||||
你是生信数据文件分析助手。用户提供一份数据目录的文件特征清单(含路径、大小、扩展名、magic 类型、压缩包内嵌文件名),请你推断哪些文件是 RNA-seq 测序数据(FASTQ),并按配对样本输出。
|
||||
|
||||
只输出 JSON,格式:
|
||||
{"samples":[{"name":"样本名","r1":"R1文件路径或内嵌名","r2":"R2文件路径或内嵌名","run":"测序run标识","group_role":"control或compare或留空","condition":"实验条件或留空","genotype":"基因型或留空","replicate":"重复编号或留空"}]}
|
||||
|
||||
规则:
|
||||
1. r1/r2 必须是清单中真实出现的路径或内嵌文件名,禁止臆造不存在的路径。
|
||||
2. 成对判断依据:文件名(或内嵌名)中含 R1/f1/1/forward 与 R2/f2/2/reverse 的配对;若无明确配对线索,单端文件只填 r1。
|
||||
3. 无法确定任何样本时,输出 {"samples":[]}。
|
||||
4. 不确定的字段留空字符串,不要编造。
|
||||
5. 只输出 JSON,不要任何解释或 markdown 围栏。
|
||||
@@ -31,10 +31,16 @@ from ..schemas import ArchiveCreateRequest, ArchiveRunRequest
|
||||
from ..services.bioflow import store
|
||||
from ..services.bioflow.archive import (
|
||||
archive_from_job,
|
||||
effective_stage_params,
|
||||
normalize_archive_name,
|
||||
remove_archive_dir,
|
||||
)
|
||||
from ..services.bioflow.reports import extract_figure_refs, locate_report, report_payload
|
||||
from ..services.bioflow.reports import (
|
||||
extract_figure_refs,
|
||||
locate_report,
|
||||
report_payload,
|
||||
validate_report_path,
|
||||
)
|
||||
from .bioflow import job_or_404
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -65,6 +71,10 @@ async def reports_register(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
job_id = str(payload.get("job_id") or "")
|
||||
record = await load_record(job_id) if job_id else None
|
||||
report_path = payload.get("md_path") or payload.get("report_path") or ""
|
||||
# G3:拒绝空路径/空文件(不静默登记空报告)
|
||||
ok, err = validate_report_path(report_path)
|
||||
if not ok:
|
||||
return _err("EMPTY_REPORT", err, 422)
|
||||
figures: list[str] = list(payload.get("figures") or [])
|
||||
if record is not None and not figures:
|
||||
located = locate_report(record.workspace_dir)
|
||||
@@ -96,12 +106,14 @@ async def create_archive(body: ArchiveCreateRequest) -> dict[str, Any]:
|
||||
existing = await store.find_archive_by_name(ctx.db, name)
|
||||
if existing is not None:
|
||||
return _err("ARCHIVE_EXISTS", f"归档名已存在: {name}", 409)
|
||||
# G6:保存各阶段实际生效参数(非创建时 record.params)
|
||||
effective_params = effective_stage_params(record.stage_chain, record._stage_states)
|
||||
report_path = locate_report(record.workspace_dir)
|
||||
result = await archive_from_job(
|
||||
{
|
||||
"id": record.id, "session_id": record.session_id,
|
||||
"title": record.title, "stage_chain": record.stage_chain,
|
||||
"params": record.params,
|
||||
"params": effective_params,
|
||||
},
|
||||
record.job_dir,
|
||||
get_settings().archives_dir_abs,
|
||||
@@ -116,7 +128,7 @@ async def create_archive(body: ArchiveCreateRequest) -> dict[str, Any]:
|
||||
"name": name,
|
||||
"description": body.description,
|
||||
"stage_chain": record.stage_chain,
|
||||
"params": record.params,
|
||||
"params": effective_params,
|
||||
"report_artifact": result.report_artifact,
|
||||
"disk_path": result.disk_path,
|
||||
"size_bytes": result.size_bytes,
|
||||
@@ -166,7 +178,10 @@ async def delete_archive(archive_id: str) -> JSONResponse:
|
||||
|
||||
@router.post("/archives/{archive_id}/run", status_code=201)
|
||||
async def archive_run(archive_id: str, body: ArchiveRunRequest) -> dict[str, Any]:
|
||||
"""复用发起(FR-BIO-023):装载阶段链+参数 → 新建任务(表单预填存档值)。"""
|
||||
"""复用发起(FR-BIO-023):装载阶段链+参数 → 新建任务(表单预填存档值)。
|
||||
|
||||
params 为存档时落库的各阶段实际生效参数(G6:非创建时 record.params)。
|
||||
"""
|
||||
ctx: Context = context_mod.context
|
||||
archive = await store.get_archive(ctx.db, archive_id)
|
||||
if archive is None:
|
||||
|
||||
@@ -156,6 +156,48 @@ async def chat_stream(body: ChatStreamRequest) -> StreamingResponse:
|
||||
last_seq = body.last_seq if body.last_seq is not None else -1
|
||||
|
||||
async def gen():
|
||||
# 0. 扩展意图两层判定(契约 §5)——必须在 route_intent 之前:
|
||||
# 「创建一个扩展用于…」这类消息会被生信词库判为 qa,永远到不了
|
||||
# bioflow 分支;扩展判定先行,命中发 extension_* 帧后返回,
|
||||
# 不命中走原 route_intent 逻辑(漏判安全,现网零影响)。
|
||||
from ..services.bioflow import extensions as ext_mod
|
||||
from ..services.bioflow.intent import detect_extension_intent
|
||||
try:
|
||||
ext_store = ext_mod.get_store()
|
||||
published = ext_store.list_extensions()
|
||||
ext_intent = detect_extension_intent(body.message, published)
|
||||
layer = str(ext_intent.get("layer") or "none")
|
||||
if layer == "create":
|
||||
yield _sse({"type": "extension_upload", "seq": -2, "data": {
|
||||
"layer": "create", "message": body.message,
|
||||
}})
|
||||
return
|
||||
if layer == "run":
|
||||
candidates = [
|
||||
{
|
||||
"ext_id": x.get("ext_id"),
|
||||
"display_name": x.get("display_name"),
|
||||
"purpose": x.get("analysis_type") or "",
|
||||
"score": x.get("score"),
|
||||
"reasons": x.get("reasons") or [],
|
||||
}
|
||||
for x in (ext_intent.get("candidates") or [])
|
||||
]
|
||||
yield _sse({"type": "extension_candidates", "seq": -2, "data": {
|
||||
"candidates": candidates,
|
||||
"high_confidence": len(candidates) == 1,
|
||||
}})
|
||||
return
|
||||
if layer == "create_ambiguous":
|
||||
yield _sse({"type": "extension_upload", "seq": -2, "data": {
|
||||
"layer": "create_ambiguous",
|
||||
"message": "没有匹配的已发布扩展,是否创建?",
|
||||
"raw_message": body.message,
|
||||
}})
|
||||
return
|
||||
except Exception as exc: # noqa: BLE001 - 扩展链路异常不回退生信主链
|
||||
logger.warning("chat_stream: extension intent check failed: %s", exc)
|
||||
|
||||
# 0. 生信流内流程咨询(has_bioflow_context + FLOW_QA_WORD)→ 知识库引用
|
||||
intent = route_intent(
|
||||
body.message,
|
||||
@@ -178,6 +220,10 @@ async def chat_stream(body: ChatStreamRequest) -> StreamingResponse:
|
||||
return
|
||||
|
||||
# 1. 规划(候选路径;复用/LLM/降级)
|
||||
# 规划开始帧:前端据此显示「分析路径规划中…」占位(planning 状态)
|
||||
yield _sse({"type": "planning", "seq": -2, "data": {
|
||||
"message": "正在分析需求并规划分析路径…",
|
||||
}})
|
||||
try:
|
||||
plan: PlanResult = await ctx.planner.plan(body.message)
|
||||
except PlannerError as exc:
|
||||
|
||||
@@ -22,7 +22,9 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -212,23 +214,43 @@ async def upload_file(
|
||||
dest = input_dir / _safe_name(file.filename)
|
||||
written = 0
|
||||
max_bytes = int(settings.max_upload_gb * 1024 ** 3)
|
||||
# G2:先落任务临时目录,再原子 rename 发布;同名冲突显式报错,不静默覆盖。
|
||||
tmp_dir = record.job_dir / ".tmp"
|
||||
tmp_dir.mkdir(parents=True, exist_ok=True)
|
||||
tmp_path = tmp_dir / f"{uuid.uuid4().hex[:8]}_{dest.name}"
|
||||
try:
|
||||
with dest.open("wb") as fh:
|
||||
# G5:逐块流式写盘(不整块进内存)
|
||||
with tmp_path.open("wb") as fh:
|
||||
while True:
|
||||
chunk = await file.read(1024 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
written += len(chunk)
|
||||
if written > max_bytes:
|
||||
fh.close()
|
||||
dest.unlink(missing_ok=True)
|
||||
return _err("UPLOAD_TOO_LARGE",
|
||||
f"超过上传上限 {settings.max_upload_gb}GB(已收 {written} 字节)", 413)
|
||||
raise _UploadTooLarge(written)
|
||||
fh.write(chunk)
|
||||
except _UploadTooLarge as exc:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
return _err("UPLOAD_TOO_LARGE",
|
||||
f"超过上传上限 {settings.max_upload_gb}GB(已收 {exc.written} 字节)", 413)
|
||||
except Exception as exc: # noqa: BLE001 - 写盘失败不登记(TC-BIO-API-050)
|
||||
logger.exception("files: 上传写盘失败 %s", job_id)
|
||||
dest.unlink(missing_ok=True)
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
return _err("INTERNAL", f"上传写盘失败: {exc}", 500)
|
||||
# 原子发布 + 同名冲突检测(G2)
|
||||
try:
|
||||
os.link(tmp_path, dest)
|
||||
except FileExistsError:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
return _err("FILE_EXISTS", f"同名文件已存在,拒绝覆盖: {dest.name}", 409)
|
||||
except OSError:
|
||||
# os.link 不支持的文件系统 → 退化:exists 检查 + rename(同事件循环内无 await 交错)
|
||||
if dest.exists() or dest.is_symlink():
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
return _err("FILE_EXISTS", f"同名文件已存在,拒绝覆盖: {dest.name}", 409)
|
||||
os.rename(tmp_path, dest)
|
||||
else:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
file_id = f"{job_id}:{dest.name}"
|
||||
# BUG-P7-01:文件落盘后自动续跑(await_inputs 任务重推;非等待态 no-op)
|
||||
await ctx.runner.retry_after_files(record)
|
||||
@@ -238,23 +260,122 @@ async def upload_file(
|
||||
}}
|
||||
|
||||
|
||||
class _UploadTooLarge(Exception):
|
||||
"""上传超限信号(携带已收字节数)。"""
|
||||
|
||||
def __init__(self, written: int) -> None:
|
||||
super().__init__("upload too large")
|
||||
self.written = written
|
||||
|
||||
|
||||
async def _register_path(record: Any, server_path: str, settings: Any) -> dict[str, Any]:
|
||||
"""服务器路径登记(FR-BIO-007):限 BIOFLOW_DATA_ROOT/任务目录内。"""
|
||||
"""服务器路径登记(FR-BIO-007):单文件或目录展开,任意数据位置(敏感路径黑名单拦截)。
|
||||
|
||||
目录:递归收集目录内可登记文件(扩展名白名单),登记为多条记录。
|
||||
白名单根校验已按用户要求移除,改为敏感路径黑名单(见 settings)。
|
||||
"""
|
||||
path = Path(server_path).expanduser().resolve()
|
||||
if not path.exists():
|
||||
return _err("PATH_NOT_EXIST", f"路径不存在: {server_path}", 400)
|
||||
if not path.is_file():
|
||||
return _err("NOT_A_FILE", f"非文件: {server_path}", 400)
|
||||
allowed_roots = [Path(settings.data_root_abs).resolve(), record.job_dir.resolve()]
|
||||
if not any(path.is_relative_to(root) for root in allowed_roots):
|
||||
return _err("PATH_NOT_ALLOWED", f"路径不在白名单根内: {server_path}", 403)
|
||||
if not path.is_file() and not path.is_dir():
|
||||
return _err("NOT_A_FILE", f"非文件或目录: {server_path}", 400)
|
||||
blocked = _blocked_reason(path, settings)
|
||||
if blocked:
|
||||
return _err("PATH_NOT_ALLOWED", f"路径不允许: {server_path}({blocked})", 403)
|
||||
|
||||
files: list[Path] = []
|
||||
if path.is_file():
|
||||
files = [path]
|
||||
else:
|
||||
ext_set = settings.upload_ext_set
|
||||
for p in sorted(path.rglob("*")):
|
||||
if p.is_file():
|
||||
if ext_set and p.suffix.lower() not in ext_set:
|
||||
continue
|
||||
if _blocked_reason(p, settings):
|
||||
continue
|
||||
files.append(p)
|
||||
if not files:
|
||||
return _err("VALIDATION",
|
||||
f"目录内无可登记文件(允许扩展名 {settings.upload_ext_allow[:80]}…)", 400)
|
||||
|
||||
# 登记文件 symlink 进任务 input(CLI 扫 /job/input 发现数据)+ 持久化挂载清单
|
||||
input_dir = record.input_dir
|
||||
input_dir.mkdir(parents=True, exist_ok=True)
|
||||
for p in files:
|
||||
link = input_dir / p.name
|
||||
try:
|
||||
if not link.exists():
|
||||
link.symlink_to(p)
|
||||
except OSError as exc: # noqa: BLE001
|
||||
logger.warning("登记 symlink 失败(不阻断): %s -> %s (%s)", p, link, exc)
|
||||
# rnaseq 系阶段需要 samples.yaml(CLI -c 必填):描述文件优先解析 + 实况分析生成,
|
||||
# 返回引用 fastq 的父目录(并入挂载清单,跨重启持久化)
|
||||
try:
|
||||
from ..services.bioflow.runner import ensure_samples_yaml
|
||||
source_roots = []
|
||||
if path.is_dir():
|
||||
source_roots.append(path)
|
||||
else:
|
||||
source_roots.append(path.parent)
|
||||
extra_mounts = ensure_samples_yaml(input_dir, source_roots=source_roots)
|
||||
if not extra_mounts:
|
||||
# 规则无法归纳(无描述文件、实况配对 0 样本) → LLM 兜底
|
||||
try:
|
||||
from ..services.bioflow.runner import llm_sample_fallback
|
||||
if await llm_sample_fallback(input_dir, source_roots):
|
||||
logger.info("samples.yaml LLM 兜底生成: %s", input_dir)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("samples.yaml LLM 兜底失败(不阻断登记): %s", exc)
|
||||
if extra_mounts:
|
||||
logger.info("samples.yaml 分析完成(挂载父目录 %d 个): %s", len(extra_mounts), input_dir)
|
||||
record.registered_server_paths = list(
|
||||
dict.fromkeys(
|
||||
(getattr(record, "registered_server_paths", None) or [])
|
||||
+ [str(Path(m)) for m in extra_mounts]
|
||||
)
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("samples.yaml 分析失败(不阻断登记): %s", exc)
|
||||
registered = [str(p) for p in files]
|
||||
try:
|
||||
_persist_registered_paths(record, registered)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("登记路径持久化失败(不阻断登记): %s", exc)
|
||||
|
||||
# BUG-P7-01:路径登记后自动续跑(await_inputs 任务重推;非等待态 no-op)
|
||||
ctx: Context = context_mod.context
|
||||
await ctx.runner.retry_after_files(record)
|
||||
return {"ok": True, "data": {
|
||||
"file_id": f"{record.id}:{path.name}", "name": path.name,
|
||||
"size": path.stat().st_size, "path": str(path), "source": "server_path",
|
||||
}}
|
||||
items = [
|
||||
{"file_id": f"{record.id}:{p.name}", "name": p.name,
|
||||
"size": p.stat().st_size, "path": str(p), "source": "server_path"}
|
||||
for p in files
|
||||
]
|
||||
return {"ok": True, "data": {"items": items, "total": len(items)}}
|
||||
|
||||
|
||||
def _persist_registered_paths(record: Any, registered: list[str]) -> None:
|
||||
"""登记路径落盘:任务目录 server_registered.json(跨重启可恢复挂载清单)。"""
|
||||
record.registered_server_paths = list(
|
||||
dict.fromkeys((getattr(record, "registered_server_paths", None) or []) + registered)
|
||||
)
|
||||
payload = {"registered_server_paths": record.registered_server_paths}
|
||||
(record.job_dir / "server_registered.json").write_text(
|
||||
json.dumps(payload, ensure_ascii=False), encoding="utf-8"
|
||||
)
|
||||
|
||||
|
||||
def _blocked_reason(path: Path, settings: Any) -> str | None:
|
||||
"""敏感路径黑名单:命中返回原因字符串,未命中返回 None。"""
|
||||
s = str(path)
|
||||
for prefix in settings.sensitive_path_prefixes:
|
||||
if s == prefix or s.startswith(prefix + "/"):
|
||||
return f"敏感路径前缀 {prefix}"
|
||||
name = path.name
|
||||
for pat in settings.sensitive_filename_patterns:
|
||||
if re.match(pat, name):
|
||||
return f"敏感文件名 {name}"
|
||||
return None
|
||||
|
||||
|
||||
@router.get("/jobs/{job_id}/files")
|
||||
@@ -395,6 +516,46 @@ async def stage_manifest(job_id: str, stage: str) -> dict[str, Any]:
|
||||
}}
|
||||
|
||||
|
||||
@router.get("/jobs/{job_id}/stages/{stage}/logs")
|
||||
async def stage_logs(
|
||||
job_id: str,
|
||||
stage: str,
|
||||
offset: int = Query(default=0, ge=0),
|
||||
limit: int = Query(default=500, ge=1, le=2000),
|
||||
) -> dict[str, Any]:
|
||||
"""阶段日志分页读取(offset=已读行数;执行中可读,前端轮询滚动 log)。
|
||||
|
||||
- kernel 流式 fsync 追加落盘,本端点按行切片返回增量;
|
||||
- offset 超尾 → 空 lines + eof:true;
|
||||
- 返回 total_bytes 供 eof 判定(文件总行数随执行增长)。
|
||||
"""
|
||||
record = await job_or_404(job_id)
|
||||
_stage_exists(record, stage)
|
||||
log_path = record.logs_dir / f"{stage}.log"
|
||||
if not log_path.exists():
|
||||
return {"ok": True, "data": {
|
||||
"job_id": job_id, "stage": stage,
|
||||
"offset": offset, "limit": limit, "lines": [],
|
||||
"total_bytes": 0, "eof": True, "log_tail": "",
|
||||
}}
|
||||
try:
|
||||
text = log_path.read_text(encoding="utf-8", errors="replace")
|
||||
except OSError:
|
||||
return _err("INTERNAL", "日志读取失败", 500)
|
||||
lines = text.splitlines()
|
||||
total = len(lines)
|
||||
page = lines[offset:offset + limit]
|
||||
eof = offset + len(page) >= total
|
||||
return {"ok": True, "data": {
|
||||
"job_id": job_id, "stage": stage,
|
||||
"offset": offset, "limit": limit,
|
||||
"lines": page,
|
||||
"total_bytes": log_path.stat().st_size if log_path.exists() else 0,
|
||||
"eof": eof,
|
||||
"log_tail": "\n".join(lines[-get_settings().log_tail_lines:]) if total else "",
|
||||
}}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# events SSE(FR-BIO-014 断线重连回放)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -0,0 +1,484 @@
|
||||
"""app/routers/extension.py —— 用户自定义扩展分析端点(契约 §6 全集)。
|
||||
|
||||
- POST /api/bioflow/extensions 上传 zip(multipart)+ display_name
|
||||
+ command + runtime_image? + golden?
|
||||
- GET /api/bioflow/extensions 列表(published 优先)
|
||||
- GET /api/bioflow/extensions/{id} 详情
|
||||
- DELETE /api/bioflow/extensions/{id} 删除
|
||||
- POST /api/bioflow/extensions/{id}/build AI 转换(LLM 失败回退规则版)
|
||||
- POST /api/bioflow/extensions/{id}/trial-run 试运行(docker 执行占位 Executor)
|
||||
- POST /api/bioflow/extensions/{id}/validate 金标/试运行校验
|
||||
- POST /api/bioflow/extensions/{id}/confirm 用户确认/参数调整(nl_triggers 可编辑)
|
||||
- POST /api/bioflow/extensions/{id}/repair LLM 修复(上限 20 轮可配)
|
||||
- POST /api/bioflow/extensions/{id}/publish 发布(仅 published 可被 NL 触发)
|
||||
- POST /api/bioflow/extensions/{id}/run 正式运行(同步返回 run_result;
|
||||
SSE 由 Dev-B 执行对接)
|
||||
- GET /api/bioflow/extensions/{id}/runs 运行历史
|
||||
- GET /api/bioflow/extensions/{id}/runs/{run_id}/events SSE(占位:先返回
|
||||
已完成 run 的快照帧;live 流由
|
||||
Dev-B 执行链对接后补齐)
|
||||
|
||||
docker 执行:trial/run 接 extensions.ExtensionExecutor 接口(Dev-B 的
|
||||
extension_exec.py 实现;inject_executor 注入)。未接入时返回 503(不是 500)。
|
||||
|
||||
服务层全部经 services/bioflow/extensions.py 状态机;LLM 转换不可用自动降级
|
||||
规则版(extension_llm.py 保证),端点绝不因 LLM 挂掉而 500。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import tarfile
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
|
||||
from ..schemas import ExtensionConfirmRequest, ExtensionRunRequest
|
||||
from ..services.bioflow import extensions as ext_mod
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/api/bioflow/extensions", tags=["bioflow-extensions"])
|
||||
|
||||
_MAX_ZIP_ENTRIES = 500 # zip 条目数上限(防炸弹)
|
||||
_MAX_ZIP_TOTAL = 200 * 1024 * 1024 # 解压总量上限 200MB
|
||||
|
||||
|
||||
def _err(code: str, message: str, status: int) -> JSONResponse:
|
||||
return JSONResponse(status_code=status, content={
|
||||
"ok": False,
|
||||
"error": {"code": code, "message": message, "degraded": False},
|
||||
})
|
||||
|
||||
|
||||
def _store() -> ext_mod.ExtensionStore:
|
||||
return ext_mod.get_store()
|
||||
|
||||
|
||||
def _ext_or_404(ext_id: str) -> dict[str, Any]:
|
||||
try:
|
||||
return _store().load_extension(ext_id)
|
||||
except ext_mod.ExtensionNotFoundError:
|
||||
raise HTTPException(status_code=404, detail={
|
||||
"ok": False,
|
||||
"error": {"code": "NOT_FOUND", "message": "扩展不存在", "degraded": False},
|
||||
}) from None
|
||||
|
||||
|
||||
def _state_error(exc: ext_mod.ExtensionError, status: int = 400) -> JSONResponse:
|
||||
code = "EXT_STATUS" if isinstance(exc, ext_mod.ExtensionStatusError) else (
|
||||
"EXT_BUILD" if isinstance(exc, ext_mod.ExtensionBuildError) else (
|
||||
"EXT_RUN" if isinstance(exc, ext_mod.ExtensionRunError) else "EXT_ERROR"))
|
||||
return _err(code, str(exc), status)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 上传/列表/详情/删除
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.post("")
|
||||
async def create_extension(
|
||||
request: Request,
|
||||
display_name: str | None = Form(default=None),
|
||||
command: str | None = Form(default=None),
|
||||
runtime_image: str = Form(default=""),
|
||||
golden: str | None = Form(default=None),
|
||||
code_zip: UploadFile | None = File(default=None),
|
||||
test_file: UploadFile | None = File(default=None),
|
||||
dir_files: list[UploadFile] | None = File(default=None, alias="code_files"),
|
||||
relative_paths: str | None = Form(default=None),
|
||||
) -> JSONResponse:
|
||||
"""上传 zip(代码包)+ 可选测试数据/金标 → 创建 draft 扩展(契约 §6)。
|
||||
|
||||
multipart:display_name/command 为表单字段;JSON body:同名字段在 JSON 内
|
||||
(Form 参数可选化以兼容两种形态——fastapi 在 multipart 缺失表单时不再 422)。
|
||||
golden 为 JSON 字符串({test_input, expected_outputs}),测试数据可随
|
||||
test_file 上传;code_files 为 base64 清单(Dev-C 前端上传卡可裁剪)。
|
||||
"""
|
||||
code_files: list[tuple[str, bytes]] = []
|
||||
data_files: list[tuple[str, bytes]] = []
|
||||
golden_dict: dict[str, Any] | None = None
|
||||
|
||||
content_type = request.headers.get("content-type", "")
|
||||
try:
|
||||
if "multipart/form-data" in content_type:
|
||||
if golden:
|
||||
try:
|
||||
golden_dict = json.loads(golden)
|
||||
except ValueError:
|
||||
return _err("VALIDATION", "golden 非合法 JSON", 422)
|
||||
if test_file is not None and test_file.filename:
|
||||
data_files.append((test_file.filename, await test_file.read()))
|
||||
rel_paths: list[str] = []
|
||||
if relative_paths:
|
||||
rel_paths = [str(x).strip() for x in relative_paths.split("\n") if str(x).strip()]
|
||||
if dir_files is not None:
|
||||
# 目录直传:多文件 + 相对路径(filename 优先,relative_paths 兜底)
|
||||
if rel_paths and len(rel_paths) == len(dir_files):
|
||||
pairs = list(zip(rel_paths, dir_files))
|
||||
else:
|
||||
pairs = [(str(f.filename or ""), f) for f in dir_files]
|
||||
raw_items: list[tuple[str, bytes]] = []
|
||||
for rel, f in pairs:
|
||||
blob = await f.read()
|
||||
if blob:
|
||||
raw_items.append((rel or str(f.filename or ""), blob))
|
||||
if raw_items:
|
||||
code_files, data_files = _classify_items(
|
||||
raw_items, command or "", code_files, data_files)
|
||||
elif code_zip is not None and code_zip.filename:
|
||||
items = _read_archive(await code_zip.read())
|
||||
code_files, data_files = _classify_items(
|
||||
items, command or "", code_files, data_files)
|
||||
else:
|
||||
# JSON 形态(前端可先不传 zip;代码文件经 base64 清单)
|
||||
try:
|
||||
body = await request.json()
|
||||
except Exception: # noqa: BLE001
|
||||
return _err("VALIDATION", "请求体须为 multipart 或 JSON", 400)
|
||||
display_name = str(body.get("display_name") or display_name or "")
|
||||
command = str(body.get("command") or command or "")
|
||||
runtime_image = str(body.get("runtime_image") or runtime_image)
|
||||
golden_dict = body.get("golden")
|
||||
for item in body.get("code_files") or []:
|
||||
if not isinstance(item, dict) or not item.get("path"):
|
||||
continue
|
||||
try:
|
||||
blob = base64.b64decode(str(item.get("data") or ""))
|
||||
except Exception: # noqa: BLE001
|
||||
continue
|
||||
code_files.append((str(item["path"]), blob))
|
||||
for item in body.get("data_files") or []:
|
||||
if not isinstance(item, dict) or not item.get("path"):
|
||||
continue
|
||||
try:
|
||||
blob = base64.b64decode(str(item.get("data") or ""))
|
||||
except Exception: # noqa: BLE001
|
||||
continue
|
||||
data_files.append((str(item["path"]), blob))
|
||||
except ext_mod.ExtensionError as exc:
|
||||
return _state_error(exc, 422)
|
||||
|
||||
if not code_files:
|
||||
return _err("VALIDATION", "未识别到任何代码文件(请检查上传内容或运行命令)", 422)
|
||||
|
||||
try:
|
||||
ext = _store().create_extension(
|
||||
display_name=str(display_name or ""),
|
||||
command=str(command or ""),
|
||||
runtime_image=str(runtime_image),
|
||||
code_files=code_files,
|
||||
data_files=data_files,
|
||||
golden=golden_dict if isinstance(golden_dict, dict) else None,
|
||||
)
|
||||
except ext_mod.ExtensionError as exc:
|
||||
return _state_error(exc, 422)
|
||||
return JSONResponse(status_code=201, content={"ok": True, "data": ext})
|
||||
|
||||
|
||||
def _safe_rel_path(raw: str) -> str:
|
||||
"""相对路径净化:拒绝 .. 与绝对路径,空段丢弃,返回 posix 形式;不可净化返回 ''。"""
|
||||
parts = str(raw or "").replace("\\", "/").split("/")
|
||||
clean: list[str] = []
|
||||
for seg in parts:
|
||||
seg = seg.strip()
|
||||
if not seg or seg == ".":
|
||||
continue
|
||||
if seg == "..":
|
||||
# 路径穿越段丢弃(与 safe_filename 语义一致:../../evil.py → evil.py)
|
||||
continue
|
||||
clean.append(seg)
|
||||
return "/".join(clean)
|
||||
|
||||
|
||||
def _read_archive(blob: bytes) -> list[tuple[str, bytes]]:
|
||||
"""归档解压(zip/tar.gz/tar,魔数分派;护栏:条目数/总量/路径净化,保留相对路径)。
|
||||
|
||||
失败抛 ExtensionError,语义由调用方处理。"""
|
||||
items: list[tuple[str, bytes]] = []
|
||||
total = 0
|
||||
|
||||
def _append(name: str, data: bytes) -> None:
|
||||
nonlocal total
|
||||
total += len(data)
|
||||
if total > _MAX_ZIP_TOTAL:
|
||||
raise ext_mod.ExtensionError(f"压缩包解压总量超限(> {_MAX_ZIP_TOTAL // (1024 * 1024)}MB)")
|
||||
rel = _safe_rel_path(name)
|
||||
if rel:
|
||||
items.append((rel, data))
|
||||
|
||||
if blob[:4] == b"PK\x03\x04":
|
||||
try:
|
||||
zf = zipfile.ZipFile(io.BytesIO(blob))
|
||||
except zipfile.BadZipFile:
|
||||
raise ext_mod.ExtensionError("zip 文件损坏") from None
|
||||
with zf:
|
||||
infos = zf.infolist()
|
||||
if len(infos) > _MAX_ZIP_ENTRIES:
|
||||
raise ext_mod.ExtensionError(f"压缩包条目过多(> {_MAX_ZIP_ENTRIES})")
|
||||
for info in infos:
|
||||
if info.is_dir():
|
||||
continue
|
||||
_append(info.filename, zf.read(info))
|
||||
return items
|
||||
|
||||
is_gz = blob[:2] == b"\x1f\x8b"
|
||||
is_ustar = len(blob) > 262 and blob[257:262] == b"ustar"
|
||||
if is_gz or is_ustar:
|
||||
mode = "r:gz" if is_gz else "r:"
|
||||
try:
|
||||
tf = tarfile.open(fileobj=io.BytesIO(blob), mode=mode)
|
||||
except tarfile.TarError:
|
||||
raise ext_mod.ExtensionError("tar 文件损坏") from None
|
||||
with tf:
|
||||
members = tf.getmembers()
|
||||
if len(members) > _MAX_ZIP_ENTRIES:
|
||||
raise ext_mod.ExtensionError(f"压缩包条目过多(> {_MAX_ZIP_ENTRIES})")
|
||||
for m in members:
|
||||
if not m.isfile():
|
||||
continue
|
||||
f = tf.extractfile(m)
|
||||
if f is None:
|
||||
continue
|
||||
_append(m.name, f.read())
|
||||
return items
|
||||
|
||||
raise ext_mod.ExtensionError("不支持的压缩包格式(仅支持 zip/tar/tar.gz)")
|
||||
|
||||
|
||||
_CODE_DIR_PREFIXES = ("raw_code/", "code/", "scripts/")
|
||||
_DATA_DIR_PREFIXES = ("raw_test_data/", "test_data/", "data/")
|
||||
_CODE_EXT = re.compile(r"\.(py|r|sh|bash|pl|js|ts|java|c|cc|cpp|go|rs|jl|rb|m|smk|nf|ipynb)$", re.I)
|
||||
_DATA_EXT = re.compile(r"\.(fa|fasta|fna|faa|pep|cds|gff|gff3|gtf|fastq|fq|gz|bam|sam|vcf|csv|tsv|txt|xlsx?|nwk|nw|maf)$", re.I)
|
||||
|
||||
|
||||
def _classify_items(items, command, code_files, data_files):
|
||||
"""结构自适应分类(§3.5,优先级递减;永不因结构不符拒绝上传)。
|
||||
|
||||
1. 目录前缀(raw_code|code|scripts → code;raw_test_data|test_data|data → data)
|
||||
2. 命令解析:入口文件恒归 code;命令参数出现的文件名归 data
|
||||
3. 扩展名模式;未知扩展名归 code
|
||||
返回 (code_files, data_files)。"""
|
||||
import re as _re
|
||||
|
||||
cmd_tokens: set[str] = set()
|
||||
if command:
|
||||
for tok in str(command).split():
|
||||
tok = tok.strip().strip('"\'')
|
||||
if not tok:
|
||||
continue
|
||||
base = tok.split("/")[-1]
|
||||
if "." in base:
|
||||
cmd_tokens.add(base)
|
||||
|
||||
for rel, blob in items:
|
||||
rl = rel.lower()
|
||||
target = None
|
||||
if any(rl.startswith(p) for p in _DATA_DIR_PREFIXES):
|
||||
target = "data"
|
||||
elif any(rl.startswith(p) for p in _CODE_DIR_PREFIXES):
|
||||
target = "code"
|
||||
else:
|
||||
base = rel.split("/")[-1]
|
||||
if _CODE_EXT.search(base):
|
||||
target = "code"
|
||||
elif _DATA_EXT.search(base):
|
||||
target = "data"
|
||||
else:
|
||||
target = "code" # 未知扩展名归 code
|
||||
# 剥已知前缀(仅标准前缀;其余前缀保持原路径)
|
||||
out_rel = rel
|
||||
for p in _DATA_DIR_PREFIXES + _CODE_DIR_PREFIXES:
|
||||
if rl.startswith(p):
|
||||
out_rel = rel[len(p):]
|
||||
break
|
||||
if target == "code":
|
||||
code_files.append((out_rel, blob))
|
||||
else:
|
||||
data_files.append((out_rel, blob))
|
||||
return code_files, data_files
|
||||
|
||||
|
||||
@router.get("")
|
||||
async def list_extensions() -> dict[str, Any]:
|
||||
items = _store().list_extensions()
|
||||
return {"ok": True, "data": {"items": items, "total": len(items)}}
|
||||
|
||||
|
||||
@router.get("/{ext_id}")
|
||||
async def get_extension(ext_id: str) -> dict[str, Any]:
|
||||
ext = _ext_or_404(ext_id)
|
||||
return {"ok": True, "data": ext}
|
||||
|
||||
|
||||
@router.delete("/{ext_id}")
|
||||
async def delete_extension(ext_id: str) -> JSONResponse:
|
||||
try:
|
||||
_store().delete(ext_id)
|
||||
except ext_mod.ExtensionNotFoundError:
|
||||
raise HTTPException(status_code=404, detail={
|
||||
"ok": False, "error": {"code": "NOT_FOUND", "message": "扩展不存在", "degraded": False},
|
||||
}) from None
|
||||
except ext_mod.ExtensionStatusError as exc:
|
||||
return _state_error(exc, 409)
|
||||
return JSONResponse(status_code=204, content=None)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 生命周期动作(build/trial/validate/repair/confirm/publish)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.post("/{ext_id}/build")
|
||||
async def build_extension(ext_id: str) -> JSONResponse:
|
||||
_ext_or_404(ext_id)
|
||||
try:
|
||||
ext = await _store().build_extension(ext_id)
|
||||
except ext_mod.ExtensionBuildError as exc:
|
||||
return _state_error(exc, 400)
|
||||
except ext_mod.ExtensionStatusError as exc:
|
||||
return _state_error(exc, 409)
|
||||
return JSONResponse(content={"ok": True, "data": ext})
|
||||
|
||||
|
||||
@router.post("/{ext_id}/trial-run")
|
||||
async def trial_run(ext_id: str) -> JSONResponse:
|
||||
_ext_or_404(ext_id)
|
||||
try:
|
||||
ext = await _store().trial_run(ext_id)
|
||||
except ext_mod.ExtensionStatusError as exc:
|
||||
return _state_error(exc, 409)
|
||||
except ext_mod.ExtensionRunError as exc:
|
||||
return _state_error(exc, 503)
|
||||
return JSONResponse(content={"ok": True, "data": ext})
|
||||
|
||||
|
||||
@router.post("/{ext_id}/validate")
|
||||
async def validate(ext_id: str) -> JSONResponse:
|
||||
_ext_or_404(ext_id)
|
||||
try:
|
||||
ext = _store().validate(ext_id)
|
||||
except ext_mod.ExtensionStatusError as exc:
|
||||
return _state_error(exc, 409)
|
||||
return JSONResponse(content={"ok": True, "data": ext})
|
||||
|
||||
|
||||
@router.post("/{ext_id}/repair")
|
||||
async def repair(ext_id: str) -> JSONResponse:
|
||||
_ext_or_404(ext_id)
|
||||
try:
|
||||
ext = await _store().repair(ext_id)
|
||||
except ext_mod.ExtensionBuildError as exc:
|
||||
return _state_error(exc, 400)
|
||||
except ext_mod.ExtensionStatusError as exc:
|
||||
return _state_error(exc, 409)
|
||||
return JSONResponse(content={"ok": True, "data": ext})
|
||||
|
||||
|
||||
@router.post("/{ext_id}/confirm")
|
||||
async def confirm(ext_id: str, body: ExtensionConfirmRequest) -> JSONResponse:
|
||||
_ext_or_404(ext_id)
|
||||
try:
|
||||
ext = _store().confirm(
|
||||
ext_id,
|
||||
nl_triggers=body.nl_triggers,
|
||||
parameters=body.parameters,
|
||||
runtime_image=body.runtime_image,
|
||||
)
|
||||
except ext_mod.ExtensionStatusError as exc:
|
||||
return _state_error(exc, 409)
|
||||
return JSONResponse(content={"ok": True, "data": ext})
|
||||
|
||||
|
||||
@router.post("/{ext_id}/publish")
|
||||
async def publish(ext_id: str) -> JSONResponse:
|
||||
_ext_or_404(ext_id)
|
||||
try:
|
||||
ext = _store().publish(ext_id)
|
||||
except ext_mod.ExtensionStatusError as exc:
|
||||
return _state_error(exc, 409)
|
||||
return JSONResponse(content={"ok": True, "data": ext})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 运行(run/runs/events)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.post("/{ext_id}/run")
|
||||
async def run(ext_id: str, body: ExtensionRunRequest) -> JSONResponse:
|
||||
"""正式运行:同步执行并返回 run_result(run 事件 SSE 由 Dev-B 执行对接)。"""
|
||||
_ext_or_404(ext_id)
|
||||
try:
|
||||
result = await _store().run(ext_id, params=body.params, files=body.files)
|
||||
except ext_mod.ExtensionStatusError as exc:
|
||||
return _state_error(exc, 409)
|
||||
except ext_mod.ExtensionRunError as exc:
|
||||
return _state_error(exc, 400 if "不存在" in str(exc) or "白名单" in str(exc) else 503)
|
||||
return JSONResponse(content={"ok": True, "data": result})
|
||||
|
||||
|
||||
@router.get("/{ext_id}/runs")
|
||||
async def list_runs(ext_id: str) -> dict[str, Any]:
|
||||
_ext_or_404(ext_id)
|
||||
items = _store().list_runs(ext_id)
|
||||
return {"ok": True, "data": {"items": items, "total": len(items)}}
|
||||
|
||||
|
||||
@router.get("/{ext_id}/runs/{run_id}/events")
|
||||
async def run_events(ext_id: str, run_id: str, request: Request) -> StreamingResponse:
|
||||
"""运行事件 SSE(契约 §6;首帧快照)。
|
||||
|
||||
P0 占位实现:run 为同步执行,此处返回已完成 run 的快照帧
|
||||
(run_started/run_log/run_done|run_failed)后关闭;live 进度流在 Dev-B
|
||||
把执行改造为后台任务 + 事件推送后,本端点改为订阅事件总线。
|
||||
"""
|
||||
ext = _ext_or_404(ext_id)
|
||||
result_file = _store().runs_dir(ext_id) / run_id / "run_result.json"
|
||||
log_file = _store().runs_dir(ext_id) / run_id / "run.log"
|
||||
|
||||
def _frames():
|
||||
result: dict[str, Any] = {}
|
||||
if result_file.is_file():
|
||||
try:
|
||||
result = json.loads(result_file.read_text(encoding="utf-8"))
|
||||
except (ValueError, OSError):
|
||||
result = {}
|
||||
log_text = ""
|
||||
if log_file.is_file():
|
||||
try:
|
||||
log_text = log_file.read_text(encoding="utf-8")
|
||||
except OSError:
|
||||
log_text = ""
|
||||
yield _sse_frame("extension_run_started", {
|
||||
"run_id": run_id, "ext_id": ext_id,
|
||||
"started_at": result.get("started_at") or "",
|
||||
})
|
||||
if log_text:
|
||||
yield _sse_frame("extension_run_log", {"run_id": run_id, "log": log_text})
|
||||
if result.get("status") == "done":
|
||||
yield _sse_frame("extension_run_done", {
|
||||
"run_id": run_id, "outputs": result.get("outputs") or [],
|
||||
})
|
||||
elif result:
|
||||
yield _sse_frame("extension_run_failed", {
|
||||
"run_id": run_id, "exit_code": result.get("exit_code"),
|
||||
"error": result.get("error") or "", "log_tail": log_text[-2000:],
|
||||
})
|
||||
|
||||
return StreamingResponse(
|
||||
_frames(),
|
||||
media_type="text/event-stream",
|
||||
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
||||
)
|
||||
|
||||
|
||||
def _sse_frame(event: str, data: dict[str, Any]) -> bytes:
|
||||
return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n".encode("utf-8")
|
||||
@@ -110,3 +110,26 @@ class ReportReviseRequest(BaseModel):
|
||||
target: str | None = Field(default=None, max_length=500)
|
||||
instruction: str = Field(min_length=1, max_length=2000)
|
||||
params: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
# ---- extensions(用户自定义扩展分析;契约 §2/§6)----
|
||||
|
||||
EXTENSION_STATUSES = (
|
||||
"draft", "building", "built", "trial_running", "trial_passed", "trial_failed",
|
||||
"validating", "validated", "repairing", "confirming", "published", "build_failed",
|
||||
)
|
||||
ExtensionStatus = Literal[
|
||||
"draft", "building", "built", "trial_running", "trial_passed", "trial_failed",
|
||||
"validating", "validated", "repairing", "confirming", "published", "build_failed",
|
||||
]
|
||||
|
||||
|
||||
class ExtensionRunRequest(BaseModel):
|
||||
params: dict[str, Any] = Field(default_factory=dict)
|
||||
files: list[str] = Field(default_factory=list, max_length=50)
|
||||
|
||||
|
||||
class ExtensionConfirmRequest(BaseModel):
|
||||
nl_triggers: list[str] | None = Field(default=None, max_length=50)
|
||||
parameters: list[dict[str, Any]] | None = None
|
||||
runtime_image: str | None = Field(default=None, max_length=128)
|
||||
|
||||
@@ -19,11 +19,27 @@ from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .manifests import ManifestError, artifact_paths, parse_manifest
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
INTERMEDIATE_DIR_NAMES = ("analysis",)
|
||||
|
||||
|
||||
def effective_stage_params(
|
||||
stage_chain: list[str], stage_states: dict[str, Any]
|
||||
) -> dict[str, Any]:
|
||||
"""G6:各阶段「实际生效」参数(confirm/rerun 更新后的 stage 参数)。
|
||||
|
||||
区别于创建时 record.params(创建后 confirm/rerun 只更新 _stage_states[stage].params,
|
||||
不回流 record.params)。存档与复用以此为单一事实源,避免参数失真。
|
||||
"""
|
||||
return {
|
||||
s: dict((stage_states.get(s) or {}).get("params") or {})
|
||||
for s in stage_chain
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class ArchiveResult:
|
||||
id: str
|
||||
@@ -88,6 +104,11 @@ async def archive_from_job(
|
||||
|
||||
# 报告:任务目录内报告 HTML(final 内 *.html / report.html)随 final 已复制;
|
||||
# report_artifact 为 doc_artifacts 登记 id(由 reports 模块/BFF 登记后回传)。
|
||||
|
||||
# G1:按 manifest 必需产物清单补归档(不只 workspace/final/)。
|
||||
# 相对路径以 workspace 根解析并改写为归档内路径(去 workspace 前缀)。
|
||||
_copy_manifest_artifacts(job, job_dir, dest)
|
||||
|
||||
size_bytes = _dir_size(dest)
|
||||
return ArchiveResult(
|
||||
id=archive_id,
|
||||
@@ -98,6 +119,42 @@ async def archive_from_job(
|
||||
)
|
||||
|
||||
|
||||
def _copy_manifest_artifacts(
|
||||
job: dict[str, Any], job_dir: Path, dest: Path
|
||||
) -> list[str]:
|
||||
"""G1:按 manifest 的必需产物清单复制归档产物(漏归档修复)。
|
||||
|
||||
每个阶段的 manifest artifacts(full_path/path/relpath_or_hint)逐条复制;
|
||||
相对路径以 workspace 根解析,并改写为归档内路径(去 workspace 前缀)。
|
||||
跳过不存在/失败阶段无 manifest 的产物。返回已复制的归档内相对路径列表。
|
||||
"""
|
||||
stage_chain = job.get("stage_chain") or []
|
||||
workspace = Path(job_dir) / "workspace"
|
||||
copied: list[str] = []
|
||||
for stage in stage_chain:
|
||||
try:
|
||||
manifest = parse_manifest(workspace, stage)
|
||||
except ManifestError:
|
||||
continue
|
||||
for raw in artifact_paths(manifest.artifacts):
|
||||
p = Path(raw)
|
||||
if not p.is_absolute():
|
||||
p = workspace / p
|
||||
if not p.is_file():
|
||||
continue
|
||||
# 相对路径改写为归档内路径(去 workspace 前缀)
|
||||
try:
|
||||
rel = p.resolve().relative_to(workspace.resolve())
|
||||
except ValueError:
|
||||
rel = Path(p.name)
|
||||
target = dest / rel
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
if not target.exists():
|
||||
shutil.copy2(p, target)
|
||||
copied.append(str(rel))
|
||||
return copied
|
||||
|
||||
|
||||
def _copytree_into(src: Path, dest: Path) -> None:
|
||||
"""复制目录内容到 dest(dest 可能已存在,逐项合并)。"""
|
||||
dest.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
@@ -343,12 +343,16 @@ _RNASEQ_ALIGN_PARAMS = [
|
||||
"help": "样本列表(逗号分隔)"},
|
||||
{"name": "ref_genome", "type": "string", "default": "", "source": "user",
|
||||
"help": "参考基因组路径/索引名"},
|
||||
{"name": "hisat2_index", "type": "string", "default": "", "source": "user",
|
||||
"help": "HISAT2 索引前缀(容器内可访问路径,如 /bioflow-data/hmm/hisat2/xxx)"},
|
||||
{"name": "cpu", "type": "choice", "choices": ["4", "8", "16", "32"], "default": "8",
|
||||
"source": "runtime", "help": "容器 CPU 数"},
|
||||
]
|
||||
_RNASEQ_QUANT_PARAMS = [
|
||||
{"name": "samples", "type": "string", "default": "", "source": "user",
|
||||
"help": "样本列表(逗号分隔)"},
|
||||
{"name": "gtf", "type": "string", "default": "", "source": "user",
|
||||
"help": "参考基因组 GTF 注释(容器内可访问路径)"},
|
||||
{"name": "cpu", "type": "choice", "choices": ["4", "8", "16", "32"], "default": "6",
|
||||
"source": "runtime", "help": "容器 CPU 数"},
|
||||
]
|
||||
@@ -361,6 +365,8 @@ _RNASEQ_SAMPLE_QC_PARAMS = [
|
||||
_RNASEQ_DE_PARAMS = [
|
||||
{"name": "design", "type": "string", "default": "", "source": "user",
|
||||
"help": "实验设计矩阵文件(对比组声明)"},
|
||||
{"name": "prefix", "type": "string", "default": "", "source": "contract",
|
||||
"help": "分析前缀"},
|
||||
{"name": "cpu", "type": "choice", "choices": ["4", "8", "16", "32"], "default": "4",
|
||||
"source": "runtime", "help": "容器 CPU 数"},
|
||||
]
|
||||
@@ -373,6 +379,8 @@ _RNASEQ_WGCNA_PARAMS = [
|
||||
_RNASEQ_ENRICH_PARAMS = [
|
||||
{"name": "gene_list", "type": "string", "default": "", "source": "contract",
|
||||
"help": "差异基因列表(上游 DE 结果)"},
|
||||
{"name": "prefix", "type": "string", "default": "", "source": "contract",
|
||||
"help": "分析前缀"},
|
||||
{"name": "cpu", "type": "choice", "choices": ["4", "8", "16", "32"], "default": "4",
|
||||
"source": "runtime", "help": "容器 CPU 数"},
|
||||
]
|
||||
@@ -387,6 +395,46 @@ _RNASEQ_REPORT_PARAMS = [
|
||||
"help": "分析前缀"},
|
||||
]
|
||||
|
||||
# wgs/pop/gwas 参数 schema(契约 param: 键 → 确认表单字段)
|
||||
_CPU4 = {"name": "cpu", "type": "choice", "choices": ["4", "8", "16", "32"],
|
||||
"default": "4", "source": "runtime", "help": "容器 CPU 数"}
|
||||
|
||||
|
||||
def _u(name: str, help_: str) -> dict[str, str]:
|
||||
return {"name": name, "type": "string", "default": "", "source": "user", "help": help_}
|
||||
|
||||
|
||||
_WGS_COMMON = [_u("prefix", "分析前缀"), _CPU4]
|
||||
_WGS_QC_PARAMS = [*_WGS_COMMON, _u("wgs_samples", "WGS 样本清单(逗号分隔或路径)")]
|
||||
_WGS_INDEX_PARAMS = [*_WGS_COMMON, _u("reference_fasta", "参考基因组 FASTA 路径(容器内可访问)")]
|
||||
_WGS_ALIGN_PARAMS = [*_WGS_COMMON]
|
||||
_WGS_MARKDUP_PARAMS = [*_WGS_COMMON,
|
||||
{"name": "max_jobs", "type": "string", "default": "8", "source": "user",
|
||||
"help": "并行任务上限"},
|
||||
{"name": "remove_duplicates", "type": "choice", "choices": ["true"],
|
||||
"default": "true", "source": "user", "help": "去重复(固定 true)"}]
|
||||
_WGS_BAM_QC_PARAMS = [*_WGS_COMMON,
|
||||
{"name": "threads", "type": "string", "default": "8", "source": "user",
|
||||
"help": "线程数"},
|
||||
{"name": "max_jobs", "type": "string", "default": "8", "source": "user",
|
||||
"help": "并行任务上限"}]
|
||||
_WGS_HAPLOTYPE_PARAMS = [*_WGS_COMMON,
|
||||
{"name": "native_pair_hmm_threads", "type": "string", "default": "4",
|
||||
"source": "user", "help": "GATK PairHMM 线程"},
|
||||
{"name": "max_jobs", "type": "string", "default": "8", "source": "user",
|
||||
"help": "并行任务上限"}]
|
||||
_WGS_IMPUTE_PARAMS = [*_WGS_COMMON, _u("beagle_jar", "Beagle JAR 路径(容器内可访问)")]
|
||||
_POP_COMMON = [*_WGS_COMMON, _u("vcf", "VCF 文件路径(容器内可访问)")]
|
||||
_GWAS_TRAIT_PARAMS = [*_WGS_COMMON,
|
||||
_u("phenotype", "表型数据文件路径"),
|
||||
_u("sample_column", "样本列名"),
|
||||
_u("trait_column", "性状列名"),
|
||||
_u("trait_name", "性状名")]
|
||||
_GWAS_COV_PARAMS = [*_WGS_COMMON,
|
||||
_u("q", "Q 矩阵路径"),
|
||||
_u("fam", "FAM 文件路径"),
|
||||
_u("k", "主成分数 K")]
|
||||
|
||||
# 通用输入/输出描述(契约文档摘要;契约文件为权威,这里给目录层最小描述)
|
||||
_IDENTIFY_INPUTS = [
|
||||
{"name": "protein_fasta", "type": "fasta", "required": True,
|
||||
@@ -460,6 +508,32 @@ _STAGE_PARAM_SCHEMAS: dict[str, tuple[list[dict[str, Any]], list[dict[str, Any]]
|
||||
{"name": "prefix", "type": "string", "default": "", "source": "contract", "help": "分析前缀"},
|
||||
], _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"rnaseq_qc": (_RNASEQ_QC_PARAMS, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"wgs_qc": (_WGS_QC_PARAMS, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"wgs_index": (_WGS_INDEX_PARAMS, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"wgs_align": (_WGS_ALIGN_PARAMS, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"wgs_markdup": (_WGS_MARKDUP_PARAMS, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"wgs_bam_qc": (_WGS_BAM_QC_PARAMS, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"wgs_haplotype_calling": (_WGS_HAPLOTYPE_PARAMS, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"wgs_gvcf_combine": (_WGS_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"wgs_joint_genotyping": (_WGS_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"wgs_variant_select": (_WGS_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"wgs_variant_filter": (_WGS_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"wgs_impute": (_WGS_IMPUTE_PARAMS, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"wgs_maf_missing_filter": (_WGS_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"pop_pca": (_POP_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"pop_phylogeny": (_POP_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"pop_structure": (_POP_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"pop_ld_decay": (_POP_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"pop_snp_density": (_POP_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"gwas_fmt": (_POP_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"gwas_kinship": (_POP_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"gwas_trait": (_GWAS_TRAIT_PARAMS, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"gwas_cov": (_GWAS_COV_PARAMS, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"gwas_mlm": (_WGS_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"gwas_plot": (_WGS_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"gwas_emmax": (_WGS_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"gwas_emplot": (_WGS_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"gwas_glm": (_WGS_COMMON, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"rnaseq_align": (_RNASEQ_ALIGN_PARAMS, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"rnaseq_quant": (_RNASEQ_QUANT_PARAMS, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
"rnaseq_sample_qc": (_RNASEQ_SAMPLE_QC_PARAMS, _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
@@ -477,6 +551,38 @@ _STAGE_PARAM_SCHEMAS: dict[str, tuple[list[dict[str, Any]], list[dict[str, Any]]
|
||||
], _DOWNSTREAM_INPUTS, _DOWNSTREAM_OUTPUTS),
|
||||
}
|
||||
|
||||
# ---- G9:单一参数 schema 访问点(Catalog/表单/预检/工具/验收都读这一份)----
|
||||
|
||||
|
||||
def _schema_params_inputs_outputs(stage: str) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any]]]:
|
||||
"""阶段参数/输入/输出的单一读取点(含未知阶段默认兜底)。"""
|
||||
if stage in _STAGE_PARAM_SCHEMAS:
|
||||
return _STAGE_PARAM_SCHEMAS[stage]
|
||||
params = [{"name": "cpu", "type": "choice", "choices": ["4", "8", "16", "32"],
|
||||
"default": "4", "source": "runtime", "help": "容器 CPU 数"}]
|
||||
inputs = _DOWNSTREAM_INPUTS if stage not in ("identify", "genome_prepare") else []
|
||||
outputs = _DOWNSTREAM_OUTPUTS
|
||||
return params, inputs, outputs
|
||||
|
||||
|
||||
def stage_schema(stage: str) -> dict[str, Any]:
|
||||
"""G9:单一参数 schema 访问点——阶段名/耗时/依赖/参数/输入/输出/输入角色 合并为一份。
|
||||
|
||||
Catalog/表单/预检/工具/验收都读这一份(唯一事实源)。
|
||||
"""
|
||||
from .stage_assets import stage_input_roles
|
||||
|
||||
params, inputs, outputs = _schema_params_inputs_outputs(stage)
|
||||
return {
|
||||
"name": STAGE_INFO_SNAPSHOT.get(stage, {}).get("name") or stage,
|
||||
"estimated_time": int(STAGE_INFO_SNAPSHOT.get(stage, {}).get("estimated_time") or 0),
|
||||
"dependencies": STAGE_DEPENDENCIES_SNAPSHOT.get(stage, []),
|
||||
"params": params,
|
||||
"inputs": inputs,
|
||||
"outputs": outputs,
|
||||
"input_roles": stage_input_roles(stage),
|
||||
}
|
||||
|
||||
|
||||
class StageCatalog:
|
||||
"""阶段目录(数据来自 stage_image_map.json + 手写快照,不运行时反射)。"""
|
||||
@@ -504,29 +610,19 @@ class StageCatalog:
|
||||
if not self.has(stage):
|
||||
return None
|
||||
entry = self._map.image_for(stage) or {}
|
||||
params, inputs, outputs = _STAGE_PARAM_SCHEMAS.get(
|
||||
stage, ([], _DOWNSTREAM_INPUTS if stage not in ("identify", "genome_prepare") else [], [])
|
||||
)
|
||||
if stage not in _STAGE_PARAM_SCHEMAS:
|
||||
params, inputs, outputs = (
|
||||
[{"name": "cpu", "type": "choice", "choices": ["4", "8", "16", "32"],
|
||||
"default": "4", "source": "runtime", "help": "容器 CPU 数"}],
|
||||
_DOWNSTREAM_INPUTS if stage not in ("identify", "genome_prepare") else [],
|
||||
_DOWNSTREAM_OUTPUTS,
|
||||
)
|
||||
info = STAGE_INFO_SNAPSHOT.get(stage, {})
|
||||
schema = stage_schema(stage)
|
||||
return StageDef(
|
||||
name=stage,
|
||||
description=str(info.get("name") or stage),
|
||||
description=str(schema["name"] or stage),
|
||||
image=str(entry.get("image") or ""),
|
||||
module=str(entry.get("module") or ""),
|
||||
executable=bool(entry.get("executable")),
|
||||
estimated_time=int(info.get("estimated_time") or 0),
|
||||
params=params,
|
||||
inputs=inputs,
|
||||
outputs=outputs,
|
||||
estimated_time=int(schema["estimated_time"]),
|
||||
params=schema["params"],
|
||||
inputs=schema["inputs"],
|
||||
outputs=schema["outputs"],
|
||||
note=str(entry.get("note") or ""),
|
||||
dependencies=STAGE_DEPENDENCIES_SNAPSHOT.get(stage, []),
|
||||
dependencies=schema["dependencies"],
|
||||
)
|
||||
|
||||
def resolve_chain(self, chain: str | None) -> list[str] | None:
|
||||
@@ -538,32 +634,52 @@ class StageCatalog:
|
||||
return None
|
||||
return list(alias)
|
||||
|
||||
def auto_fill_prerequisites(self, stages: list[str]) -> tuple[list[str], list[str]]:
|
||||
def auto_fill_prerequisites(
|
||||
self, stages: list[str], satisfied_roles: list[str] | None = None
|
||||
) -> tuple[list[str], list[str]]:
|
||||
"""前置补全(FR-BIO-003):缺失的前置阶段自动补全。
|
||||
|
||||
G7/G8:satisfied_roles(已有 validated 资产的输入角色)对应的上游阶段
|
||||
跳过不补跑(如 role=counts 满足时做 DE 不补跑 qc/align/quant/sample_qc)。
|
||||
|
||||
返回 (补全后的链, 被补全的阶段名列表)。
|
||||
链内阶段按依赖拓扑顺序排列(依赖在前)。
|
||||
"""
|
||||
from .stage_assets import upstream_stages_for_roles
|
||||
|
||||
skip_stages = upstream_stages_for_roles(satisfied_roles or [])
|
||||
missing: list[str] = []
|
||||
added: list[str] = []
|
||||
chain_set = set(stages)
|
||||
# 反复补全直到闭合(依赖的依赖也补)
|
||||
# 反复补全直到闭合(依赖的依赖也补;被角色资产满足的上游跳过)
|
||||
changed = True
|
||||
while changed:
|
||||
changed = False
|
||||
for s in list(stages):
|
||||
for dep in STAGE_DEPENDENCIES_SNAPSHOT.get(s, []):
|
||||
if dep not in chain_set and self.has(dep):
|
||||
if dep in chain_set:
|
||||
continue
|
||||
if dep in skip_stages:
|
||||
continue
|
||||
if self.has(dep):
|
||||
chain_set.add(dep)
|
||||
missing.append(dep)
|
||||
added.append(dep)
|
||||
changed = True
|
||||
ordered = topo_order([s for s in GENE_FAMILY_CHAIN if s in chain_set]
|
||||
or list(chain_set))
|
||||
# G7/G8:被角色资产满足的上游阶段从最终链中剔除(不补跑)
|
||||
if skip_stages:
|
||||
ordered = [s for s in ordered if s not in skip_stages]
|
||||
return ordered, added
|
||||
|
||||
def dependencies(self, stage: str) -> list[str]:
|
||||
return STAGE_DEPENDENCIES_SNAPSHOT.get(stage, [])
|
||||
return stage_schema(stage)["dependencies"]
|
||||
|
||||
def stage_input_roles(self, stage: str) -> list[str]:
|
||||
"""G8:阶段依赖的输入角色(数据角色)。"""
|
||||
from .stage_assets import stage_input_roles
|
||||
return stage_input_roles(stage)
|
||||
|
||||
|
||||
def topo_order(stages: list[str]) -> list[str]:
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
"""app/services/bioflow/delivery.py —— 交付清单(G21)+ 终报告范围(G22)+ 资源预算(G23)。
|
||||
|
||||
- G21 交付清单完整性:delivery_validate(manifest) 做八项完整性检查
|
||||
(selected_stages / run_id / input_sources / output_sources / config /
|
||||
required_results / validation_values / review_status),返回
|
||||
{ok, checks{八项:bool}, incomplete[{item, reason, error(G13)}]}。
|
||||
- G22 终报告范围:classify_report_scope(status, source) 五分类
|
||||
(user_provided / reused / not_executed / skipped / failed)。
|
||||
- G23 资源预算:check_download_budget(size_bytes, budget_bytes)(超限→
|
||||
DOWNLOAD_BUDGET_EXCEEDED)+ check_disk_available(dest_dir, needed_bytes)
|
||||
(不足→DISK_FULL),均 ok=False + 结构化错误。
|
||||
|
||||
服务层纪律:纯 dict/stdlib,无 pydantic。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .errors import error_object, error_result
|
||||
from .source_fetcher import check_disk_space
|
||||
|
||||
# G21 八项完整性检查项
|
||||
DELIVERY_ITEMS = (
|
||||
"selected_stages", "run_id", "input_sources", "output_sources",
|
||||
"config", "required_results", "validation_values", "review_status",
|
||||
)
|
||||
|
||||
# G22 五分类
|
||||
SCOPE_USER_PROVIDED = "user_provided"
|
||||
SCOPE_REUSED = "reused"
|
||||
SCOPE_NOT_EXECUTED = "not_executed"
|
||||
SCOPE_SKIPPED = "skipped"
|
||||
SCOPE_FAILED = "failed"
|
||||
|
||||
|
||||
def delivery_validate(manifest: dict[str, Any]) -> dict[str, Any]:
|
||||
"""G21:校验交付清单完整性(八项)。
|
||||
|
||||
返回 {"ok": bool, "checks": {item: bool},
|
||||
"incomplete": [{"item", "reason", "error": {G13 对象}}]}。
|
||||
"""
|
||||
manifest = manifest or {}
|
||||
checks: dict[str, bool] = {}
|
||||
incomplete: list[dict[str, Any]] = []
|
||||
for item in DELIVERY_ITEMS:
|
||||
value = manifest.get(item)
|
||||
ok = bool(value)
|
||||
checks[item] = ok
|
||||
if not ok:
|
||||
incomplete.append({
|
||||
"item": item,
|
||||
"reason": f"缺 {item}",
|
||||
"error": error_object("INCOMPLETE_DELIVERY", f"交付清单缺项: {item}",
|
||||
category="input"),
|
||||
})
|
||||
return {"ok": not incomplete, "checks": checks, "incomplete": incomplete}
|
||||
|
||||
|
||||
def publish_delivery(manifest: dict[str, Any]) -> dict[str, Any]:
|
||||
"""G21:发布交付清单(校验通过置 review_status=approved)。"""
|
||||
check = delivery_validate(manifest)
|
||||
if not check["ok"]:
|
||||
return {"ok": False, "published": False, **check}
|
||||
manifest["review_status"] = "approved"
|
||||
return {"ok": True, "published": True, "review_status": "approved",
|
||||
"incomplete": []}
|
||||
|
||||
|
||||
# ---- G22 终报告范围 ----
|
||||
|
||||
def classify_report_scope(status: str, source: str | None = None) -> str:
|
||||
"""G22:阶段状态 → 终报告范围五分类。
|
||||
|
||||
failed→failed;skipped→skipped;done + source=reuse→reused;
|
||||
done + source=user/空→user_provided;其余非终态→not_executed。
|
||||
"""
|
||||
s = str(status or "")
|
||||
if s == "failed":
|
||||
return SCOPE_FAILED
|
||||
if s == "skipped":
|
||||
return SCOPE_SKIPPED
|
||||
if s == "done":
|
||||
return SCOPE_REUSED if source == "reuse" else SCOPE_USER_PROVIDED
|
||||
return SCOPE_NOT_EXECUTED
|
||||
|
||||
|
||||
def build_report(
|
||||
manifest: dict[str, Any],
|
||||
stage_sources: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""G22:终报告结构(执行范围 / 数据来源与核验 / 方法版本 / QC / 科学前提 /
|
||||
失败跳过未覆盖 / 产物清单)。"""
|
||||
stages = list(manifest.get("selected_stages") or [])
|
||||
sources = dict(stage_sources or {})
|
||||
scope = [{"stage": s, "source": sources.get(s, SCOPE_USER_PROVIDED)} for s in stages]
|
||||
|
||||
failed = [s for s in stages if sources.get(s) == SCOPE_FAILED]
|
||||
skipped = [s for s in stages if sources.get(s) == SCOPE_SKIPPED]
|
||||
not_executed = [s for s in stages if sources.get(s) == SCOPE_NOT_EXECUTED]
|
||||
|
||||
config = manifest.get("config") or {}
|
||||
return {
|
||||
"execution_scope": scope,
|
||||
"data_sources": {
|
||||
"inputs": list(manifest.get("input_sources") or []),
|
||||
"outputs": list(manifest.get("output_sources") or []),
|
||||
},
|
||||
"method_version": str(manifest.get("plan_revision") or manifest.get("run_id") or ""),
|
||||
"qc": {"validation_values": dict(manifest.get("validation_values") or {})},
|
||||
"scientific_preconditions": config.get("scientific_preconditions", []),
|
||||
"failed_skipped_uncovered": {
|
||||
"failed": failed, "skipped": skipped, "not_executed": not_executed,
|
||||
},
|
||||
"artifacts": list(manifest.get("required_results") or []),
|
||||
}
|
||||
|
||||
|
||||
# ---- G23 资源预算门禁 ----
|
||||
|
||||
def check_download_budget(size_bytes: int, budget_bytes: int) -> dict[str, Any]:
|
||||
"""G23:下载预算门禁——超限阻断并返回结构化错误。"""
|
||||
if size_bytes > budget_bytes:
|
||||
return error_result(
|
||||
"DOWNLOAD_BUDGET_EXCEEDED",
|
||||
f"下载预算超限:本次 {size_bytes} > 预算 {budget_bytes}",
|
||||
)
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
def check_disk_available(dest_dir: str | Path, needed_bytes: int) -> dict[str, Any]:
|
||||
"""G23:磁盘门禁——剩余空间不足阻断并返回结构化错误。"""
|
||||
try:
|
||||
ok, free = check_disk_space(dest_dir, needed_bytes)
|
||||
except OSError as exc:
|
||||
return error_result("PATH_NOT_EXIST", f"目标目录不可用:{exc}")
|
||||
if not ok:
|
||||
return error_result(
|
||||
"DISK_FULL", f"磁盘空间不足:需 {needed_bytes} 字节,剩余 {free} 字节",
|
||||
)
|
||||
return {"ok": True, "free_bytes": free}
|
||||
@@ -0,0 +1,196 @@
|
||||
"""app/services/bioflow/errors.py —— 统一结构化错误对象(G13)+ 确定性分类器(G14)。
|
||||
|
||||
失败返回统一携带结构化错误对象(替代裸字符串 message):
|
||||
{"ok": False, "code": "<码>", "message": "<摘要>",
|
||||
"error": {error_id, category, code, summary, evidence_refs,
|
||||
diagnosis_status, retryable_without_change, suggested_actions}}
|
||||
|
||||
8 类 category(G14:确定性映射):
|
||||
input / scientific_precondition / reference / environment /
|
||||
transient / resource / adapter / artifact
|
||||
|
||||
G14 分类纪律:退出码/日志是「证据」不是「结论」——
|
||||
- 日志含 OOM → resource;日志含镜像缺失 → environment;日志含科学前提不满足 → scientific_precondition。
|
||||
- 仅凭「非零退出码」无法定类 → adapter + undiagnosed(不臆测)。
|
||||
|
||||
服务层纪律:纯 dict/stdlib,无 pydantic。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
CATEGORIES = (
|
||||
"input", "scientific_precondition", "reference", "environment",
|
||||
"transient", "resource", "adapter", "artifact",
|
||||
)
|
||||
|
||||
# 明显映射:错误码 → (category, diagnosis_status, retryable_without_change, suggested_actions)
|
||||
_CODE_META: dict[str, tuple[str, str, bool, list[str]]] = {
|
||||
# ---- resource(资源耗尽)----
|
||||
"DISK_FULL": ("resource", "diagnosed", False, ["清理磁盘空间后重试"]),
|
||||
"UPLOAD_TOO_LARGE": ("resource", "diagnosed", False, ["减小文件大小或调大上传上限"]),
|
||||
"DOWNLOAD_BUDGET_EXCEEDED": ("resource", "diagnosed", False, ["调大下载预算或减小下载量"]),
|
||||
# ---- environment(运行环境)----
|
||||
"OSS_UNCONFIGURED": ("environment", "diagnosed", False, ["配置 BIOFLOW_OSS_ENDPOINT 或改用预签名 URL"]),
|
||||
"SSHPASS_MISSING": ("environment", "diagnosed", False, ["本机安装 sshpass 或改用密钥认证"]),
|
||||
# ---- reference(参考数据)----
|
||||
"REFERENCE_NOT_FOUND": ("reference", "diagnosed", False, ["补齐参考基因组/HMM/索引后重试"]),
|
||||
# ---- transient(临时/可重试)----
|
||||
"NETWORK_ERROR": ("transient", "diagnosed", True, ["重试下载"]),
|
||||
"LAN_TIMEOUT": ("transient", "diagnosed", True, ["重试局域网拉取"]),
|
||||
"TRUNCATED": ("transient", "diagnosed", True, ["重试下载"]),
|
||||
"LAN_ERROR": ("transient", "diagnosed", True, ["重试局域网拉取"]),
|
||||
"HTTP_ERROR": ("transient", "diagnosed", True, ["确认 URL 可达或稍后重试"]),
|
||||
# ---- input(用户输入/文件/参数)----
|
||||
"EMPTY_SOURCE": ("input", "diagnosed", False, ["提供非空源地址"]),
|
||||
"BAD_SCHEME": ("input", "diagnosed", False, ["改用 http/https/oss 协议"]),
|
||||
"BAD_OSS": ("input", "diagnosed", False, ["修正 OSS 地址为 oss://bucket/key"]),
|
||||
"PATH_NOT_EXIST": ("input", "diagnosed", False, ["确认路径存在且可达"]),
|
||||
"IS_DIR": ("input", "diagnosed", False, ["改为单文件路径"]),
|
||||
"NOT_FOUND": ("input", "diagnosed", False, ["确认目标存在"]),
|
||||
"BAD_FORMAT": ("input", "diagnosed", False, ["选择支持的目标格式 fasta/gff3/tsv/samples.yaml"]),
|
||||
"EMPTY_TABLE": ("input", "diagnosed", False, ["提供非空表格数据"]),
|
||||
"MISSING_COLUMNS": ("input", "diagnosed", False, ["确认列映射(名称/序列/r1/r2)"]),
|
||||
"VALIDATION": ("input", "diagnosed", False, ["修正输入以满足校验"]),
|
||||
"VALIDATION_FAILED": ("input", "diagnosed", False, ["修正数据后重转"]),
|
||||
"UNKNOWN_STAGE": ("input", "diagnosed", False, ["选择有效阶段"]),
|
||||
"UNKNOWN_PARAMS": ("input", "diagnosed", False, ["移除未知参数"]),
|
||||
"BAD_ROLE": ("input", "diagnosed", False, ["使用合法角色名(非空、无路径分隔符/..)"]),
|
||||
"EMPTY_REPORT": ("input", "diagnosed", False, ["提供非空报告路径/文件"]),
|
||||
"EMPTY_SAMPLES": ("input", "diagnosed", False, ["提供含样本的描述"]),
|
||||
"INCOMPLETE_DELIVERY": ("input", "diagnosed", False, ["补齐交付清单缺失项(selected_stages/run_id/必需结果等)后重试"]),
|
||||
# ---- adapter(适配/映射)----
|
||||
"UNSUPPORTED": ("adapter", "diagnosed", False, ["改用支持的转换路径"]),
|
||||
"CROSS_CHAIN": ("adapter", "diagnosed", False, ["选择本能力单元内的阶段"]),
|
||||
# ---- artifact(产物)----
|
||||
"ARCHIVE_EXISTS": ("artifact", "diagnosed", False, ["更换归档名"]),
|
||||
"ARCHIVE_FAILED": ("artifact", "diagnosed", False, ["检查任务目录完整性后重试"]),
|
||||
"ARCHIVE_NOT_FOUND": ("artifact", "diagnosed", False, ["确认归档名存在"]),
|
||||
"ASSET_EXISTS": ("artifact", "diagnosed", False, ["更换角色或文件名"]),
|
||||
"PUBLISH_FAILED": ("artifact", "diagnosed", False, ["检查目标目录可写性后重试"]),
|
||||
"NO_MANIFEST": ("artifact", "diagnosed", False, ["检查阶段产物 manifest 是否产出"]),
|
||||
# ---- scientific_precondition(科学前提)----
|
||||
"SCIENTIFIC_PRECONDITION": ("scientific_precondition", "diagnosed", False, ["确认科学前提满足后重试"]),
|
||||
}
|
||||
|
||||
# 关键词兜底(summary 命中即映射;用于非枚举码)
|
||||
_KEYWORD_META: list[tuple[tuple[str, ...], tuple[str, str, bool, list[str]]]] = [
|
||||
(("out of memory", "oom", "memoryerror", "oom-killed", "memory limit", "内存"),
|
||||
("resource", "diagnosed", False, ["调大内存或减小数据规模"])),
|
||||
(("no such image", "image not found", "pull access denied", "镜像", "docker"),
|
||||
("environment", "diagnosed", False, ["确认镜像存在且可拉取"])),
|
||||
(("no significant", "not enough", "insufficient", "科学前提", "no differentially", "not found enough"),
|
||||
("scientific_precondition", "diagnosed", False, ["确认科学前提满足后重试"])),
|
||||
(("timeout", "超时", "network", "网络", "unreachable", "connecterror", "timed out"),
|
||||
("transient", "diagnosed", True, ["重试"])),
|
||||
(("reference", "参考", "hmm", "pfam", "gtf", "index"),
|
||||
("reference", "diagnosed", False, ["补齐参考数据后重试"])),
|
||||
(("manifest", "no_manifest", "产物缺失"),
|
||||
("artifact", "diagnosed", False, ["检查阶段产物 manifest"])),
|
||||
]
|
||||
|
||||
# 阶段失败 reason(runner _fail_stage)→ (category, diagnosis_status, retryable, actions)
|
||||
_REASON_META: dict[str, tuple[str, str, bool, list[str]]] = {
|
||||
"timeout": ("transient", "diagnosed", True, ["重跑该阶段"]),
|
||||
"exit_code": ("adapter", "undiagnosed", False, ["查看阶段日志定位失败原因", "修正后重跑"]),
|
||||
"no_manifest": ("artifact", "diagnosed", False, ["检查阶段是否产出 manifest"]),
|
||||
"manifest_error": ("artifact", "diagnosed", False, ["检查 manifest JSON 完整性"]),
|
||||
"manifest_failed": ("artifact", "diagnosed", False, ["检查阶段产物 manifest status"]),
|
||||
}
|
||||
|
||||
# 退出码 137 = 128 + SIGKILL(常为 OOM kill)
|
||||
_OOM_EXIT_CODES = {137, 139}
|
||||
|
||||
|
||||
def classify(code: str | None, summary: str = "") -> tuple[str, str, bool, list[str]]:
|
||||
"""错误码/摘要 → (category, diagnosis_status, retryable_without_change, suggested_actions)。"""
|
||||
code = str(code or "")
|
||||
summary = str(summary or "")
|
||||
if code in _CODE_META:
|
||||
cat, diag, retry, actions = _CODE_META[code]
|
||||
return cat, diag, retry, list(actions)
|
||||
low = summary.lower()
|
||||
for kws, meta in _KEYWORD_META:
|
||||
if any(k in low for k in kws):
|
||||
cat, diag, retry, actions = meta
|
||||
return cat, diag, retry, list(actions)
|
||||
return "adapter", "undiagnosed", False, []
|
||||
|
||||
|
||||
def classify_reason(
|
||||
reason: str | None,
|
||||
*,
|
||||
exit_code: int | None = None,
|
||||
log_tail: str | None = None,
|
||||
) -> tuple[str, str, bool, list[str]]:
|
||||
"""阶段失败原因 + 证据(退出码/日志)→ 确定性 8 类。
|
||||
|
||||
确定性 reason 优先(timeout / manifest 系列);泛化 exit_code 用证据(退出码/日志)
|
||||
定类——仅凭非零退出码不定类(undiagnosed),不臆测。
|
||||
"""
|
||||
reason = str(reason or "")
|
||||
log = (log_tail or "").lower()
|
||||
|
||||
# 1. 确定性 reason(非 exit_code)优先:timeout / manifest 系列
|
||||
if reason in _REASON_META and reason != "exit_code":
|
||||
cat, diag, retry, actions = _REASON_META[reason]
|
||||
return cat, diag, retry, list(actions)
|
||||
|
||||
# 2. 泛化 exit_code / 未知 reason:退出码是证据(137/139 = OOM kill)
|
||||
if exit_code in _OOM_EXIT_CODES:
|
||||
return "resource", "diagnosed", False, ["容器可能被 OOM kill,调大内存后重跑"]
|
||||
|
||||
# 3. 日志证据定类
|
||||
for kws, meta in _KEYWORD_META:
|
||||
if any(k in log for k in kws):
|
||||
cat, diag, retry, actions = meta
|
||||
return cat, diag, retry, list(actions)
|
||||
|
||||
# 4. 兜底:exit_code → adapter undiagnosed;未知 → 同
|
||||
if reason in _REASON_META:
|
||||
cat, diag, retry, actions = _REASON_META[reason]
|
||||
return cat, diag, retry, list(actions)
|
||||
return "adapter", "undiagnosed", False, []
|
||||
|
||||
|
||||
def error_object(
|
||||
code: str | None,
|
||||
summary: str = "",
|
||||
*,
|
||||
category: str | None = None,
|
||||
evidence_refs: list[str] | None = None,
|
||||
diagnosis_status: str | None = None,
|
||||
retryable_without_change: bool | None = None,
|
||||
suggested_actions: list[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""构造统一结构化错误对象(8 字段)。"""
|
||||
cat, diag, retry, actions = classify(code, summary)
|
||||
return {
|
||||
"error_id": uuid.uuid4().hex[:16],
|
||||
"category": category or cat,
|
||||
"code": code or "INTERNAL",
|
||||
"summary": summary,
|
||||
"evidence_refs": list(evidence_refs or []),
|
||||
"diagnosis_status": diagnosis_status or diag,
|
||||
"retryable_without_change": retry if retryable_without_change is None else retryable_without_change,
|
||||
"suggested_actions": list(suggested_actions or actions),
|
||||
}
|
||||
|
||||
|
||||
def error_result(
|
||||
code: str | None,
|
||||
summary: str = "",
|
||||
*,
|
||||
evidence_refs: list[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""失败返回统一信封:{"ok": False, "code", "message", "error"}。
|
||||
|
||||
code 同时出现在顶层与 error 对象内(兼容 Phase A 的 r["code"] 断言)。
|
||||
"""
|
||||
return {
|
||||
"ok": False,
|
||||
"code": code or "INTERNAL",
|
||||
"message": summary,
|
||||
"error": error_object(code, summary, evidence_refs=evidence_refs),
|
||||
}
|
||||
@@ -23,9 +23,15 @@ from typing import Any
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
EVENT_TYPES = (
|
||||
"job_created", "plan_ready", "stage_started", "stage_log", "stage_done",
|
||||
"stage_failed", "stage_skipped", "stage_await_confirm", "stage_await_files",
|
||||
"job_created", "plan_ready", "stage_started", "stage_log", "stage_progress",
|
||||
"stage_done", "stage_failed", "stage_skipped", "stage_partial",
|
||||
"stage_await_confirm", "stage_await_files",
|
||||
"job_done", "job_paused", "job_archived", "job_recovered",
|
||||
# 用户自定义扩展分析事件帧(契约 §4;并行链路,不触碰现有 job_* 语义)
|
||||
"extension_created", "extension_build", "extension_trial", "extension_repair",
|
||||
"extension_validate", "extension_published", "extension_candidates",
|
||||
"extension_confirm", "extension_run_started", "extension_run_log",
|
||||
"extension_run_done", "extension_run_failed",
|
||||
)
|
||||
|
||||
# 断线重连时需要回放的关键状态事件(含历史事件全量)
|
||||
@@ -71,9 +77,34 @@ class EventBus:
|
||||
self._subscribers: set[Subscriber] = set()
|
||||
self._sub_lock = asyncio.Lock()
|
||||
self._buffer: deque[dict[str, Any]] = deque(maxlen=MEMORY_BUFFER_LIMIT)
|
||||
self._seq = 0
|
||||
# 跨重启 seq 单调递增(BUG-FIX 2026-09-17):旧进程事件已持久化在
|
||||
# events/*.jsonl,重启后 _seq 归零会让新事件 seq 小于浏览器 last_seq
|
||||
# 水位,前端回放被过滤、收不到 stage_await_files 等关键帧。
|
||||
# 启动时扫描历史日志取最大 seq 作为起始值。
|
||||
self._seq = self._scan_max_seq()
|
||||
self._file_lock = threading.Lock()
|
||||
|
||||
def _scan_max_seq(self) -> int:
|
||||
"""扫描 events_dir 下全部 jsonl,取历史最大 seq(无文件/坏行 → 0)。"""
|
||||
max_seq = 0
|
||||
try:
|
||||
for path in self.events_dir.glob("*.jsonl"):
|
||||
try:
|
||||
for line in path.read_text(encoding="utf-8").splitlines():
|
||||
if not line.strip():
|
||||
continue
|
||||
try:
|
||||
seq = int(json.loads(line).get("seq") or 0)
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
if seq > max_seq:
|
||||
max_seq = seq
|
||||
except OSError:
|
||||
continue
|
||||
except OSError: # pragma: no cover - 目录扫描失败(权限等)不影响启动
|
||||
return 0
|
||||
return max_seq
|
||||
|
||||
# ---- 内部 ----
|
||||
|
||||
def _events_file(self, job_id: str) -> Path:
|
||||
@@ -177,6 +208,17 @@ class EventBus:
|
||||
elif t == "stage_skipped":
|
||||
st = stage_states.setdefault(str(d.get("stage")), {"stage": str(d.get("stage"))})
|
||||
st["status"] = "skipped"
|
||||
elif t == "stage_partial":
|
||||
st = stage_states.setdefault(str(d.get("stage")), {"stage": str(d.get("stage"))})
|
||||
st["status"] = "partial_success"
|
||||
elif t == "stage_progress":
|
||||
st = stage_states.setdefault(str(d.get("stage")), {"stage": str(d.get("stage"))})
|
||||
pct = d.get("percent")
|
||||
if isinstance(pct, (int, float)):
|
||||
st["progress"] = max(float(st.get("progress") or 0.0), float(pct))
|
||||
msg = d.get("message")
|
||||
if msg:
|
||||
st["progressMessage"] = str(msg)
|
||||
elif t == "stage_await_confirm":
|
||||
st = stage_states.setdefault(str(d.get("stage")), {"stage": str(d.get("stage"))})
|
||||
st["status"] = "waiting_confirm"
|
||||
|
||||
@@ -0,0 +1,286 @@
|
||||
"""app/services/bioflow/extension_exec.py —— 扩展分析 docker 执行器(并行链路,零触碰 kernel.py)。
|
||||
|
||||
- 宿主挂载通用入口:run_extension.py 以 -v <bioflow_kernel/cli>:/opt/bioflow-entry:ro
|
||||
挂进任意镜像执行,8 个 base 镜像与自动生成镜像零改动(契约 §3)。
|
||||
- 命令拼装:list-form docker run(无 shell)—— --rm/--user/三数据卷 + 入口卷/
|
||||
--cpus/--memory 白名单/镜像:dev/python -B /opt/bioflow-entry/run_extension.py。
|
||||
- 白名单 helper 为 kernel.py 同款纪律的独立复制(硬约束:不 import kernel.py;
|
||||
并行链路各自演进,互不耦合)。
|
||||
- 超时:terminate → 3s → kill(与 kernel.run_stage 同约定)。
|
||||
- 返回 {exit_code, timed_out, log_tail, cmd};退出码 3 = 输出契约校验失败。
|
||||
|
||||
服务层纪律:纯 dataclass/dict/stdlib,无 pydantic;执行经 async subprocess。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import shlex
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CONTAINER_EXT_DIR = "/job/ext"
|
||||
CONTAINER_INPUT_DIR = "/job/input"
|
||||
CONTAINER_OUTPUT_DIR = "/job/output"
|
||||
CONTAINER_ENTRY_DIR = "/opt/bioflow-entry"
|
||||
CONTAINER_ENTRY = "/opt/bioflow-entry/run_extension.py"
|
||||
EXIT_OUTPUT_CONTRACT = 3
|
||||
|
||||
|
||||
class ExtensionExecError(RuntimeError):
|
||||
"""扩展执行前置校验失败(镜像空/入口脚本缺失)。"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExtensionExecConfig:
|
||||
"""扩展执行配置(由 app.config 注入;测试可构造最小实例)。
|
||||
|
||||
白名单纪律与 KernelConfig 对齐但独立持有(kernel.py 不 import)。"""
|
||||
|
||||
host_uid: int = 0
|
||||
docker_tag_suffix: str = "dev"
|
||||
entry_script_dir: str = "" # 宿主 bioflow_kernel/cli 绝对路径
|
||||
container_ext_dir: str = CONTAINER_EXT_DIR
|
||||
container_input_dir: str = CONTAINER_INPUT_DIR
|
||||
container_output_dir: str = CONTAINER_OUTPUT_DIR
|
||||
container_entry_dir: str = CONTAINER_ENTRY_DIR
|
||||
default_cpus: str = "4"
|
||||
default_memory: str = "32g"
|
||||
timeout_sec: float = 36000.0
|
||||
log_tail_lines: int = 80
|
||||
cpu_allow: str = "4,8,16,32"
|
||||
memory_allow: str = "16g,32g,64g"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExtensionRunRequest:
|
||||
"""一次扩展执行的输入形态(与契约 §1 runs/{run_id} 布局对应)。"""
|
||||
|
||||
image: str
|
||||
ext_dir: str | Path
|
||||
input_dir: str | Path
|
||||
output_dir: str | Path
|
||||
params: dict[str, Any] = field(default_factory=dict)
|
||||
cpu: str | None = None
|
||||
memory: str | None = None
|
||||
|
||||
|
||||
class ExtensionExecutor:
|
||||
"""扩展 docker 执行:命令拼装(纯函数式,可 mock 断言)+ 异步执行。"""
|
||||
|
||||
def __init__(self, config: ExtensionExecConfig) -> None:
|
||||
self.cfg = config
|
||||
|
||||
# ---- 白名单(独立复制 kernel.py 纪律;并行链路互不 import)----
|
||||
|
||||
def _safe_cpu(self, cpu: str) -> str:
|
||||
if str(cpu).strip() in {c.strip() for c in self.cfg.cpu_allow.split(",")}:
|
||||
return str(cpu).strip()
|
||||
logger.warning("extension exec: cpu=%s 不在白名单 %s,回落默认 %s",
|
||||
cpu, self.cfg.cpu_allow, self.cfg.default_cpus)
|
||||
return self.cfg.default_cpus
|
||||
|
||||
def _safe_memory(self, memory: str) -> str:
|
||||
if str(memory).strip() in {m.strip() for m in self.cfg.memory_allow.split(",")}:
|
||||
return str(memory).strip()
|
||||
logger.warning("extension exec: memory=%s 不在白名单 %s,回落默认 %s",
|
||||
memory, self.cfg.memory_allow, self.cfg.default_memory)
|
||||
return self.cfg.default_memory
|
||||
|
||||
# ---- 命令拼装 ----
|
||||
|
||||
def build_command(self, request: ExtensionRunRequest) -> list[str]:
|
||||
"""拼装 docker run 命令(契约 §3 形态;list 传参无 shell 注入)。
|
||||
|
||||
顺序:docker,run,--rm,--user,uid,-v(ext:ro/input/output/entry:ro),
|
||||
--cpus,--memory,<image>:dev,python,-B,/opt/bioflow-entry/run_extension.py,
|
||||
--ext-dir,--input-dir,--output-dir,[--param name=value…]
|
||||
"""
|
||||
if not str(request.image).strip():
|
||||
raise ExtensionExecError("扩展无镜像(runtime_image 为空)")
|
||||
entry = Path(self.cfg.entry_script_dir) / "run_extension.py"
|
||||
if not entry.is_file():
|
||||
raise ExtensionExecError(f"入口脚本缺失: {entry}")
|
||||
cmd: list[str] = ["docker", "run", "--rm", "--user", str(self.cfg.host_uid)]
|
||||
cmd += ["-v", f"{request.ext_dir}:{self.cfg.container_ext_dir}:ro"]
|
||||
cmd += ["-v", f"{request.input_dir}:{self.cfg.container_input_dir}"]
|
||||
cmd += ["-v", f"{request.output_dir}:{self.cfg.container_output_dir}"]
|
||||
cmd += ["-v", f"{self.cfg.entry_script_dir}:{self.cfg.container_entry_dir}:ro"]
|
||||
cmd += ["--cpus", self._safe_cpu(request.cpu or self.cfg.default_cpus)]
|
||||
cmd += ["--memory", self._safe_memory(request.memory or self.cfg.default_memory)]
|
||||
# 镜像名可能已含 tag(自动生成的 bioflow-ext-<fp>:<suffix>);
|
||||
# 未含 tag 时才补 suffix——否则拼成 a:b:b 的非法引用(真实 bug)。
|
||||
image_ref = str(request.image)
|
||||
cmd.append(image_ref if ":" in image_ref else f"{image_ref}:{self.cfg.docker_tag_suffix}")
|
||||
cmd += ["python", "-B", CONTAINER_ENTRY]
|
||||
cmd += ["--ext-dir", self.cfg.container_ext_dir]
|
||||
cmd += ["--input-dir", self.cfg.container_input_dir]
|
||||
cmd += ["--output-dir", self.cfg.container_output_dir]
|
||||
for key in sorted(request.params):
|
||||
value = request.params[key]
|
||||
if value is None:
|
||||
continue
|
||||
cmd += ["--param", f"{key}={value}"]
|
||||
return cmd
|
||||
|
||||
@staticmethod
|
||||
def command_display(cmd: list[str]) -> str:
|
||||
"""命令可读形态(日志用)。"""
|
||||
return " ".join(shlex.quote(str(c)) for c in cmd)
|
||||
|
||||
# ---- 执行 ----
|
||||
|
||||
async def run(self, request: ExtensionRunRequest,
|
||||
log_path: str | Path | None = None) -> dict[str, Any]:
|
||||
"""执行扩展 docker:返回 {exit_code, timed_out, log_tail, cmd}。
|
||||
|
||||
- 超时:terminate → 3s → kill(与 kernel.run_stage 同约定)
|
||||
- 日志全量落 log_path,返回尾部 log_tail_lines 行
|
||||
- 退出码 3 = 输出契约校验失败;本函数只报事实,语义由上层判定
|
||||
"""
|
||||
cmd = self.build_command(request)
|
||||
log_path = Path(log_path) if log_path else Path(request.output_dir) / "run.log"
|
||||
log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.STDOUT,
|
||||
)
|
||||
timed_out = False
|
||||
try:
|
||||
stdout, _ = await asyncio.wait_for(
|
||||
proc.communicate(), timeout=self.cfg.timeout_sec
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
timed_out = True
|
||||
try:
|
||||
proc.terminate()
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=3.0)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
stdout = b""
|
||||
exit_code = proc.returncode
|
||||
text = (stdout or b"").decode("utf-8", errors="replace")
|
||||
try:
|
||||
log_path.write_text(text, encoding="utf-8")
|
||||
except OSError:
|
||||
logger.warning("extension exec: 日志落盘失败 %s", log_path)
|
||||
return {
|
||||
"exit_code": exit_code,
|
||||
"timed_out": timed_out,
|
||||
"log_tail": _tail_lines(text, self.cfg.log_tail_lines),
|
||||
"cmd": cmd,
|
||||
}
|
||||
|
||||
async def run_stream(self, request: ExtensionRunRequest) -> Any:
|
||||
"""流式执行(增量日志推送;SSE 实时滚动)。
|
||||
|
||||
逐行读 stdout → 产出增量 chunk;结束后一并产出终态 dict
|
||||
(含 exit_code/timed_out/log_tail/full_log)。超时/取消对进程组
|
||||
terminate→3s→kill(与 run 同约定)。
|
||||
"""
|
||||
cmd = self.build_command(request)
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.STDOUT,
|
||||
)
|
||||
timed_out = False
|
||||
chunks: list[str] = []
|
||||
assert proc.stdout is not None
|
||||
while True:
|
||||
try:
|
||||
raw = await asyncio.wait_for(proc.stdout.readline(), timeout=self.cfg.timeout_sec)
|
||||
except asyncio.TimeoutError:
|
||||
timed_out = True
|
||||
await self._terminate_then_kill(proc)
|
||||
break
|
||||
except asyncio.CancelledError:
|
||||
await self._terminate_then_kill(proc)
|
||||
raise
|
||||
if not raw:
|
||||
break
|
||||
text = raw.decode("utf-8", errors="replace")
|
||||
chunks.append(text)
|
||||
yield {"chunk": text}
|
||||
await proc.wait()
|
||||
full = "".join(chunks)
|
||||
yield {
|
||||
"final": True,
|
||||
"exit_code": proc.returncode,
|
||||
"timed_out": timed_out,
|
||||
"log_tail": _tail_lines(full, self.cfg.log_tail_lines),
|
||||
"full_log": full,
|
||||
"cmd": cmd,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
async def _terminate_then_kill(proc: asyncio.subprocess.Process) -> None:
|
||||
"""terminate → 3s → kill(kernel.run_stage 同约定)。"""
|
||||
try:
|
||||
proc.terminate()
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=3.0)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
|
||||
|
||||
def _tail_lines(text: str, n: int) -> str:
|
||||
"""返回文本最后 n 行(与 kernel.tail_lines 同语义;独立复制)。"""
|
||||
lines = text.splitlines()
|
||||
if len(lines) <= n:
|
||||
return "\n".join(lines)
|
||||
return "\n".join(lines[-n:])
|
||||
|
||||
|
||||
def _default_exec_config() -> ExtensionExecConfig:
|
||||
"""从 app.config 推导(默认 uid = 服务进程 os.getuid())。"""
|
||||
from ...config import get_settings
|
||||
|
||||
s = get_settings()
|
||||
return ExtensionExecConfig(
|
||||
host_uid=s.host_uid,
|
||||
docker_tag_suffix=s.docker_image_tag_suffix,
|
||||
entry_script_dir=_resolve_extension_entry_dir(),
|
||||
default_cpus=s.default_cpus,
|
||||
default_memory=s.default_memory,
|
||||
timeout_sec=s.stage_timeout_sec,
|
||||
log_tail_lines=s.log_tail_lines,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_extension_entry_dir() -> str:
|
||||
"""宿主入口脚本目录 = 部署树 bioflow_kernel/cli(env 优先 + 多候选推导)。
|
||||
|
||||
2026-09-15 实测:parents[4] 推导在 175 仓库布局成立,176 部署树路径深度
|
||||
不同 → 找到不存在位置 → 试运行报「入口脚本缺失」。修复:
|
||||
1. BIOFLOW_EXT_ENTRY_DIR env 显式指定(部署树配置)
|
||||
2. 从 __file__ 向上逐级找 */bioflow_kernel/cli/run_extension.py
|
||||
3. 兜底 parents[4] 推导
|
||||
"""
|
||||
env_dir = os.environ.get("BIOFLOW_EXT_ENTRY_DIR", "").strip()
|
||||
if env_dir and Path(env_dir).is_dir():
|
||||
return env_dir
|
||||
here = Path(__file__).resolve()
|
||||
for parent in [here] + list(here.parents):
|
||||
candidate = parent / "bioflow_kernel" / "cli"
|
||||
if (candidate / "run_extension.py").is_file():
|
||||
return str(candidate)
|
||||
return str(here.parents[4] / "bioflow_kernel" / "cli")
|
||||
|
||||
|
||||
def build_extension_executor() -> ExtensionExecutor:
|
||||
"""应用启动:构造扩展执行器(入口脚本缺失不在此抛错,执行前置校验)。"""
|
||||
return ExtensionExecutor(_default_exec_config())
|
||||
@@ -0,0 +1,399 @@
|
||||
"""app/services/bioflow/extension_image.py —— 扩展镜像 176 本地自动生成(无 docker SDK)。
|
||||
|
||||
- ensure_image(base_image, requirements, tag_suffix):
|
||||
指纹 = sha256("<base>:<tag_suffix>\n" + 排序依赖行) 前 12 位;
|
||||
镜像名 = bioflow-ext-<fingerprint>:<tag_suffix>;
|
||||
docker inspect 命中即复用;未命中生成 Dockerfile
|
||||
(FROM base:tag + apt 安装 + /opt/conda/bin/pip install) → docker build
|
||||
(超时可配默认 30min) → 登记 registry 台账。
|
||||
- 失败分类:ImageBuildError.kind ∈ {pip_unavailable, apt_unavailable, network, timeout,
|
||||
docker_missing, inspect_failed, build_failed};可读 message 供 AI 修复循环/用户。
|
||||
- 失败清理:删除构建上下文目录与无 tag 中间镜像(docker rmi <image_id>)。
|
||||
- collect_garbage(extension_ids):registry 引用计数,无引用镜像 docker rmi 回收。
|
||||
- 全部经 subprocess 调 docker(list-form,无 shell)。
|
||||
|
||||
服务层纪律:纯 dataclass/dict/stdlib;registry 为 {data_root}/extensions/_images/registry.json。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
IMAGE_PREFIX = "bioflow-ext"
|
||||
PIP_BIN = "/opt/conda/bin/pip" # 镜像红线:镜像内无 conda 本体,直接调 pip 路径
|
||||
PIP_INDEX = "https://pypi.org/simple" # 176 对 pypi 直连
|
||||
|
||||
|
||||
class ImageBuildError(RuntimeError):
|
||||
"""镜像构建失败(kind 供 AI 修复循环/前端区分处理)。"""
|
||||
|
||||
def __init__(self, kind: str, message: str) -> None:
|
||||
super().__init__(message)
|
||||
self.kind = kind
|
||||
self.message = message
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImageConfig:
|
||||
"""镜像生成配置(由 app.config 注入;测试可构造最小实例)。"""
|
||||
|
||||
data_root: str = "" # 注册表父目录(extensions/_images)
|
||||
docker_tag_suffix: str = "dev"
|
||||
build_timeout_sec: float = 1800.0 # 默认 30min
|
||||
pip_index: str = PIP_INDEX
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImageRegistry:
|
||||
"""registry 台账内存形态 + 原子落盘(tmp 写 + os.replace)。"""
|
||||
|
||||
path: Path
|
||||
images: dict[str, dict[str, Any]] = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: Path) -> "ImageRegistry":
|
||||
if not path.exists():
|
||||
return cls(path=path)
|
||||
try:
|
||||
raw = json.loads(path.read_text(encoding="utf-8"))
|
||||
except (ValueError, OSError):
|
||||
logger.warning("extension image registry 损坏,重建: %s", path)
|
||||
return cls(path=path)
|
||||
images = raw.get("images") if isinstance(raw, dict) else {}
|
||||
return cls(path=path, images=images)
|
||||
|
||||
def save(self) -> None:
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = self.path.with_suffix(".json.tmp")
|
||||
tmp.write_text(
|
||||
json.dumps({"version": "1", "images": self.images},
|
||||
ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
tmp.replace(self.path)
|
||||
|
||||
def get(self, name: str) -> dict[str, Any] | None:
|
||||
return self.images.get(name)
|
||||
|
||||
def put(self, name: str, record: dict[str, Any]) -> None:
|
||||
self.images[name] = dict(record)
|
||||
self.save()
|
||||
|
||||
def drop(self, name: str) -> None:
|
||||
if self.images.pop(name, None) is not None:
|
||||
self.save()
|
||||
|
||||
def referenced(self, extension_ids: set[str]) -> set[str]:
|
||||
"""仍被给定扩展引用的镜像名集合(收集策略输入)。"""
|
||||
out: set[str] = set()
|
||||
for name, record in self.images.items():
|
||||
refs = record.get("refs") or []
|
||||
if any(str(ext_id) in refs for ext_id in extension_ids):
|
||||
out.add(name)
|
||||
return out
|
||||
|
||||
|
||||
def _sorted_requirements(requirements: dict[str, Any]) -> tuple[list[str], list[str]]:
|
||||
"""requirements 双形态兼容:{"pip": [...], "apt": [...]}(契约 §2)或
|
||||
{"python": [...], "system": [...]}(175 旧版)。"""
|
||||
if isinstance(requirements, dict):
|
||||
pip = list(requirements.get("pip") or requirements.get("python") or [])
|
||||
apt = list(requirements.get("apt") or requirements.get("system") or [])
|
||||
else:
|
||||
pip, apt = [], []
|
||||
return sorted({str(x).strip() for x in pip if str(x).strip()}), \
|
||||
sorted({str(x).strip() for x in apt if str(x).strip()})
|
||||
|
||||
|
||||
def image_fingerprint(base_image: str, requirements: dict[str, Any],
|
||||
tag_suffix: str = "dev") -> str:
|
||||
"""依赖指纹:sha256(base 镜像 + tag + 排序依赖)前 12 位。"""
|
||||
pip, apt = _sorted_requirements(requirements)
|
||||
payload = [f"{base_image}:{tag_suffix}"]
|
||||
for item in pip:
|
||||
payload.append(f"pip\t{item}")
|
||||
for item in apt:
|
||||
payload.append(f"apt\t{item}")
|
||||
digest = hashlib.sha256("\n".join(payload).encode("utf-8")).hexdigest()
|
||||
return digest[:12]
|
||||
|
||||
|
||||
def image_name_for(fingerprint: str, tag_suffix: str = "dev") -> str:
|
||||
"""镜像名 = bioflow-ext-<指纹12位>:<tag_suffix>(契约 §B2 去重命名)。"""
|
||||
return f"{IMAGE_PREFIX}-{fingerprint}:{tag_suffix}"
|
||||
|
||||
|
||||
def render_dockerfile(base_image: str, requirements: dict[str, Any],
|
||||
tag_suffix: str = "dev",
|
||||
pip_index: str = PIP_INDEX) -> str:
|
||||
"""生成 Dockerfile:FROM base:tag + LABEL(失败清理定位)+ apt(+可选)+ pip install。"""
|
||||
pip, apt = _sorted_requirements(requirements)
|
||||
fingerprint = image_fingerprint(base_image, requirements, tag_suffix)
|
||||
lines = ["# auto-generated by extension_image (fingerprint=%s)" % fingerprint,
|
||||
f"FROM {base_image}:{tag_suffix}",
|
||||
"LABEL plantai.ext.fingerprint=%s" % fingerprint]
|
||||
if apt:
|
||||
# 单依赖时尾反斜杠由 bash 忽略(命令结束符后接换行安全);多依赖续行
|
||||
lines.append("RUN apt-get update && apt-get install -y --no-install-recommends " +
|
||||
" ".join(apt) + " && rm -rf /var/lib/apt/lists/*")
|
||||
if pip:
|
||||
joined = " ".join(pip)
|
||||
lines.append(f"RUN {PIP_BIN} install --no-cache-dir -i {pip_index} {joined}")
|
||||
if not pip and not apt:
|
||||
lines.append("RUN true")
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
class ExtensionImageBuilder:
|
||||
"""176 本地自动镜像生成:指纹去重 + docker build + 失败清理 + registry 台账。"""
|
||||
|
||||
def __init__(self, config: ImageConfig) -> None:
|
||||
self.cfg = config
|
||||
self._registry_lock = threading.Lock()
|
||||
self._registry: ImageRegistry | None = None
|
||||
|
||||
# ---- registry(懒加载 + 锁;单进程多协程安全)----
|
||||
|
||||
def registry_path(self) -> Path:
|
||||
return Path(self.cfg.data_root) / "extensions" / "_images" / "registry.json"
|
||||
|
||||
def registry(self) -> ImageRegistry:
|
||||
with self._registry_lock:
|
||||
if self._registry is None or self._registry.path != self.registry_path():
|
||||
self._registry = ImageRegistry.load(self.registry_path())
|
||||
return self._registry
|
||||
|
||||
def refresh_registry(self) -> None:
|
||||
with self._registry_lock:
|
||||
self._registry = None
|
||||
|
||||
# ---- docker 子进程(测试注入 fake 断言命令形态)----
|
||||
|
||||
def _run_docker(self, cmd: list[str], timeout_sec: float) -> dict[str, Any]:
|
||||
"""执行 docker 命令(list-form,无 shell);返回 {exit_code, stdout, stderr}。"""
|
||||
proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout_sec)
|
||||
return {"exit_code": proc.returncode, "stdout": proc.stdout, "stderr": proc.stderr}
|
||||
|
||||
# ---- 构建 ----
|
||||
|
||||
def ensure_image(self, base_image: str, requirements: dict[str, Any],
|
||||
tag_suffix: str | None = None,
|
||||
timeout_sec: float | None = None) -> str:
|
||||
"""保证镜像存在:inspect 命中复用,否则 build + 登记。返回镜像名(带 tag)。
|
||||
|
||||
- 超时可配(默认 ImageConfig.build_timeout_sec)
|
||||
- build 失败:清理上下文目录 + 无 tag 中间镜像,抛 ImageBuildError(kind 可读)
|
||||
"""
|
||||
tag = tag_suffix or self.cfg.docker_tag_suffix
|
||||
timeout = timeout_sec or self.cfg.build_timeout_sec
|
||||
fingerprint = image_fingerprint(base_image, requirements, tag)
|
||||
name = image_name_for(fingerprint, tag)
|
||||
|
||||
hit = self._inspect_image(name)
|
||||
if hit.get("hit"):
|
||||
self._touch_registry(name, fingerprint, base_image, requirements)
|
||||
return name
|
||||
|
||||
pip, apt = _sorted_requirements(requirements)
|
||||
if not pip and not apt:
|
||||
# 无依赖:直接复用 base 镜像,不产生新镜像(仅记录映射)
|
||||
record = {"image": base_image, "fingerprint": fingerprint,
|
||||
"requirements": {"pip": pip, "apt": apt},
|
||||
"created_at": _now_iso()}
|
||||
self.registry().put(f"base:{base_image}:{tag}", record)
|
||||
return f"{base_image}:{tag}"
|
||||
|
||||
build_dir = self._build_context_dir(fingerprint)
|
||||
dockerfile = render_dockerfile(base_image, requirements, tag, self.cfg.pip_index)
|
||||
try:
|
||||
build_dir.mkdir(parents=True, exist_ok=True)
|
||||
(build_dir / "Dockerfile").write_text(dockerfile, encoding="utf-8")
|
||||
result = self._run_docker(
|
||||
["docker", "build", "-t", name, str(build_dir)], timeout
|
||||
)
|
||||
if result["exit_code"] != 0:
|
||||
raise ImageBuildError(self._classify(result), self._readable(result))
|
||||
self.registry().put(name, {
|
||||
"image": name,
|
||||
"fingerprint": fingerprint,
|
||||
"base_image": base_image,
|
||||
"requirements": {"pip": pip, "apt": apt},
|
||||
"created_at": _now_iso(),
|
||||
"refs": [],
|
||||
})
|
||||
return name
|
||||
except ImageBuildError:
|
||||
self._cleanup_failed_build(fingerprint, build_dir)
|
||||
raise
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
self._cleanup_failed_build(fingerprint, build_dir)
|
||||
raise ImageBuildError(
|
||||
"timeout",
|
||||
f"docker build 超时({timeout}s): {name};可增大 BIOFLOW_EXT_IMAGE_BUILD_TIMEOUT",
|
||||
) from exc
|
||||
except FileNotFoundError as exc:
|
||||
raise ImageBuildError("docker_missing", "docker 命令不可用(未安装或不在 PATH)") from exc
|
||||
|
||||
def _inspect_image(self, name: str) -> dict[str, Any]:
|
||||
try:
|
||||
result = self._run_docker(["docker", "inspect", name], timeout_sec=60.0)
|
||||
except FileNotFoundError:
|
||||
return {"hit": False, "error": "docker 命令不可用"}
|
||||
except subprocess.TimeoutExpired:
|
||||
return {"hit": False, "error": "docker inspect 超时"}
|
||||
if result["exit_code"] != 0:
|
||||
return {"hit": False, "error": _first_line(result["stderr"])}
|
||||
return {"hit": True}
|
||||
|
||||
def _touch_registry(self, name: str, fingerprint: str,
|
||||
base_image: str, requirements: dict[str, Any]) -> None:
|
||||
"""inspect 命中时确保 registry 有记录(registry 丢失/异地重建场景)。"""
|
||||
pip, apt = _sorted_requirements(requirements)
|
||||
record = self.registry().get(name)
|
||||
if record is None:
|
||||
self.registry().put(name, {
|
||||
"image": name, "fingerprint": fingerprint, "base_image": base_image,
|
||||
"requirements": {"pip": pip, "apt": apt},
|
||||
"created_at": _now_iso(), "refs": [],
|
||||
})
|
||||
|
||||
def _build_context_dir(self, fingerprint: str) -> Path:
|
||||
return Path(self.cfg.data_root) / "extensions" / "_images" / "build" / fingerprint
|
||||
|
||||
def _cleanup_failed_build(self, fingerprint: str, build_dir: Path) -> None:
|
||||
"""失败清理:构建上下文目录 + 无 tag 中间层镜像。
|
||||
|
||||
中间层定位:Dockerfile 内 LABEL plantai.ext.fingerprint 在失败构建时
|
||||
仍留在最后一个成功中间层(失败层无 tag 不可见),用 --filter label=
|
||||
精确 rmi 中间层,不误删共享层(base 层无该 LABEL)。
|
||||
"""
|
||||
import shutil
|
||||
|
||||
if build_dir.exists():
|
||||
shutil.rmtree(build_dir, ignore_errors=True)
|
||||
label = f"plantai.ext.fingerprint={fingerprint}"
|
||||
try:
|
||||
listed = self._run_docker(
|
||||
["docker", "images", "-q", "--filter", f"label={label}"],
|
||||
timeout=60.0,
|
||||
)
|
||||
if listed.get("exit_code") == 0:
|
||||
for line in str(listed.get("stdout") or "").splitlines():
|
||||
if line.strip():
|
||||
self._run_docker(["docker", "rmi", line.strip()], timeout_sec=120.0)
|
||||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
logger.info("extension image: 构建失败已清理 fingerprint=%s", fingerprint)
|
||||
|
||||
@staticmethod
|
||||
def _classify(result: dict[str, Any]) -> str:
|
||||
"""构建失败分类:区分 pip 不可满足 vs 网络 vs 其他。"""
|
||||
stderr = str(result.get("stderr") or "")
|
||||
stdout = str(result.get("stdout") or "")
|
||||
text = stderr + "\n" + stdout
|
||||
lower = text.lower()
|
||||
if "no matching distribution" in lower:
|
||||
return "pip_unavailable"
|
||||
if "not a valid identifier" in lower or "invalid value" in lower:
|
||||
return "pip_unavailable"
|
||||
if any(token in lower for token in ("temporary failure", "failed to fetch",
|
||||
"connection refused", "network is unreachable",
|
||||
"could not resolve")):
|
||||
return "network"
|
||||
if "unable to locate package" in lower or "404 not found" in lower:
|
||||
return "apt_unavailable"
|
||||
return "build_failed"
|
||||
|
||||
@staticmethod
|
||||
def _readable(result: dict[str, Any]) -> str:
|
||||
"""可读错误信息(给 AI 修复循环/用户;缺包提示在尾部)。"""
|
||||
stderr = str(result.get("stderr") or "")
|
||||
stdout = str(result.get("stdout") or "")
|
||||
text = stderr + "\n" + stdout
|
||||
lines = [line for line in text.splitlines() if line.strip()]
|
||||
if not lines:
|
||||
return "docker build 失败(无输出)"
|
||||
return "\n".join(lines[-20:])
|
||||
|
||||
# ---- 引用计数 / 回收 ----
|
||||
|
||||
def register_ref(self, image: str, extension_id: str) -> None:
|
||||
"""登记扩展 → 镜像引用(发布/运行扩展时调用,幂等)。"""
|
||||
record = self.registry().get(image)
|
||||
if record is None:
|
||||
return
|
||||
refs = set(record.get("refs") or [])
|
||||
if extension_id not in refs:
|
||||
refs.add(extension_id)
|
||||
record["refs"] = sorted(refs)
|
||||
self.registry().put(image, record)
|
||||
|
||||
def unregister_ref(self, image: str, extension_id: str) -> None:
|
||||
"""解除引用(扩展删除时调用,幂等)。"""
|
||||
record = self.registry().get(image)
|
||||
if record is None:
|
||||
return
|
||||
refs = set(record.get("refs") or [])
|
||||
refs.discard(extension_id)
|
||||
record["refs"] = sorted(refs)
|
||||
self.registry().put(image, record)
|
||||
|
||||
def collect_garbage(self, extension_ids: set[str],
|
||||
dry_run: bool = False) -> list[dict[str, Any]]:
|
||||
"""无引用镜像回收(registry 引用计数):保留被 extension_ids 引用的镜像。"""
|
||||
alive = self.registry().referenced(extension_ids)
|
||||
collected: list[dict[str, Any]] = []
|
||||
for name in list(self.registry().images):
|
||||
if name in alive or name.startswith("base:"):
|
||||
continue
|
||||
if dry_run:
|
||||
collected.append({"image": name, "removed": False, "dry_run": True})
|
||||
continue
|
||||
removed = False
|
||||
try:
|
||||
result = self._run_docker(["docker", "rmi", name], timeout_sec=300.0)
|
||||
removed = result["exit_code"] == 0
|
||||
if removed:
|
||||
self.registry().drop(name)
|
||||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
collected.append({"image": name, "removed": removed})
|
||||
return collected
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
return time.strftime("%Y-%m-%dT%H:%M:%S%z")
|
||||
|
||||
|
||||
def _first_line(text: str) -> str:
|
||||
for line in text.splitlines():
|
||||
if line.strip():
|
||||
return line.strip()
|
||||
return ""
|
||||
|
||||
|
||||
def _default_image_config() -> ImageConfig:
|
||||
"""从 app.config 推导(registry 根 = data_root/extensions/_images)。"""
|
||||
from ...config import get_settings
|
||||
|
||||
s = get_settings()
|
||||
timeout = getattr(s, "ext_image_build_timeout_sec", 0.0) or 1800.0
|
||||
return ImageConfig(
|
||||
data_root=s.data_root_abs,
|
||||
docker_tag_suffix=s.docker_image_tag_suffix,
|
||||
build_timeout_sec=timeout,
|
||||
)
|
||||
|
||||
|
||||
def build_extension_image_builder() -> ExtensionImageBuilder:
|
||||
"""应用启动:构造镜像生成器(懒加载 registry,不触碰 docker)。"""
|
||||
return ExtensionImageBuilder(_default_image_config())
|
||||
@@ -0,0 +1,511 @@
|
||||
"""app/services/bioflow/extension_llm.py —— 扩展 build/repair 的 AI 转换(DeepSeek 直连)。
|
||||
|
||||
- LLM 配置键:DOCGEN_LLM_*(沿用 docgen 直连模式);BIOFLOW_EXTENSION_*
|
||||
为可选别名(缺 DOCGEN 值时尝试)。
|
||||
- LLM 不可用/失败:一律回退规则版(_command_to_template/_detect_requirements/
|
||||
_infer_language 概念移植自 gfa_core/extension_stage/manager.py),**绝不抛 500**;
|
||||
回退无法完成才抛 LLMUnavailable(上层收敛 build_failed)。
|
||||
- 复用 app/services/bioflow/llm.py 的 complete_json(httpx 直连 + JSON 修复重试)。
|
||||
|
||||
服务层纪律:纯 stdlib + httpx,无 pydantic。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shlex
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from . import extension_spec as spec
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMUnavailable(RuntimeError):
|
||||
"""LLM 不可用且规则版无法完成转换(上层收敛 build_failed,不冒 500)。"""
|
||||
|
||||
|
||||
# 标准库白名单(_detect_requirements 移植:这些 import 不视为第三方依赖)
|
||||
STDLIB_MODULES = {
|
||||
"os", "sys", "re", "json", "csv", "argparse", "pathlib", "math",
|
||||
"statistics", "subprocess", "time", "datetime", "collections",
|
||||
"itertools", "functools", "glob", "gzip", "shutil", "logging", "random",
|
||||
"string", "typing", "io", "tempfile", "warnings", "hashlib", "urllib",
|
||||
}
|
||||
|
||||
CODE_SUFFIXES = {".py", ".r", ".sh", ".bash", ".command"}
|
||||
|
||||
_LLM_SYSTEM = (
|
||||
"You convert an already-runnable bioinformatics script package into a "
|
||||
"PlantAI custom extension. Reply with strict JSON only, no markdown fences."
|
||||
)
|
||||
|
||||
|
||||
def llm_config() -> dict[str, Any]:
|
||||
"""DOCGEN_LLM_* 优先,BIOFLOW_EXTENSION_* 别名兜底(base/key/model/timeout)。"""
|
||||
cfg: dict[str, Any] = {}
|
||||
for field, keys in (
|
||||
("base_url", ("DOCGEN_LLM_BASE_URL", "BIOFLOW_EXTENSION_BASE_URL")),
|
||||
("api_key", ("DOCGEN_LLM_API_KEY", "BIOFLOW_EXTENSION_API_KEY")),
|
||||
("model", ("DOCGEN_LLM_MODEL", "BIOFLOW_EXTENSION_MODEL")),
|
||||
("timeout", ("DOCGEN_LLM_TIMEOUT", "BIOFLOW_EXTENSION_TIMEOUT")),
|
||||
):
|
||||
value = ""
|
||||
for key in keys:
|
||||
raw = os.environ.get(key, "")
|
||||
if raw:
|
||||
value = raw
|
||||
break
|
||||
cfg[field] = value
|
||||
cfg["base_url"] = (cfg["base_url"] or "https://api.deepseek.com/v1").rstrip("/")
|
||||
cfg["model"] = cfg["model"] or "deepseek-chat"
|
||||
try:
|
||||
cfg["timeout"] = float(cfg["timeout"] or "300")
|
||||
except ValueError:
|
||||
cfg["timeout"] = 300.0
|
||||
return cfg
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 规则版探测(_detect_requirements/_infer_language 概念移植,纯函数)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def infer_language(path: str) -> str:
|
||||
suffix = Path(str(path)).suffix.lower()
|
||||
if suffix == ".py":
|
||||
return "python"
|
||||
if suffix == ".r":
|
||||
return "r"
|
||||
if suffix in {".sh", ".bash", ".command"}:
|
||||
return "shell"
|
||||
return suffix.lstrip(".") or "unknown"
|
||||
|
||||
|
||||
def detect_requirements(code_dir: str | Path) -> dict[str, list[str]]:
|
||||
"""规则版依赖探测:grep .py import / requirements.txt / .r library / .sh 命令。
|
||||
|
||||
返回 {"pip": [...], "apt": [...]}(契约 §2 requirements 形态)。
|
||||
"""
|
||||
pip: set[str] = set()
|
||||
apt: set[str] = set()
|
||||
root = Path(code_dir)
|
||||
if not root.is_dir():
|
||||
return {"pip": [], "apt": []}
|
||||
for path in sorted(root.rglob("*")):
|
||||
if not path.is_file():
|
||||
continue
|
||||
name = path.name.lower()
|
||||
if name in {"requirements.txt", "pip_requirements.txt"}:
|
||||
for line in path.read_text(encoding="utf-8", errors="ignore").splitlines():
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
dep = re.split(r"[=<>!~\s]", line)[0].strip()
|
||||
if dep:
|
||||
pip.add(dep)
|
||||
continue
|
||||
if path.suffix.lower() == ".py":
|
||||
text = path.read_text(encoding="utf-8", errors="ignore")
|
||||
for match in re.finditer(r"^\s*(?:import|from)\s+([A-Za-z_][A-Za-z0-9_]*)", text, re.M):
|
||||
mod = match.group(1)
|
||||
if mod not in STDLIB_MODULES:
|
||||
pip.add(mod)
|
||||
continue
|
||||
if path.suffix.lower() == ".r":
|
||||
text = path.read_text(encoding="utf-8", errors="ignore")
|
||||
for match in re.finditer(r"\b(?:library|require)\(([^)]+)\)", text):
|
||||
pkg = match.group(1).strip().strip("'\"")
|
||||
if pkg:
|
||||
pip.add(pkg)
|
||||
continue
|
||||
if path.suffix.lower() in {".sh", ".bash", ".command"}:
|
||||
text = path.read_text(encoding="utf-8", errors="ignore")
|
||||
apt.update(_shell_commands(text))
|
||||
return {"pip": sorted(pip), "apt": sorted(apt)}
|
||||
|
||||
|
||||
_SHELL_COMMAND_WHITELIST = {
|
||||
"python", "python3", "rscript", "perl", "bash", "sh", "awk", "sed", "grep",
|
||||
"cut", "sort", "uniq", "head", "tail", "cat", "echo", "cd", "mkdir", "rm",
|
||||
"cp", "mv", "ls", "wc", "tr", "paste", "join", "find", "xargs", "tar", "gzip",
|
||||
"gunzip", "zcat", "chmod", "export", "conda", "source", "time", "date", "ln",
|
||||
"set", "if", "then", "else", "fi", "for", "while", "do", "done", "exit",
|
||||
"return", "case", "esac", "unset", "shift", "read",
|
||||
}
|
||||
|
||||
|
||||
def _shell_commands(text: str) -> set[str]:
|
||||
""".sh 脚本中的可执行命令探测(_detect_shell_commands 移植裁剪)。
|
||||
|
||||
白名单外且形如命令名者计为 apt 依赖(如 fastqc/samtools)。
|
||||
"""
|
||||
commands: set[str] = set()
|
||||
for raw in text.splitlines():
|
||||
line = raw.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
try:
|
||||
parts = shlex.split(line, comments=True)
|
||||
except ValueError:
|
||||
continue
|
||||
if not parts:
|
||||
continue
|
||||
cmd = Path(parts[0]).name.lower()
|
||||
if not re.fullmatch(r"[a-z][a-z0-9_+.-]*", cmd or ""):
|
||||
continue
|
||||
if cmd in _SHELL_COMMAND_WHITELIST:
|
||||
continue
|
||||
commands.add(cmd)
|
||||
return commands
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 规则版命令模板化(LLM 失败回退;_command_to_template 概念移植)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def command_to_template(command: str, code_dir: str | Path, *, data_files: list[str] | None = None) -> dict[str, Any]:
|
||||
"""规则版:把 original_command 转为 command_template + 依赖探测 + 参数派生。
|
||||
|
||||
步骤:
|
||||
1. 已知脚本路径片段 → {code_dir}/<脚本>;
|
||||
2. 输入文件 token → 语义占位符(扩展名优先,其次 {input_file});
|
||||
3. results/ 类输出目录 → {output_dir};裸 --out/--output_dir 类 token 补 {output_dir};
|
||||
4. 无 {output_dir} 时追加 " > {output_dir}/output.txt"(output_dir 必含硬约束)。
|
||||
"""
|
||||
command = str(command or "").strip()
|
||||
code_dir = Path(code_dir)
|
||||
data_files = data_files or []
|
||||
|
||||
# 1. 代码脚本 → {code_dir}/<name>
|
||||
template = command
|
||||
for path in code_dir.rglob("*"):
|
||||
if not path.is_file() or path.suffix.lower() not in CODE_SUFFIXES:
|
||||
continue
|
||||
name = path.name
|
||||
template = re.sub(
|
||||
rf"(?<![A-Za-z0-9_./-]){re.escape(name)}\b",
|
||||
f"{{code_dir}}/{name}",
|
||||
template,
|
||||
)
|
||||
template = re.sub(r"\{code_dir\}/?\{code_dir\}", "{code_dir}", template)
|
||||
|
||||
# 2. 输入文件 → 语义占位符(扩展名优先)
|
||||
known_inputs = {Path(x).name for x in data_files}
|
||||
for fname in sorted(known_inputs, key=len, reverse=True):
|
||||
placeholder = _placeholder_for_filename(fname)
|
||||
template = re.sub(
|
||||
rf"(?<![A-Za-z0-9_./-]){re.escape(fname)}\b",
|
||||
placeholder,
|
||||
template,
|
||||
)
|
||||
|
||||
# 3. 输出目录 → {output_dir}
|
||||
for token in ("results/", "result/", "out_dir/", "output_dir/", "outputs/", "out/", "output/"):
|
||||
template = template.replace(token, "{output_dir}/")
|
||||
template = re.sub(r"\{output_dir\}+/+", "{output_dir}/", template)
|
||||
template = template.replace("{output_dir}//", "{output_dir}/")
|
||||
|
||||
# 4. 末尾裸目录 token(--out xxx 等)→ {output_dir}
|
||||
for token in ("--out", "--outdir", "--output", "--output_dir", "--output-dir", "-o"):
|
||||
template = re.sub(
|
||||
rf"({re.escape(token)})\s+([A-Za-z0-9_./-]+/?)",
|
||||
lambda m: f"{m.group(1)} {{output_dir}}",
|
||||
template,
|
||||
)
|
||||
|
||||
# 5. 无 {output_dir}:追加重定向兜底(output_dir 必含硬约束)
|
||||
if "{output_dir}" not in template:
|
||||
template = f"{template.rstrip()} > {{output_dir}}/output.txt"
|
||||
|
||||
# 6. 依赖探测 + 参数派生
|
||||
requirements = detect_requirements(code_dir)
|
||||
parameters: list[dict[str, Any]] = []
|
||||
seen: set[str] = set()
|
||||
for match in re.finditer(r"\{param_([A-Za-z_][A-Za-z0-9_]*)\}", template):
|
||||
name = match.group(1)
|
||||
if name not in seen:
|
||||
seen.add(name)
|
||||
parameters.append({"name": name, "type": "string", "required": True,
|
||||
"description": "", "default": None})
|
||||
return {
|
||||
"command_template": template,
|
||||
"requirements": requirements,
|
||||
"parameters": parameters,
|
||||
"source": "rules",
|
||||
}
|
||||
|
||||
|
||||
def _placeholder_for_filename(fname: str) -> str:
|
||||
lower = fname.lower()
|
||||
if lower.endswith(".hmm"):
|
||||
return "{input_hmm}"
|
||||
if lower.endswith((".pep.fasta", ".pep.fa", ".pep")):
|
||||
return "{input_pep_fasta}"
|
||||
if lower.endswith((".cds.fasta", ".cds.fa")):
|
||||
return "{input_cds_fasta}"
|
||||
if lower.endswith((".fasta", ".fa", ".faa", ".fna")):
|
||||
return "{input_fasta}"
|
||||
if lower.endswith((".id", ".ids", ".geneid", ".geneids", ".list", ".lst")):
|
||||
return "{input_id_file}"
|
||||
return "{input_file}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# build:LLM 转换(失败回退规则版)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _sanitize_conversion(raw: dict[str, Any]) -> dict[str, Any]:
|
||||
"""LLM 输出净化:字段形态归一 + 白名单校验(spec_v2 sanitize 概念裁剪)。"""
|
||||
out: dict[str, Any] = {}
|
||||
template = str(raw.get("command_template") or "").strip()
|
||||
spec.validate_command_template(template) # 不通过抛 ExtensionSpecError
|
||||
out["command_template"] = template
|
||||
|
||||
requirements: dict[str, Any] = {}
|
||||
raw_req = raw.get("requirements") or {}
|
||||
if isinstance(raw_req, dict):
|
||||
for key in ("pip", "apt"):
|
||||
values = raw_req.get(key) or []
|
||||
if isinstance(values, str):
|
||||
values = [values]
|
||||
requirements[key] = sorted({str(x).strip() for x in values if str(x).strip()})
|
||||
out["requirements"] = {"pip": requirements.get("pip") or [], "apt": requirements.get("apt") or []}
|
||||
|
||||
parameters: list[dict[str, Any]] = []
|
||||
raw_params = raw.get("parameters") or []
|
||||
if isinstance(raw_params, dict):
|
||||
raw_params = [{"name": k, **({"description": v} if isinstance(v, str) else v)} for k, v in raw_params.items()]
|
||||
if isinstance(raw_params, list):
|
||||
for item in raw_params:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
name = spec.safe_parameter_name(item.get("name") or "")
|
||||
if not name:
|
||||
continue
|
||||
ptype = str(item.get("type") or "string")
|
||||
if ptype not in {"string", "file", "int", "float", "bool"}:
|
||||
ptype = "string"
|
||||
parameters.append({
|
||||
"name": name,
|
||||
"type": ptype,
|
||||
"required": bool(item.get("required", False)),
|
||||
"description": str(item.get("description") or ""),
|
||||
"default": item.get("default"),
|
||||
})
|
||||
# 模板 {param_*} 未在 parameters 声明的一律补齐(契约完整性)
|
||||
out["parameters"] = spec.template_parameters(template, parameters)
|
||||
|
||||
input_roles = raw.get("input_roles")
|
||||
out["input_roles"] = input_roles if isinstance(input_roles, dict) else {}
|
||||
|
||||
nl_triggers = raw.get("nl_triggers") or []
|
||||
if isinstance(nl_triggers, str):
|
||||
nl_triggers = [nl_triggers]
|
||||
out["nl_triggers"] = [str(x).strip() for x in nl_triggers if str(x).strip()][:50]
|
||||
|
||||
out["analysis_type"] = spec.normalize_analysis_type(str(raw.get("analysis_type") or ""))
|
||||
runtime_image = str(raw.get("runtime_image") or raw.get("runtime_env") or "")
|
||||
out["runtime_image"] = spec.RUNTIME_IMAGE_MAP.get(runtime_image, runtime_image)
|
||||
out["source"] = str(raw.get("source") or "llm")
|
||||
return out
|
||||
|
||||
|
||||
async def _llm_build(ext: dict[str, Any], store: Any) -> dict[str, Any]:
|
||||
"""LLM 转换:把 original_command 转 command_template + 元数据。"""
|
||||
from ...config import get_settings
|
||||
|
||||
s = get_settings()
|
||||
cfg = llm_config()
|
||||
if not cfg["api_key"]:
|
||||
raise LLMUnavailable("扩展转换 LLM 未配置(DOCGEN_LLM_API_KEY 为空)")
|
||||
|
||||
from .llm import complete_json, LLMError, LLMTimeout, LLMUnavailable as PlanLLMUnavailable
|
||||
|
||||
raw_code = store.ext_dir(ext["ext_id"]) / "raw_code"
|
||||
code_files = _inventory(raw_code)
|
||||
data_files = _inventory(store.ext_dir(ext["ext_id"]) / "raw_test_data")
|
||||
user = _build_prompt(ext, code_files=code_files, data_files=data_files)
|
||||
try:
|
||||
payload = await complete_json(
|
||||
_LLM_SYSTEM,
|
||||
user,
|
||||
temperature=0,
|
||||
max_tokens=6000,
|
||||
timeout=float(s.docgen_llm_timeout),
|
||||
)
|
||||
except (PlanLLMUnavailable, LLMTimeout, LLMError, ValueError):
|
||||
raise
|
||||
return _sanitize_conversion(payload if isinstance(payload, dict) else {})
|
||||
|
||||
|
||||
def _build_prompt(ext: dict[str, Any], *, code_files: list[str], data_files: list[str]) -> str:
|
||||
lines = [
|
||||
"Convert the following extension upload into an extension spec JSON.",
|
||||
"The JSON must contain exactly these keys:",
|
||||
" command_template: the original_command rewritten with placeholders only from",
|
||||
" the whitelist: {code_dir} {input_dir} {input_file} {input_hmm} {input_fasta}",
|
||||
" {input_query_fasta} {input_reference_fasta} {input_pep_fasta} {input_cds_fasta}",
|
||||
" {input_id_file} {output_dir} {work_dir}, or {param_<name>} for CLI parameters.",
|
||||
" {output_dir} is REQUIRED (the extension must write outputs below it).",
|
||||
" requirements: {\"pip\": [...], \"apt\": [...]} - third-party dependencies only",
|
||||
" (stdlib imports are not dependencies).",
|
||||
" parameters: [{\"name\", \"type\": string|file|int|float|bool, \"required\",",
|
||||
" \"description\", \"default\"}] for each {param_*} used.",
|
||||
" input_roles: {\"<placeholder>\": \"<role>\"} for semantic placeholders used.",
|
||||
" nl_triggers: 3-6 short Chinese phrases users would say to invoke this",
|
||||
" extension in conversation (plus the display_name itself).",
|
||||
" analysis_type: one of " + ", ".join(sorted(spec.ANALYSIS_TYPES)) + ".",
|
||||
" runtime_image: one of " + ", ".join(sorted(spec.RUNTIME_IMAGE_MAP.values())) + ".",
|
||||
"",
|
||||
f"display_name: {ext.get('display_name')}",
|
||||
f"original_command: {ext.get('original_command')}",
|
||||
f"runtime_image (uploaded): {ext.get('runtime_image')}",
|
||||
"code files: " + (", ".join(code_files[:30]) or "(none)"),
|
||||
"test data files: " + (", ".join(data_files[:20]) or "(none)"),
|
||||
]
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _inventory(root: Path) -> list[str]:
|
||||
if not root.is_dir():
|
||||
return []
|
||||
return sorted(
|
||||
str(p.relative_to(root)) for p in root.rglob("*") if p.is_file()
|
||||
)[:200]
|
||||
|
||||
|
||||
def _rules_build(ext: dict[str, Any], store: Any) -> dict[str, Any]:
|
||||
"""规则版回退:command_to_template + 依赖探测 + analysis_type 推断。"""
|
||||
raw_code = store.ext_dir(ext["ext_id"]) / "raw_code"
|
||||
data_files = _inventory(store.ext_dir(ext["ext_id"]) / "raw_test_data")
|
||||
converted = command_to_template(
|
||||
ext.get("original_command") or "",
|
||||
raw_code,
|
||||
data_files=[Path(x).name for x in data_files],
|
||||
)
|
||||
spec.validate_command_template(converted["command_template"])
|
||||
converted["parameters"] = spec.template_parameters(converted["command_template"], converted.get("parameters"))
|
||||
converted["input_roles"] = {}
|
||||
converted["nl_triggers"] = _default_triggers(ext, converted["command_template"])
|
||||
converted["analysis_type"] = spec.infer_analysis_type(
|
||||
template=converted["command_template"],
|
||||
display_name=str(ext.get("display_name") or ""),
|
||||
requirements=converted["requirements"],
|
||||
)
|
||||
converted["runtime_image"] = str(ext.get("runtime_image") or spec.DEFAULT_RUNTIME_IMAGE)
|
||||
converted["source"] = "rules"
|
||||
return converted
|
||||
|
||||
|
||||
def _default_triggers(ext: dict[str, Any], template: str) -> list[str]:
|
||||
"""无 LLM 时的 nl_triggers 兜底:display_name + analysis_type 关键词。"""
|
||||
name = str(ext.get("display_name") or "").strip()
|
||||
triggers = [name] if name else []
|
||||
atype = spec.infer_analysis_type(
|
||||
template=template, display_name=name,
|
||||
requirements=ext.get("requirements") or {},
|
||||
)
|
||||
keywords = {
|
||||
"blast_search": "blast 比对",
|
||||
"hmmer_search": "hmm 搜索",
|
||||
"table_filter": "表格过滤",
|
||||
"sequence_extraction": "序列提取",
|
||||
"plot_render": "绘图",
|
||||
}
|
||||
if atype in keywords and keywords[atype] not in triggers:
|
||||
triggers.append(keywords[atype])
|
||||
return [t for t in triggers if t]
|
||||
|
||||
|
||||
async def build_conversion(ext: dict[str, Any], *, store: Any) -> dict[str, Any]:
|
||||
"""build 主入口:LLM 优先,失败回退规则版;两者皆失败抛 LLMUnavailable。"""
|
||||
try:
|
||||
return await _llm_build(ext, store)
|
||||
except Exception as exc: # noqa: BLE001 - 全部 LLM 失败路径(连接/超时/鉴权/坏 JSON)
|
||||
logger.warning("extensions: LLM build 失败,回退规则版: %s", exc)
|
||||
try:
|
||||
return _rules_build(ext, store)
|
||||
except Exception: # noqa: BLE001
|
||||
logger.exception("extensions: 规则版 build 也失败")
|
||||
raise LLMUnavailable("扩展转换失败:LLM 与规则版回退均无法完成") from None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# repair:失败上下文 → LLM 修复(失败回退规则版)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def repair_conversion(ext: dict[str, Any], *, failure_context: dict[str, Any], store: Any) -> dict[str, Any]:
|
||||
"""repair:把 trial 失败日志交给 LLM 修复 command_template(或 requirements)。
|
||||
|
||||
LLM 失败回退规则版(log_tail 里出现 ModuleNotFoundError 时把缺失模块并入
|
||||
requirements——B2「AI 修复循环兜底」的最小规则实现;镜像重建由 Dev-B)。
|
||||
"""
|
||||
cfg = llm_config()
|
||||
if cfg["api_key"]:
|
||||
from ...config import get_settings
|
||||
from .llm import complete_json
|
||||
|
||||
log_tail = str(failure_context.get("log_tail") or "")[-4000:]
|
||||
user = (
|
||||
"The extension trial run failed. Repair the extension spec JSON "
|
||||
"(same keys as build: command_template with whitelist placeholders "
|
||||
"and REQUIRED {output_dir}, requirements, parameters, input_roles, "
|
||||
"nl_triggers, analysis_type, runtime_image, plus a \"suggestion\" "
|
||||
"string describing the fix in Chinese).\n\n"
|
||||
f"current extension.json: {json.dumps(_repair_view(ext), ensure_ascii=False)}\n"
|
||||
f"trial log tail:\n{log_tail}\n"
|
||||
)
|
||||
try:
|
||||
payload = await complete_json(
|
||||
_LLM_SYSTEM,
|
||||
user,
|
||||
temperature=0,
|
||||
max_tokens=6000,
|
||||
timeout=float(get_settings().docgen_llm_timeout),
|
||||
)
|
||||
if isinstance(payload, dict) and payload.get("command_template"):
|
||||
converted = _sanitize_conversion(payload)
|
||||
converted["suggestion"] = str(payload.get("suggestion") or "")
|
||||
converted["source"] = "repair_llm"
|
||||
return converted
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning("extensions: LLM repair 失败,回退规则版: %s", exc)
|
||||
|
||||
# 规则版:ModuleNotFoundError → requirements 补缺;模板保持
|
||||
template = str(ext.get("command_template") or "")
|
||||
if not template:
|
||||
raise LLMUnavailable("repair 需要已有 command_template")
|
||||
log_tail = str(failure_context.get("log_tail") or "")
|
||||
added: list[str] = []
|
||||
for match in re.finditer(r"No module named ['\"]?([A-Za-z_][A-Za-z0-9_]*)", log_tail):
|
||||
added.append(match.group(1))
|
||||
requirements = dict(ext.get("requirements") or {"pip": [], "apt": []})
|
||||
requirements["pip"] = sorted({str(x) for x in (requirements.get("pip") or []) if str(x)} | set(added))
|
||||
if not added:
|
||||
raise LLMUnavailable("repair 失败:LLM 不可用且日志中未发现缺失依赖,无法规则修复")
|
||||
return {
|
||||
"command_template": template,
|
||||
"requirements": {"pip": requirements.get("pip") or [], "apt": requirements.get("apt") or []},
|
||||
"parameters": ext.get("parameters") or [],
|
||||
"input_roles": ext.get("input_roles") or {},
|
||||
"nl_triggers": ext.get("nl_triggers") or [],
|
||||
"analysis_type": ext.get("analysis_type") or spec.DEFAULT_ANALYSIS_TYPE,
|
||||
"runtime_image": ext.get("runtime_image") or spec.DEFAULT_RUNTIME_IMAGE,
|
||||
"suggestion": "已并入缺失依赖: " + ", ".join(added),
|
||||
"source": "repair_rules",
|
||||
}
|
||||
|
||||
|
||||
def _repair_view(ext: dict[str, Any]) -> dict[str, Any]:
|
||||
"""传给 LLM 的扩展视图(剔除易超长字段:source_files 明细等)。"""
|
||||
keys = ("ext_id", "display_name", "status", "command_template", "original_command",
|
||||
"runtime_image", "requirements", "parameters", "input_roles",
|
||||
"nl_triggers", "analysis_type", "repair_count")
|
||||
return {k: ext.get(k) for k in keys}
|
||||
@@ -0,0 +1,184 @@
|
||||
"""app/services/bioflow/extension_spec.py —— 扩展 command_template 白名单契约(spec_v2 移植)。
|
||||
|
||||
从 gfa_core/extension_stage/spec_v2.py 只移植契约常量与纯校验函数(旧模块只读):
|
||||
- 占位符白名单 12 个 + {param_<name>} 受控参数占位符;
|
||||
- 硬约束:{output_dir} 必含;白名单外占位符拒绝;
|
||||
- analysis_type 枚举(发布扩展的 intent 打分维度之一);
|
||||
- template → parameters 派生(未配置 parameters 时从 {param_*} 占位符生成)。
|
||||
|
||||
服务层纪律:纯 stdlib,无 pydantic、无 httpx。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
SPEC_VERSION = "extension_stage_v1"
|
||||
|
||||
# ---- 占位符白名单(契约 §2,spec_v2.PLACEHOLDERS 原表 12 个)----
|
||||
|
||||
PLACEHOLDERS = {
|
||||
"code_dir",
|
||||
"input_dir",
|
||||
"input_file",
|
||||
"input_hmm",
|
||||
"input_fasta",
|
||||
"input_query_fasta",
|
||||
"input_reference_fasta",
|
||||
"input_pep_fasta",
|
||||
"input_cds_fasta",
|
||||
"input_id_file",
|
||||
"output_dir",
|
||||
"work_dir",
|
||||
}
|
||||
|
||||
# 语义占位符(与旧版语义同表:绑定具体输入文件的角色化占位符)
|
||||
SEMANTIC_PLACEHOLDERS = {
|
||||
"input_hmm",
|
||||
"input_fasta",
|
||||
"input_query_fasta",
|
||||
"input_reference_fasta",
|
||||
"input_pep_fasta",
|
||||
"input_cds_fasta",
|
||||
"input_id_file",
|
||||
}
|
||||
|
||||
# ---- analysis_type 枚举(spec_v2.ANALYSIS_TYPES 移植)----
|
||||
|
||||
ANALYSIS_TYPES = {
|
||||
"sequence_extraction",
|
||||
"blast_search",
|
||||
"hmmer_search",
|
||||
"table_filter",
|
||||
"id_cleaning",
|
||||
"fasta_header_cleaning",
|
||||
"gff_processing",
|
||||
"plot_render",
|
||||
"enrichment_or_stats",
|
||||
"generic_cli",
|
||||
}
|
||||
|
||||
DEFAULT_ANALYSIS_TYPE = "generic_cli"
|
||||
|
||||
# 默认运行时镜像(B2 决策:工具链最全的 bioflow-gf)
|
||||
DEFAULT_RUNTIME_IMAGE = "bioflow-gf"
|
||||
|
||||
# 运行时环境 → 镜像映射(B2 定稿;上传/LLM 推荐时翻译旧 conda 环境名)
|
||||
RUNTIME_IMAGE_MAP = {
|
||||
"test4py311": "bioflow-gf",
|
||||
"gfpy312": "bioflow-gfpy",
|
||||
"rnaseq": "bioflow-rnaseq",
|
||||
"rnaR": "bioflow-rnaR",
|
||||
"ASAW": "bioflow-wgs",
|
||||
"interproscanpy311": "bioflow-interproscan",
|
||||
"panfamily": "bioflow-pan",
|
||||
"r_de_analysis": "bioflow-rde",
|
||||
}
|
||||
|
||||
_PARAM_RE = re.compile(r"\{([A-Za-z_][A-Za-z0-9_]*)\}")
|
||||
|
||||
|
||||
class ExtensionSpecError(RuntimeError):
|
||||
"""command_template 违反扩展契约(白名单/必含约束)。"""
|
||||
|
||||
|
||||
def used_placeholders(template: str) -> set[str]:
|
||||
"""模板中出现的全部 {占位符} 名(不含 {param_*} 前缀剥离)。"""
|
||||
return set(_PARAM_RE.findall(template or ""))
|
||||
|
||||
|
||||
def unsupported_placeholders(template: str) -> set[str]:
|
||||
"""白名单外且非 {param_*} 的占位符名集合。"""
|
||||
names = used_placeholders(template)
|
||||
return {name for name in names if name not in PLACEHOLDERS and not name.startswith("param_")}
|
||||
|
||||
|
||||
def param_placeholder_names(template: str) -> set[str]:
|
||||
"""{param_<name>} 占位符的参数名集合(剔除前缀)。"""
|
||||
return {
|
||||
name[len("param_"):]
|
||||
for name in used_placeholders(template)
|
||||
if name.startswith("param_") and len(name) > len("param_")
|
||||
}
|
||||
|
||||
|
||||
def validate_command_template(template: str) -> None:
|
||||
"""command_template 契约校验:{output_dir} 必含;白名单外占位符拒绝。
|
||||
|
||||
通过无异常;违反抛 ExtensionSpecError(spec_v2.normalize_spec_v2 同序)。
|
||||
"""
|
||||
if not (template or "").strip():
|
||||
raise ExtensionSpecError("command_template 不能为空")
|
||||
if "{output_dir}" not in template:
|
||||
raise ExtensionSpecError("command_template 必须包含 {output_dir}(输出契约硬约束)")
|
||||
unsupported = sorted(unsupported_placeholders(template))
|
||||
if unsupported:
|
||||
raise ExtensionSpecError("command_template 含白名单外占位符: " + ", ".join(unsupported))
|
||||
|
||||
|
||||
def normalize_analysis_type(value: str | None) -> str:
|
||||
"""analysis_type 白名单归一:非法值回落 generic_cli(漏判安全)。"""
|
||||
v = str(value or "").strip().lower()
|
||||
return v if v in ANALYSIS_TYPES else DEFAULT_ANALYSIS_TYPE
|
||||
|
||||
|
||||
def infer_analysis_type(*, template: str, display_name: str, requirements: dict[str, Any] | None = None) -> str:
|
||||
"""规则版 analysis_type 探测(spec_v2.infer_analysis_type 移植裁剪:模板/依赖/名称关键词)。"""
|
||||
t = (template or "").lower()
|
||||
pip_deps = " ".join(
|
||||
str(x).lower() for x in ((requirements or {}).get("pip") or [])
|
||||
)
|
||||
name = (display_name or "").lower()
|
||||
if "blastp" in t or "blastp" in pip_deps:
|
||||
return "blast_search"
|
||||
if any(k in t or k in pip_deps for k in ("hmmsearch", "hmmscan", "hmmbuild", "hmmer")):
|
||||
return "hmmer_search"
|
||||
if any(k in t for k in ("awk", "cut ", "grep ")):
|
||||
return "table_filter"
|
||||
if any(k in name for k in ("clean_id", "id_clean", "header_clean")):
|
||||
return "fasta_header_cleaning"
|
||||
if "seqtk" in t and "subseq" in t:
|
||||
return "sequence_extraction"
|
||||
if "matplotlib" in t or "plot" in name or any(k in name for k in ("画图", "绘图", "作图")):
|
||||
return "plot_render"
|
||||
return DEFAULT_ANALYSIS_TYPE
|
||||
|
||||
|
||||
def template_parameters(template: str, existing: list[dict[str, Any]] | None = None) -> list[dict[str, Any]]:
|
||||
"""由 {param_*} 占位符派生/补齐 parameters(已配置的定义优先,缺省的补齐为空定义)。"""
|
||||
merged: dict[str, dict[str, Any]] = {}
|
||||
for item in existing or []:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
name = safe_parameter_name(item.get("name") or "")
|
||||
if not name:
|
||||
continue
|
||||
merged[name] = {
|
||||
"name": name,
|
||||
"type": str(item.get("type") or "string"),
|
||||
"required": bool(item.get("required", False)),
|
||||
"description": str(item.get("description") or ""),
|
||||
"default": item.get("default"),
|
||||
}
|
||||
for name in sorted(param_placeholder_names(template)):
|
||||
merged.setdefault(name, {
|
||||
"name": name,
|
||||
"type": "string",
|
||||
"required": False,
|
||||
"description": "",
|
||||
"default": None,
|
||||
})
|
||||
return [merged[name] for name in sorted(merged)]
|
||||
|
||||
|
||||
def safe_parameter_name(value: Any) -> str:
|
||||
"""参数名净化(spec_v2.safe_parameter_name 同语义:非字母数字下划线→_,去 param_ 前缀)。"""
|
||||
text = re.sub(r"[^A-Za-z0-9_]+", "_", str(value or "").strip())
|
||||
text = re.sub(r"_+", "_", text).strip("_")
|
||||
if not text:
|
||||
return ""
|
||||
if text[0].isdigit():
|
||||
text = "p_" + text
|
||||
if text.startswith("param_"):
|
||||
text = text[len("param_"):]
|
||||
return text
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user