Compare commits
107
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d9abd0b86a | ||
|
|
282415fcb1 | ||
|
|
1ca2241a97 | ||
|
|
c78df3a6a6 | ||
|
|
ac07cb2008 | ||
|
|
d077a5b930 | ||
|
|
c4dc344a31 | ||
|
|
cf5ccb3a37 | ||
|
|
d99d5ef74a | ||
|
|
4add9f176b | ||
|
|
1319328b85 | ||
|
|
2d0e3a4858 | ||
|
|
ea0f2b3fad | ||
|
|
19fa12eb39 | ||
|
|
11337a788f | ||
|
|
55410b663a | ||
|
|
dddb70ff6f | ||
|
|
649e66f41d | ||
|
|
42aeb8fa4c | ||
|
|
c747ebd77a | ||
|
|
60375b04d7 | ||
|
|
1bd4aea5fd | ||
|
|
51532a30bc | ||
|
|
5a71e29cf7 | ||
|
|
d43d30ee3f | ||
|
|
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 |
+37
-1
@@ -1,4 +1,4 @@
|
||||
# 运行时产物与环境
|
||||
# 运行时产物与环境
|
||||
node_modules
|
||||
dist/
|
||||
.venv/
|
||||
@@ -21,3 +21,39 @@ 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
|
||||
|
||||
# review_service 运行时产物(不入库)
|
||||
review_service/review.db
|
||||
review_service/review.db-wal
|
||||
review_service/review.db-shm
|
||||
review_service/review.log
|
||||
review_service/data/reformat_jobs/
|
||||
# 期刊种子 journals.json 是源码:顶层 **/data/ 会误伤 review_service/data/,显式放行
|
||||
!review_service/data/
|
||||
review_service/tools/tectonic
|
||||
review_service/data/tectonic-cache/
|
||||
review_service/data/format_cache/
|
||||
|
||||
# 部署方案文档(含主机凭据,不入库)
|
||||
docs/DEPLOY_PAPER_176.md
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
# AGENTS.md —— 协作约束(科研审稿模块)
|
||||
|
||||
## 项目背景
|
||||
|
||||
本项目(PlantAI-rag-new)是植物基因组学文献知识工作台,拥有自研 RAG 与知识库
|
||||
(Qdrant + Neo4j + vstore)。2026-10 起新增「科研投稿与智能审稿助手」功能
|
||||
(审稿周期/期刊分析/期刊推荐/投稿前自审/审稿意见辅助/格式检查)。
|
||||
|
||||
## 审稿模块(review_service)约束 —— 必须遵守
|
||||
|
||||
1. **所有审稿相关新代码只允许写在 `review_service/` 目录下新增**。
|
||||
- 服务自身:FastAPI 独立进程,端口 8100(模式同 `paper_writer_service` 8097 / `bioflow_service` 8098)。
|
||||
- 环境变量前缀统一 `REVIEW_`;LLM 凭据复用 `DOCGEN_LLM_*` 裸键(回退链见 `review_service/app/services/llm.py`)。
|
||||
2. **BFF(`sdk/`)层只允许条件注册代理路由**:
|
||||
- 允许的全部改动 = `sdk/app/config.py` 增加 `review_enabled` / `review_service_url` 等开关键
|
||||
+ `sdk/app/main.py` 按开关 `include_router(review_proxy.router)`
|
||||
+ 新文件 `sdk/app/routers/review_proxy.py`(反代实现)。
|
||||
- **严禁**修改 sdk 下其他 router/service/db 的任何既有逻辑。
|
||||
3. **严禁修改现有 RAG 核心与既有服务代码**,包括但不限于:
|
||||
- `agent/`(qa 8091 / methodology 8092 / legacy-rag 8093)
|
||||
- `vstore/`(8094,审稿服务只经 HTTP 只读消费其 `/api/vstore/*` 接口)
|
||||
- `paper_writer_service/`、`bioflow_service/`、`bioflow_kernel/`
|
||||
- `ui/`(前端入口另行按 UI_SPEC 流程评审)
|
||||
- `knowledge/`、`kb-libraries/`、`tooling/`(依赖锁另议)
|
||||
4. **依赖纪律**:审稿服务只使用统一 venv(`tooling/requirements.lock`)中已有的库
|
||||
(fastapi / httpx / pydantic-settings 等)。如确需新增全局依赖,
|
||||
须先更新 `tooling/requirements.lock` 并经人工确认。
|
||||
5. **上游对接只读**:PDF 解析走 vstore 8094(`POST /api/vstore/ingest` +
|
||||
`GET/DELETE /api/vstore/documents/{id}`);LLM 走 DeepSeek OpenAI 兼容接口;
|
||||
后续期刊数据/期刊推荐等子功能同样以新 service + 新路由在 `review_service/` 内扩展。
|
||||
6. **format_check 例外**:format_check 允许用 pymupdf 本地解档获取版式元数据
|
||||
(字体/边距/行距/页数),因为 vstore 只回文本不返回版式信息。正文提取仍走 vstore。
|
||||
|
||||
## 目录速览(审稿模块)
|
||||
|
||||
```
|
||||
review_service/
|
||||
├── app/main.py # FastAPI 入口(/health + 挂载业务路由)
|
||||
├── app/config.py # pydantic-settings(REVIEW_ 前缀)
|
||||
├── app/routers/presubmit.py # POST /api/review/presubmit(自审稿报告)
|
||||
├── app/routers/format_check.py # POST /api/review/format-check(版式与规范检查)
|
||||
├── app/services/llm.py # DeepSeek httpx 封装(complete_text/complete_json)
|
||||
├── app/services/paper_parse.py # PDF → 纯文本(经 vstore 解析接口)
|
||||
├── app/services/format_check.py# 版式规则引擎(pymupdf/docx 本地解档 + 规则检查)
|
||||
├── restart.sh # 8100 重启脚本
|
||||
└── pytest.ini
|
||||
```
|
||||
|
||||
## 启动
|
||||
|
||||
```bash
|
||||
cd review_service && bash restart.sh # 默认 8100,venv 从 ../plantai-rag-new-venv 推导
|
||||
# BFF 开关:sdk/.env 加 REVIEW_ENABLED=1 后 bff_restart(反代 /api/review/*)
|
||||
```
|
||||
@@ -0,0 +1,12 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="zh-CN" data-theme="dark">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1, viewport-fit=cover">
|
||||
<title>蜂群中枢 · PlantAI 管理后台</title>
|
||||
</head>
|
||||
<body>
|
||||
<div id="app"></div>
|
||||
<script type="module" src="/src/main.js"></script>
|
||||
</body>
|
||||
</html>
|
||||
Generated
+1147
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,18 @@
|
||||
{
|
||||
"name": "plantai-admin",
|
||||
"version": "0.1.0",
|
||||
"private": true,
|
||||
"type": "module",
|
||||
"scripts": {
|
||||
"dev": "vite",
|
||||
"build": "vite build",
|
||||
"preview": "vite preview"
|
||||
},
|
||||
"dependencies": {
|
||||
"vue": "3.5.13"
|
||||
},
|
||||
"devDependencies": {
|
||||
"vite": "5.4.11",
|
||||
"@vitejs/plugin-vue": "5.2.1"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
<script setup>
|
||||
import { ref, onMounted } from 'vue';
|
||||
import { getToken, setToken, api } from './api/admin.js';
|
||||
import Login from './views/Login.vue';
|
||||
import Users from './views/Users.vue';
|
||||
import Logs from './views/Logs.vue';
|
||||
import HubSettings from './components/HubSettings.vue';
|
||||
|
||||
const authed = ref(false);
|
||||
const page = ref('overview');
|
||||
const user = ref(null);
|
||||
const stats = ref({ users: 0, total: 0, denied: 0, p95: null });
|
||||
const hub = ref({ hub_id: '', hub_name: '蜂群中枢' });
|
||||
const settingsOpen = ref(false);
|
||||
|
||||
const PAGES = { overview: '概览', users: '用户管理', logs: '日志查看' };
|
||||
|
||||
function onUnauthorized() { setToken(''); authed.value = false; }
|
||||
|
||||
async function loadHub() {
|
||||
try {
|
||||
const r = await api.settings();
|
||||
hub.value = r.settings;
|
||||
} catch (e) {
|
||||
if (e.status === 401) onUnauthorized();
|
||||
}
|
||||
}
|
||||
|
||||
async function loadStats() {
|
||||
try {
|
||||
const [u, l] = await Promise.all([api.users(), api.logs({ page: 1, page_size: 1 })]);
|
||||
stats.value.users = u.users.length;
|
||||
stats.value.total = l.total;
|
||||
stats.value.denied = (l.aggregates && l.aggregates.denied) || 0;
|
||||
stats.value.p95 = (l.aggregates && l.aggregates.p95_latency_ms) || null;
|
||||
} catch (e) {
|
||||
if (e.status === 401) onUnauthorized();
|
||||
}
|
||||
}
|
||||
|
||||
async function init() {
|
||||
if (!getToken()) { authed.value = false; return; }
|
||||
try {
|
||||
const r = await api.me();
|
||||
user.value = r.user;
|
||||
authed.value = true;
|
||||
loadHub();
|
||||
loadStats();
|
||||
} catch (e) {
|
||||
setToken('');
|
||||
authed.value = false;
|
||||
}
|
||||
}
|
||||
|
||||
function onLogin(u) { user.value = u; authed.value = true; page.value = 'overview'; loadHub(); loadStats(); }
|
||||
async function onLogout() { try { await api.logout(); } catch {} setToken(''); authed.value = false; }
|
||||
function onSettingsSaved(s) { hub.value = s; settingsOpen.value = false; }
|
||||
|
||||
onMounted(init);
|
||||
</script>
|
||||
|
||||
<template>
|
||||
<Login v-if="!authed" @login="onLogin" />
|
||||
<div v-else class="shell">
|
||||
<aside class="sidebar">
|
||||
<div class="side-brand">
|
||||
<div class="hex">蜂</div>
|
||||
<div><div class="brand-name">{{ hub.hub_name }}</div><div class="brand-sub">{{ hub.hub_id.slice(0, 8) }}</div></div>
|
||||
</div>
|
||||
<nav class="nav">
|
||||
<div class="nav-item" :class="{ active: page === 'overview' }" @click="page = 'overview'"><span>◈</span> 概览</div>
|
||||
<div class="nav-item" :class="{ active: page === 'users' }" @click="page = 'users'"><span>◉</span> 用户管理</div>
|
||||
<div class="nav-item" :class="{ active: page === 'logs' }" @click="page = 'logs'"><span>≣</span> 日志查看</div>
|
||||
</nav>
|
||||
<div class="side-foot">
|
||||
<div class="who">
|
||||
<div class="avatar">蜂</div>
|
||||
<div><div class="who-name">{{ user?.username }}</div><div class="who-role">系统管理员</div></div>
|
||||
<span class="logout" @click="onLogout">退出</span>
|
||||
</div>
|
||||
</div>
|
||||
</aside>
|
||||
<main class="main">
|
||||
<div class="topbar">
|
||||
<h1>{{ PAGES[page] }}</h1>
|
||||
<span style="color: var(--text-3); font-size: 13px">{{ hub.hub_name }}</span>
|
||||
<button class="btn btn-ghost" style="margin-left: auto" @click="settingsOpen = true">⚙ 设置</button>
|
||||
</div>
|
||||
<div class="content">
|
||||
<section v-if="page === 'overview'">
|
||||
<div class="stats">
|
||||
<div class="stat"><div class="k">{{ stats.users }}</div><div class="l">注册用户</div></div>
|
||||
<div class="stat"><div class="k amber">{{ stats.total }}</div><div class="l">累计请求</div></div>
|
||||
<div class="stat"><div class="k">{{ stats.denied }}</div><div class="l">越权 / 失败</div></div>
|
||||
<div class="stat"><div class="k">{{ stats.p95 ?? '—' }}</div><div class="l">P95 耗时 ms</div></div>
|
||||
</div>
|
||||
</section>
|
||||
<Users v-else-if="page === 'users'" @unauthorized="onUnauthorized" />
|
||||
<Logs v-else @unauthorized="onUnauthorized" />
|
||||
</div>
|
||||
</main>
|
||||
<HubSettings :open="settingsOpen" :hub="hub" @close="settingsOpen = false" @saved="onSettingsSaved" />
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<style>
|
||||
.shell { display: flex; height: 100vh; }
|
||||
.sidebar { width: 218px; flex: none; background: var(--bg-2); border-right: 1px solid var(--border); display: flex; flex-direction: column; }
|
||||
.side-brand { display: flex; align-items: center; gap: 10px; padding: 20px 18px 18px; }
|
||||
.hex { width: 30px; height: 30px; flex: none; background: linear-gradient(135deg, #f2b95a, #d98f22); clip-path: polygon(50% 0, 100% 25%, 100% 75%, 50% 100%, 0 75%, 0 25%); display: grid; place-items: center; color: #1a1206; font-weight: 700; font-size: 12px; }
|
||||
.brand-name { font-size: 15px; font-weight: 600; }
|
||||
.brand-sub { font-size: 11px; color: var(--text-3); }
|
||||
.nav { padding: 6px 10px; flex: 1; }
|
||||
.nav-item { display: flex; align-items: center; gap: 9px; padding: 9px 12px; margin-bottom: 3px; border-radius: var(--radius-sm); color: var(--text-2); font-size: 13.5px; cursor: pointer; }
|
||||
.nav-item:hover { background: var(--surface); color: var(--text); }
|
||||
.nav-item.active { background: var(--accent-soft); color: var(--accent); }
|
||||
.side-foot { padding: 14px 18px; border-top: 1px solid var(--border); }
|
||||
.who { display: flex; align-items: center; gap: 10px; }
|
||||
.avatar { width: 30px; height: 30px; border-radius: 50%; flex: none; background: linear-gradient(135deg, #3a4a5e, #243146); display: grid; place-items: center; font-size: 12px; color: var(--text); }
|
||||
.who-name { font-size: 13px; }
|
||||
.who-role { font-size: 11px; color: var(--text-3); }
|
||||
.logout { margin-left: auto; color: var(--text-3); cursor: pointer; font-size: 12px; }
|
||||
.logout:hover { color: var(--red); }
|
||||
.main { flex: 1; overflow-y: auto; display: flex; flex-direction: column; }
|
||||
.topbar { display: flex; align-items: center; gap: 14px; padding: 18px 28px; border-bottom: 1px solid var(--border); position: sticky; top: 0; background: rgba(12, 14, 20, 0.9); backdrop-filter: blur(6px); z-index: 5; }
|
||||
.topbar h1 { font-size: 17px; margin: 0; font-weight: 600; }
|
||||
.content { padding: 26px 28px 40px; }
|
||||
.stats { display: grid; grid-template-columns: repeat(4, 1fr); gap: 16px; }
|
||||
.stat { background: var(--surface); border: 1px solid var(--border); border-radius: var(--radius); padding: 18px 20px; }
|
||||
.stat .k { font-size: 26px; font-weight: 600; letter-spacing: -0.5px; }
|
||||
.stat .k.amber { color: var(--accent); }
|
||||
.stat .l { color: var(--text-2); font-size: 12.5px; margin-top: 3px; }
|
||||
.toolbar { display: flex; align-items: center; gap: 10px; margin-bottom: 16px; }
|
||||
.toolbar .spacer { flex: 1; }
|
||||
.search { display: flex; align-items: center; gap: 8px; background: var(--surface); border: 1px solid var(--border); border-radius: var(--radius-sm); padding: 8px 12px; width: 240px; }
|
||||
.search input { background: none; border: none; outline: none; color: var(--text); font-size: 13px; width: 100%; }
|
||||
select { background: var(--surface); border: 1px solid var(--border); color: var(--text); border-radius: var(--radius-sm); padding: 8px 10px; font-size: 13px; outline: none; }
|
||||
.pag { display: flex; align-items: center; gap: 8px; padding: 12px 18px; border-top: 1px solid var(--border); color: var(--text-3); font-size: 12.5px; }
|
||||
.pag .spacer { flex: 1; }
|
||||
.pag button { background: var(--surface-2); border: 1px solid var(--border); color: var(--text-2); border-radius: 6px; padding: 5px 11px; font-size: 12.5px; }
|
||||
.pag button.on { color: var(--accent); border-color: var(--accent); }
|
||||
.pag button:disabled { opacity: 0.4; cursor: default; }
|
||||
</style>
|
||||
@@ -0,0 +1,48 @@
|
||||
const TOKEN_KEY = 'pa-admin-token';
|
||||
let token = localStorage.getItem(TOKEN_KEY) || '';
|
||||
|
||||
export function getToken() { return token; }
|
||||
export function setToken(t) {
|
||||
token = t || '';
|
||||
if (t) localStorage.setItem(TOKEN_KEY, t);
|
||||
else localStorage.removeItem(TOKEN_KEY);
|
||||
}
|
||||
|
||||
async function request(path, { method = 'GET', body } = {}) {
|
||||
const res = await fetch('/api/admin' + path, {
|
||||
method,
|
||||
headers: {
|
||||
...(body ? { 'Content-Type': 'application/json' } : {}),
|
||||
...(token ? { Authorization: 'Bearer ' + token } : {}),
|
||||
},
|
||||
body: body ? JSON.stringify(body) : undefined,
|
||||
});
|
||||
let data = null;
|
||||
try { data = await res.json(); } catch { data = null; }
|
||||
if (!res.ok) {
|
||||
const err = new Error((data && data.error && data.error.message) || ('请求失败 HTTP ' + res.status));
|
||||
err.status = res.status;
|
||||
err.code = data && data.error && data.error.code;
|
||||
throw err;
|
||||
}
|
||||
return data;
|
||||
}
|
||||
|
||||
export const api = {
|
||||
login: (username, password) => request('/login', { method: 'POST', body: { username, password } }),
|
||||
logout: () => request('/logout', { method: 'POST' }),
|
||||
me: () => request('/me'),
|
||||
units: () => request('/units'),
|
||||
users: () => request('/users'),
|
||||
createUser: (body) => request('/users', { method: 'POST', body }),
|
||||
updateUser: (id, body) => request('/users/' + id, { method: 'PATCH', body }),
|
||||
deleteUser: (id) => request('/users/' + id, { method: 'DELETE' }),
|
||||
logs: (params = {}) => {
|
||||
const q = new URLSearchParams();
|
||||
Object.entries(params).forEach(([k, v]) => { if (v !== undefined && v !== null && v !== '') q.set(k, v); });
|
||||
const qs = q.toString();
|
||||
return request('/logs' + (qs ? '?' + qs : ''));
|
||||
},
|
||||
settings: () => request('/settings'),
|
||||
updateSettings: (body) => request('/settings', { method: 'PUT', body }),
|
||||
};
|
||||
@@ -0,0 +1,57 @@
|
||||
<script setup>
|
||||
import { ref, watch } from 'vue';
|
||||
import { api } from '../api/admin.js';
|
||||
|
||||
const props = defineProps({ open: Boolean, hub: Object });
|
||||
const emit = defineEmits(['close', 'saved']);
|
||||
|
||||
const hubId = ref('');
|
||||
const hubName = ref('');
|
||||
const busy = ref(false);
|
||||
const error = ref('');
|
||||
|
||||
watch(() => props.open, (open) => {
|
||||
if (open) {
|
||||
hubId.value = (props.hub && props.hub.hub_id) || '';
|
||||
hubName.value = (props.hub && props.hub.hub_name) || '';
|
||||
error.value = '';
|
||||
busy.value = false;
|
||||
}
|
||||
});
|
||||
|
||||
async function save() {
|
||||
if (busy.value) return;
|
||||
if (!hubName.value.trim()) { error.value = '中枢名称不能为空'; return; }
|
||||
busy.value = true;
|
||||
error.value = '';
|
||||
try {
|
||||
const r = await api.updateSettings({ hub_id: hubId.value.trim(), hub_name: hubName.value.trim() });
|
||||
emit('saved', r.settings);
|
||||
} catch (e) {
|
||||
error.value = e.message || '保存失败';
|
||||
} finally {
|
||||
busy.value = false;
|
||||
}
|
||||
}
|
||||
</script>
|
||||
|
||||
<template>
|
||||
<div v-if="open" class="overlay" @click.self="emit('close')">
|
||||
<div class="modal">
|
||||
<div class="modal-head">中枢设置</div>
|
||||
<div class="modal-body">
|
||||
<div class="field"><label>中枢名称</label><input v-model="hubName" placeholder="例如:湘湖实验室中枢"></div>
|
||||
<div class="field" style="margin-bottom: 0"><label>中枢 ID</label><input v-model="hubId" class="mono" placeholder="唯一标识,留空则由系统生成"></div>
|
||||
<p v-if="error" class="err">{{ error }}</p>
|
||||
</div>
|
||||
<div class="modal-foot">
|
||||
<button class="btn btn-ghost" @click="emit('close')">取消</button>
|
||||
<button class="btn btn-primary" :disabled="busy" @click="save">{{ busy ? '保存中…' : '保存' }}</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<style scoped>
|
||||
.err { color: var(--red); font-size: 12px; margin: 12px 0 0; }
|
||||
</style>
|
||||
@@ -0,0 +1,101 @@
|
||||
<script setup>
|
||||
import { ref, watch } from 'vue';
|
||||
import { api } from '../api/admin.js';
|
||||
|
||||
const props = defineProps({
|
||||
open: Boolean,
|
||||
mode: { type: String, default: 'create' },
|
||||
user: Object,
|
||||
});
|
||||
const emit = defineEmits(['close', 'saved']);
|
||||
|
||||
const UNITS = [
|
||||
['qa', '问答'], ['kg', '知识图谱'], ['extraction', '文献抽取'], ['docgen', '文档生成'],
|
||||
['figures', '图片图库'], ['private_docs', '私有资料'], ['bioflow', '生信分析'],
|
||||
['paper', '论文撰写'], ['legacy', '旧版接口'],
|
||||
];
|
||||
|
||||
const username = ref('');
|
||||
const password = ref('');
|
||||
const role = ref('researcher');
|
||||
const authorizations = ref([]);
|
||||
const busy = ref(false);
|
||||
const error = ref('');
|
||||
|
||||
watch(() => props.open, (open) => {
|
||||
if (open) {
|
||||
error.value = '';
|
||||
busy.value = false;
|
||||
if (props.mode === 'edit' && props.user) {
|
||||
username.value = props.user.username;
|
||||
password.value = '';
|
||||
role.value = props.user.role === 'admin' ? 'admin' : (props.user.role === 'guest' ? 'guest' : 'researcher');
|
||||
authorizations.value = [...(props.user.authorizations || [])];
|
||||
} else {
|
||||
username.value = '';
|
||||
password.value = '';
|
||||
role.value = 'researcher';
|
||||
authorizations.value = ['qa'];
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
async function save() {
|
||||
if (busy.value) return;
|
||||
if (props.mode === 'create' && (!username.value.trim() || !password.value)) {
|
||||
error.value = '用户名和密码必填';
|
||||
return;
|
||||
}
|
||||
busy.value = true;
|
||||
error.value = '';
|
||||
try {
|
||||
if (props.mode === 'create') {
|
||||
await api.createUser({ username: username.value.trim(), password: password.value, role: role.value, authorizations: authorizations.value });
|
||||
} else {
|
||||
const body = { role: role.value, authorizations: authorizations.value };
|
||||
if (password.value) body.password = password.value;
|
||||
await api.updateUser(props.user.id, body);
|
||||
}
|
||||
emit('saved');
|
||||
} catch (e) {
|
||||
error.value = e.message || '保存失败';
|
||||
} finally {
|
||||
busy.value = false;
|
||||
}
|
||||
}
|
||||
</script>
|
||||
|
||||
<template>
|
||||
<div v-if="open" class="overlay" @click.self="emit('close')">
|
||||
<div class="modal">
|
||||
<div class="modal-head">{{ mode === 'create' ? '新建用户' : '编辑用户' }}</div>
|
||||
<div class="modal-body">
|
||||
<div class="field"><label>用户名</label><input v-model="username" :disabled="mode === 'edit'"></div>
|
||||
<div class="field"><label>{{ mode === 'create' ? '初始密码' : '重置密码(留空不改)' }}</label><input v-model="password" type="password"></div>
|
||||
<div class="field">
|
||||
<label>角色</label>
|
||||
<div class="seg">
|
||||
<button :class="{ on: role === 'researcher' }" @click="role = 'researcher'">研究员</button>
|
||||
<button :class="{ on: role === 'guest' }" @click="role = 'guest'">访客</button>
|
||||
<button :class="{ on: role === 'admin' }" @click="role = 'admin'">管理员</button>
|
||||
</div>
|
||||
</div>
|
||||
<div class="field" style="margin-bottom: 0">
|
||||
<label>授权清单</label>
|
||||
<div class="checks">
|
||||
<label v-for="[k, label] in UNITS" :key="k" class="check"><input type="checkbox" :value="k" v-model="authorizations"> {{ label }}</label>
|
||||
</div>
|
||||
</div>
|
||||
<p v-if="error" class="err">{{ error }}</p>
|
||||
</div>
|
||||
<div class="modal-foot">
|
||||
<button class="btn btn-ghost" @click="emit('close')">取消</button>
|
||||
<button class="btn btn-primary" :disabled="busy" @click="save">{{ busy ? '保存中…' : (mode === 'create' ? '创建用户' : '保存') }}</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<style scoped>
|
||||
.err { color: var(--red); font-size: 12px; margin: 12px 0 0; }
|
||||
</style>
|
||||
@@ -0,0 +1,5 @@
|
||||
import { createApp } from 'vue';
|
||||
import App from './App.vue';
|
||||
import './style.css';
|
||||
|
||||
createApp(App).mount('#app');
|
||||
@@ -0,0 +1,96 @@
|
||||
:root {
|
||||
--bg: #0c0e14;
|
||||
--bg-2: #10131b;
|
||||
--surface: #151924;
|
||||
--surface-2: #1a1f2c;
|
||||
--border: #232a38;
|
||||
--border-2: #2e3648;
|
||||
--text: #e9edf5;
|
||||
--text-2: #9aa4b6;
|
||||
--text-3: #5f6a7d;
|
||||
--accent: #e8a33d;
|
||||
--accent-ink: #14100a;
|
||||
--accent-soft: rgba(232, 163, 61, 0.14);
|
||||
--green: #3ecf8e;
|
||||
--green-soft: rgba(62, 207, 142, 0.14);
|
||||
--red: #ef6b6b;
|
||||
--red-soft: rgba(239, 107, 107, 0.14);
|
||||
--blue: #6ea8ff;
|
||||
--radius: 10px;
|
||||
--radius-sm: 7px;
|
||||
}
|
||||
|
||||
* { box-sizing: border-box; }
|
||||
html, body, #app { margin: 0; height: 100%; }
|
||||
body {
|
||||
background: var(--bg);
|
||||
color: var(--text);
|
||||
font-family: -apple-system, BlinkMacSystemFont, "PingFang SC", "Microsoft YaHei", "Segoe UI", Roboto, sans-serif;
|
||||
font-size: 14px;
|
||||
line-height: 1.5;
|
||||
-webkit-font-smoothing: antialiased;
|
||||
}
|
||||
.mono { font-family: ui-monospace, "SF Mono", "JetBrains Mono", Menlo, Consolas, monospace; }
|
||||
button { font-family: inherit; cursor: pointer; }
|
||||
|
||||
.btn {
|
||||
display: inline-flex; align-items: center; justify-content: center; gap: 7px;
|
||||
padding: 9px 15px; font-size: 13.5px; border-radius: var(--radius-sm);
|
||||
border: 1px solid transparent; transition: filter .15s; background: transparent;
|
||||
}
|
||||
.btn:focus-visible { outline: 2px solid var(--accent); outline-offset: 2px; }
|
||||
.btn-primary { background: var(--accent); color: var(--accent-ink); font-weight: 600; }
|
||||
.btn-primary:hover { filter: brightness(1.06); }
|
||||
.btn-ghost { border-color: var(--border-2); color: var(--text-2); }
|
||||
.btn-ghost:hover { color: var(--text); border-color: var(--text-3); }
|
||||
.btn-danger { background: var(--red-soft); color: var(--red); }
|
||||
.btn-block { width: 100%; }
|
||||
|
||||
.field { margin-bottom: 15px; }
|
||||
.field label { display: block; font-size: 12px; color: var(--text-2); margin-bottom: 7px; }
|
||||
.field input, .field select {
|
||||
width: 100%; padding: 10px 12px; font-size: 14px; color: var(--text);
|
||||
background: var(--bg-2); border: 1px solid var(--border); border-radius: var(--radius-sm); outline: none;
|
||||
}
|
||||
.field input:focus, .field select:focus { border-color: var(--accent); box-shadow: 0 0 0 3px var(--accent-soft); }
|
||||
|
||||
.chip {
|
||||
display: inline-block; padding: 3px 9px; border-radius: 20px; font-size: 11.5px;
|
||||
border: 1px solid var(--border-2); color: var(--text-2); margin: 1.5px 3px 1.5px 0;
|
||||
}
|
||||
.chip.amber { border-color: transparent; background: var(--accent-soft); color: var(--accent); }
|
||||
.badge { display: inline-block; padding: 2px 8px; border-radius: 5px; font-size: 11.5px; font-family: ui-monospace, Menlo, monospace; }
|
||||
.badge.green { background: var(--green-soft); color: var(--green); }
|
||||
.badge.red { background: var(--red-soft); color: var(--red); }
|
||||
.badge.amber { background: var(--accent-soft); color: var(--accent); }
|
||||
.badge.muted { background: var(--surface-2); color: var(--text-3); }
|
||||
|
||||
.panel { background: var(--surface); border: 1px solid var(--border); border-radius: var(--radius); }
|
||||
table { width: 100%; border-collapse: collapse; }
|
||||
th { text-align: left; font-size: 11.5px; font-weight: 500; color: var(--text-3); padding: 10px 18px; border-bottom: 1px solid var(--border); }
|
||||
td { padding: 13px 18px; border-bottom: 1px solid var(--border); font-size: 13.5px; }
|
||||
tr:last-child td { border-bottom: none; }
|
||||
tbody tr:hover { background: rgba(255,255,255,0.015); }
|
||||
.cell-id { color: var(--text-2); }
|
||||
.row-actions { display: flex; gap: 6px; }
|
||||
.row-actions .act { color: var(--blue); cursor: pointer; font-size: 13px; padding: 3px 6px; border-radius: 5px; }
|
||||
.row-actions .act:hover { background: rgba(110,168,255,0.1); }
|
||||
.row-actions .act.danger { color: var(--red); }
|
||||
.row-actions .act.danger:hover { background: var(--red-soft); }
|
||||
|
||||
.role { font-weight: 600; }
|
||||
.role.admin { color: var(--accent); }
|
||||
.role.researcher { color: var(--blue); }
|
||||
.role.guest { color: var(--text-2); }
|
||||
|
||||
.overlay { position: fixed; inset: 0; background: rgba(0,0,0,0.55); display: flex; align-items: center; justify-content: center; z-index: 50; }
|
||||
.modal { width: 460px; max-width: calc(100vw - 32px); background: var(--surface); border: 1px solid var(--border-2); border-radius: 14px; }
|
||||
.modal-head { padding: 18px 22px 0; font-size: 15px; font-weight: 600; }
|
||||
.modal-body { padding: 18px 22px; }
|
||||
.modal-foot { padding: 14px 22px 20px; display: flex; justify-content: flex-end; gap: 10px; border-top: 1px solid var(--border); }
|
||||
.checks { display: grid; grid-template-columns: 1fr 1fr; gap: 9px 14px; margin-top: 8px; }
|
||||
.check { display: flex; align-items: center; gap: 8px; font-size: 13px; color: var(--text-2); }
|
||||
.check input { accent-color: var(--accent); width: 15px; height: 15px; }
|
||||
.seg { display: flex; gap: 6px; }
|
||||
.seg button { background: transparent; border: 1px solid var(--border-2); color: var(--text-2); border-radius: 7px; padding: 7px 14px; font-size: 13px; }
|
||||
.seg button.on { background: var(--accent-soft); border-color: transparent; color: var(--accent); }
|
||||
@@ -0,0 +1,53 @@
|
||||
<script setup>
|
||||
import { ref } from 'vue';
|
||||
import { api, setToken } from '../api/admin.js';
|
||||
|
||||
const emit = defineEmits(['login']);
|
||||
const username = ref('');
|
||||
const password = ref('');
|
||||
const error = ref('');
|
||||
const busy = ref(false);
|
||||
|
||||
async function submit() {
|
||||
if (busy.value) return;
|
||||
if (!username.value || !password.value) { error.value = '请输入用户名和密码'; return; }
|
||||
busy.value = true;
|
||||
error.value = '';
|
||||
try {
|
||||
const r = await api.login(username.value, password.value);
|
||||
setToken(r.token);
|
||||
emit('login', r.user);
|
||||
} catch (e) {
|
||||
error.value = e.message || '登录失败';
|
||||
} finally {
|
||||
busy.value = false;
|
||||
}
|
||||
}
|
||||
</script>
|
||||
|
||||
<template>
|
||||
<div class="login">
|
||||
<div class="card">
|
||||
<div class="brand">
|
||||
<div class="hex">蜂</div>
|
||||
<div><div class="brand-name">蜂群中枢</div><div class="brand-sub">PlantAI 管理后台</div></div>
|
||||
</div>
|
||||
<div class="field"><label>用户名</label><input v-model="username" autocomplete="username" @keyup.enter="submit"></div>
|
||||
<div class="field"><label>密码</label><input v-model="password" type="password" autocomplete="current-password" @keyup.enter="submit"></div>
|
||||
<button class="btn btn-primary btn-block" :disabled="busy" @click="submit">{{ busy ? '登录中…' : '登录' }}</button>
|
||||
<p v-if="error" class="err">{{ error }}</p>
|
||||
<div class="hint">仅限中枢管理员。登录与所有操作均会留痕。</div>
|
||||
</div>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<style scoped>
|
||||
.login { height: 100vh; display: flex; align-items: center; justify-content: center; background: radial-gradient(1200px 600px at 70% -10%, rgba(232, 163, 61, 0.08), transparent 60%), var(--bg); }
|
||||
.card { width: 360px; background: var(--surface); border: 1px solid var(--border); border-radius: 14px; padding: 34px 32px 30px; }
|
||||
.brand { display: flex; align-items: center; gap: 11px; margin-bottom: 26px; }
|
||||
.hex { width: 34px; height: 34px; flex: none; background: linear-gradient(135deg, #f2b95a, #d98f22); clip-path: polygon(50% 0, 100% 25%, 100% 75%, 50% 100%, 0 75%, 0 25%); display: grid; place-items: center; color: #1a1206; font-weight: 700; font-size: 14px; }
|
||||
.brand-name { font-size: 16px; font-weight: 600; }
|
||||
.brand-sub { color: var(--text-3); font-size: 12px; }
|
||||
.err { color: var(--red); font-size: 12px; margin: 10px 0 0; }
|
||||
.hint { margin-top: 20px; padding-top: 16px; border-top: 1px solid var(--border); font-size: 12px; color: var(--text-3); }
|
||||
</style>
|
||||
@@ -0,0 +1,102 @@
|
||||
<script setup>
|
||||
import { ref, onMounted } from 'vue';
|
||||
import { api } from '../api/admin.js';
|
||||
|
||||
const emit = defineEmits(['unauthorized']);
|
||||
|
||||
const from = ref('');
|
||||
const event = ref('');
|
||||
const q = ref('');
|
||||
const page = ref(1);
|
||||
const pageSize = 30;
|
||||
const items = ref([]);
|
||||
const total = ref(0);
|
||||
const aggregates = ref({});
|
||||
const loading = ref(false);
|
||||
|
||||
const EVENT_LABEL = {
|
||||
http_request: '调用', admin_login_ok: '登录成功', admin_login_fail: '登录失败',
|
||||
call_denied: '越权', user_created: '建用户', user_updated: '改用户', user_deleted: '删用户',
|
||||
};
|
||||
|
||||
function statusBadge(s) {
|
||||
if (!s) return 'muted';
|
||||
if (s >= 200 && s < 300) return 'green';
|
||||
if (s >= 500) return 'red';
|
||||
return 'amber';
|
||||
}
|
||||
|
||||
async function load() {
|
||||
loading.value = true;
|
||||
try {
|
||||
const r = await api.logs({
|
||||
from_date: from.value || undefined,
|
||||
event: event.value || undefined,
|
||||
q: q.value || undefined,
|
||||
page: page.value,
|
||||
page_size: pageSize,
|
||||
});
|
||||
items.value = r.items;
|
||||
total.value = r.total;
|
||||
aggregates.value = r.aggregates || {};
|
||||
} catch (e) {
|
||||
if (e.status === 401) emit('unauthorized');
|
||||
} finally {
|
||||
loading.value = false;
|
||||
}
|
||||
}
|
||||
|
||||
function search() { page.value = 1; load(); }
|
||||
function prev() { if (page.value > 1) { page.value -= 1; load(); } }
|
||||
function next() { if (page.value * pageSize < total.value) { page.value += 1; load(); } }
|
||||
|
||||
onMounted(load);
|
||||
</script>
|
||||
|
||||
<template>
|
||||
<div>
|
||||
<div class="stats" style="margin-bottom: 16px">
|
||||
<div class="stat"><div class="k">{{ total }}</div><div class="l">匹配日志</div></div>
|
||||
<div class="stat"><div class="k amber">{{ aggregates.denied ?? 0 }}</div><div class="l">越权 / 失败</div></div>
|
||||
<div class="stat"><div class="k">{{ aggregates.avg_latency_ms ?? '—' }}</div><div class="l">平均耗时 ms</div></div>
|
||||
<div class="stat"><div class="k">{{ aggregates.p95_latency_ms ?? '—' }}</div><div class="l">P95 耗时 ms</div></div>
|
||||
</div>
|
||||
<div class="toolbar">
|
||||
<input class="date" v-model="from" type="date">
|
||||
<select v-model="event">
|
||||
<option value="">全部事件</option>
|
||||
<option v-for="(l, k) in EVENT_LABEL" :key="k" :value="k">{{ l }}</option>
|
||||
</select>
|
||||
<div class="search"><input v-model="q" placeholder="按用户 / 路径过滤…" @keyup.enter="search"></div>
|
||||
<button class="btn btn-ghost" @click="search">查询</button>
|
||||
</div>
|
||||
<div class="panel">
|
||||
<table>
|
||||
<thead><tr><th>时间</th><th>用户</th><th>事件</th><th>方法</th><th>路径</th><th>状态</th><th>耗时</th></tr></thead>
|
||||
<tbody>
|
||||
<tr v-if="loading"><td colspan="7" class="cell-id">加载中…</td></tr>
|
||||
<tr v-else-if="!items.length"><td colspan="7" class="cell-id">暂无日志</td></tr>
|
||||
<tr v-for="(e, i) in items" :key="i">
|
||||
<td class="cell-id mono">{{ (e.ts || '').replace('T', ' ').slice(0, 19) }}</td>
|
||||
<td class="mono cell-id">{{ e.user_id ? e.user_id.slice(0, 8) : '—' }}</td>
|
||||
<td>{{ EVENT_LABEL[e.event] || e.event }}</td>
|
||||
<td class="cell-id mono">{{ e.method || '—' }}</td>
|
||||
<td class="mono">{{ e.path || '—' }}</td>
|
||||
<td><span v-if="e.status" class="badge" :class="statusBadge(e.status)">{{ e.status }}</span><span v-else class="badge muted">—</span></td>
|
||||
<td class="cell-id mono">{{ e.latency_ms != null ? e.latency_ms + 'ms' : '—' }}</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
<div class="pag">
|
||||
<span>共 {{ total }} 条</span><span class="spacer"></span>
|
||||
<button :disabled="page <= 1" @click="prev">上一页</button>
|
||||
<button class="on">{{ page }}</button>
|
||||
<button :disabled="page * pageSize >= total" @click="next">下一页</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<style scoped>
|
||||
input.date { background: var(--surface); border: 1px solid var(--border); color: var(--text); border-radius: var(--radius-sm); padding: 8px 10px; font-size: 13px; outline: none; }
|
||||
</style>
|
||||
@@ -0,0 +1,72 @@
|
||||
<script setup>
|
||||
import { ref, onMounted } from 'vue';
|
||||
import { api } from '../api/admin.js';
|
||||
import UserForm from '../components/UserForm.vue';
|
||||
|
||||
const emit = defineEmits(['unauthorized']);
|
||||
|
||||
const users = ref([]);
|
||||
const loading = ref(true);
|
||||
const formOpen = ref(false);
|
||||
const formMode = ref('create');
|
||||
const editing = ref(null);
|
||||
|
||||
const ROLE_LABEL = { admin: '管理员', researcher: '研究员', guest: '访客' };
|
||||
const UNIT_LABEL = { qa: '问答', kg: '知识图谱', extraction: '文献抽取', docgen: '文档生成', figures: '图片图库', private_docs: '私有资料', bioflow: '生信分析', paper: '论文撰写', legacy: '旧版接口' };
|
||||
|
||||
async function load() {
|
||||
loading.value = true;
|
||||
try {
|
||||
const r = await api.users();
|
||||
users.value = r.users;
|
||||
} catch (e) {
|
||||
if (e.status === 401) emit('unauthorized');
|
||||
} finally {
|
||||
loading.value = false;
|
||||
}
|
||||
}
|
||||
|
||||
function create() { formMode.value = 'create'; editing.value = null; formOpen.value = true; }
|
||||
function edit(u) { formMode.value = 'edit'; editing.value = u; formOpen.value = true; }
|
||||
async function del(u) {
|
||||
if (!confirm(`删除用户「${u.username}」?此操作不可撤销。`)) return;
|
||||
try { await api.deleteUser(u.id); load(); } catch (e) { alert(e.message || '删除失败'); }
|
||||
}
|
||||
function onSaved() { formOpen.value = false; load(); }
|
||||
|
||||
onMounted(load);
|
||||
</script>
|
||||
|
||||
<template>
|
||||
<div>
|
||||
<div class="toolbar">
|
||||
<div class="spacer"></div>
|
||||
<button class="btn btn-primary" @click="create">+ 新建用户</button>
|
||||
</div>
|
||||
<div class="panel">
|
||||
<table>
|
||||
<thead><tr><th>用户</th><th>角色</th><th>授权范围</th><th>创建时间</th><th style="width: 120px"></th></tr></thead>
|
||||
<tbody>
|
||||
<tr v-if="loading"><td colspan="5" class="cell-id">加载中…</td></tr>
|
||||
<tr v-else-if="!users.length"><td colspan="5" class="cell-id">暂无用户</td></tr>
|
||||
<tr v-for="u in users" :key="u.id">
|
||||
<td>
|
||||
<div style="display: flex; align-items: center; gap: 10px">
|
||||
<div class="avatar">{{ u.username[0] }}</div>
|
||||
<div><div>{{ u.username }}</div><div class="cell-id" style="font-size: 12px">{{ u.id.slice(0, 8) }}</div></div>
|
||||
</div>
|
||||
</td>
|
||||
<td><span class="role" :class="u.role">{{ ROLE_LABEL[u.role] || u.role }}</span></td>
|
||||
<td>
|
||||
<span v-if="u.role === 'admin'" class="chip amber">全部权限</span>
|
||||
<template v-else><span v-for="a in u.authorizations" :key="a" class="chip">{{ UNIT_LABEL[a] || a }}</span></template>
|
||||
</td>
|
||||
<td class="cell-id mono">{{ (u.created_at || '').slice(0, 10) }}</td>
|
||||
<td><div class="row-actions"><span class="act" @click="edit(u)">编辑</span><span class="act danger" @click="del(u)">删除</span></div></td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
<UserForm :open="formOpen" :mode="formMode" :user="editing" @close="formOpen = false" @saved="onSaved" />
|
||||
</div>
|
||||
</template>
|
||||
@@ -0,0 +1,15 @@
|
||||
import { defineConfig } from 'vite';
|
||||
import vue from '@vitejs/plugin-vue';
|
||||
|
||||
export default defineConfig({
|
||||
base: '/admin/',
|
||||
plugins: [vue()],
|
||||
server: {
|
||||
host: '127.0.0.1',
|
||||
port: 5181,
|
||||
proxy: {
|
||||
'/api': { target: 'http://127.0.0.1:8090', changeOrigin: false },
|
||||
},
|
||||
},
|
||||
build: { outDir: 'dist', target: 'es2020' },
|
||||
});
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
|
||||
@@ -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)}"
|
||||
)
|
||||
|
||||
@@ -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,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)
|
||||
@@ -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)
|
||||
@@ -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"],
|
||||
|
||||
@@ -17,7 +17,7 @@ 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 .routers import bioarchives, bioflow, biojobs, extension
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -74,6 +74,7 @@ def create_app() -> FastAPI:
|
||||
app.include_router(bioflow.router)
|
||||
app.include_router(biojobs.router)
|
||||
app.include_router(bioarchives.router)
|
||||
app.include_router(extension.router)
|
||||
return app
|
||||
|
||||
|
||||
|
||||
@@ -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 围栏。
|
||||
@@ -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:
|
||||
|
||||
@@ -239,22 +239,113 @@ async def upload_file(
|
||||
|
||||
|
||||
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 +486,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)
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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_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,14 @@ 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_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
|
||||
@@ -0,0 +1,275 @@
|
||||
"""app/services/bioflow/extension_streams.py —— 扩展试运行/运行 SSE 事件流(契约 §4)。
|
||||
|
||||
供 Dev-A(试运行/运行端点)与 Dev-C(前端聊天流)对接的 async 生成器接口:
|
||||
- trial_stream(ext_id, executor, request, events, log_path):
|
||||
事件 extension_trial running → log(增量) → done/failed(契约 §4 形态)。
|
||||
- run_stream(run_id, ext_id, executor, request, events, log_path):
|
||||
事件 extension_run_started → extension_run_log(增量) → done/failed。
|
||||
事件帧 = {"event": ..., "data": {...}};发布走 events.publish(兼容现有
|
||||
EventBus 的 {type, seq, data} 帧,data 即契约 §4 的 data)。
|
||||
|
||||
零新依赖;仅依赖 extension_exec(Dev-B 同 agent)。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any, AsyncIterator
|
||||
|
||||
from .extension_exec import EXIT_OUTPUT_CONTRACT, ExtensionExecConfig, ExtensionExecutor, ExtensionRunRequest
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
LOG_CHUNK_LINES = 40 # 每帧日志行数上限(前端滚动可读)
|
||||
LOG_MAX_LINES = 300 # 单帧 log_tail 累计尾部行数上限
|
||||
LOG_FLUSH_LINES = 80 # 累计行数阈值:超出即刷一帧(不等子进程结束)
|
||||
|
||||
|
||||
def _frame(event: str, data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""契约 §4 事件帧形态 {event, data}。"""
|
||||
return {"event": event, "data": dict(data)}
|
||||
|
||||
|
||||
class _suppress:
|
||||
"""抑制 aclose 异常(生成器关闭期的 RuntimeError/取消竞态不向上抛)。"""
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return True
|
||||
|
||||
|
||||
def _tail(text: str, n: int) -> str:
|
||||
lines = text.splitlines()
|
||||
if len(lines) <= n:
|
||||
return text
|
||||
return "\n".join(lines[-n:])
|
||||
|
||||
|
||||
class _LogAccumulator:
|
||||
"""增量日志缓冲:按行累计,超阈值产出增量帧(log 字段只含新增行)。"""
|
||||
|
||||
def __init__(self, flush_lines: int = LOG_FLUSH_LINES) -> None:
|
||||
self.full: list[str] = []
|
||||
self.sent = 0
|
||||
self.flush_lines = flush_lines
|
||||
|
||||
def add(self, text: str) -> str:
|
||||
self.full.extend(text.splitlines())
|
||||
return self.full
|
||||
|
||||
def pending_lines(self) -> int:
|
||||
return len(self.full) - self.sent
|
||||
|
||||
def drain(self, max_lines: int = LOG_CHUNK_LINES) -> str:
|
||||
"""取出未发送的增量(至多 max_lines 行)。"""
|
||||
if self.sent >= len(self.full):
|
||||
return ""
|
||||
chunk = self.full[self.sent:self.sent + max_lines]
|
||||
self.sent += len(chunk)
|
||||
return "\n".join(chunk)
|
||||
|
||||
def tail(self, n: int = LOG_MAX_LINES) -> str:
|
||||
return _tail("\n".join(self.full), n)
|
||||
|
||||
def write_file(self, path: Path) -> None:
|
||||
try:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text("\n".join(self.full), encoding="utf-8")
|
||||
except OSError:
|
||||
logger.warning("extension stream: 日志落盘失败 %s", path)
|
||||
|
||||
|
||||
def _final_trial_frame(ext_id: str, outcome: dict[str, Any],
|
||||
log_buf: _LogAccumulator, log_path: Path) -> dict[str, Any]:
|
||||
"""试运行终帧:done(含 exit_code 透传,3=输出契约校验失败)或 failed。"""
|
||||
log_buf.write_file(log_path)
|
||||
exit_code = int(outcome.get("exit_code") or 0)
|
||||
timed_out = bool(outcome.get("timed_out"))
|
||||
data: dict[str, Any] = {
|
||||
"ext_id": ext_id,
|
||||
"log_tail": log_buf.tail(),
|
||||
"exit_code": exit_code,
|
||||
}
|
||||
if timed_out:
|
||||
data["status"] = "failed"
|
||||
data["error"] = f"试运行超时({exit_code} 进程已终止)"
|
||||
return _frame("extension_trial", data)
|
||||
if exit_code == EXIT_OUTPUT_CONTRACT:
|
||||
data["status"] = "failed"
|
||||
data["error"] = "输出契约校验失败:输出目录为空(退出码 3)"
|
||||
return _frame("extension_trial", data)
|
||||
if exit_code == 0:
|
||||
data["status"] = "done"
|
||||
return _frame("extension_trial", data)
|
||||
data["status"] = "failed"
|
||||
data["error"] = f"试运行失败:命令退出码 {exit_code}"
|
||||
return _frame("extension_trial", data)
|
||||
|
||||
|
||||
async def trial_stream(
|
||||
ext_id: str,
|
||||
executor: ExtensionExecutor,
|
||||
request: ExtensionRunRequest,
|
||||
events: Any,
|
||||
*,
|
||||
log_path: Path | None = None,
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
"""试运行事件流:extension_trial running → log(增量,节流) → done/failed。
|
||||
|
||||
用法(Dev-A trial-run 端点):
|
||||
async for frame in trial_stream(...):
|
||||
yield to_sse(frame) # 每帧一个 SSE data 块
|
||||
"""
|
||||
log_path = Path(log_path) if log_path else Path(request.output_dir) / "trial.log"
|
||||
yield _frame("extension_trial", {"ext_id": ext_id, "status": "running", "log": ""})
|
||||
events.publish("extension_trial", {"ext_id": ext_id, "status": "running"})
|
||||
|
||||
log_buf = _LogAccumulator()
|
||||
agen = executor.run_stream(request)
|
||||
try:
|
||||
async for item in agen:
|
||||
if "chunk" in item:
|
||||
log_buf.add(item["chunk"])
|
||||
if log_buf.pending_lines() >= log_buf.flush_lines:
|
||||
frame = _frame("extension_trial", {
|
||||
"ext_id": ext_id, "status": "running",
|
||||
"log": log_buf.drain(),
|
||||
})
|
||||
events.publish("extension_trial", frame["data"])
|
||||
yield frame
|
||||
else: # final
|
||||
final = _final_trial_frame(ext_id, item, log_buf, log_path)
|
||||
events.publish("extension_trial", final["data"])
|
||||
yield final
|
||||
return
|
||||
except asyncio.CancelledError:
|
||||
log_buf.write_file(log_path)
|
||||
raise
|
||||
except Exception as exc: # executor 异常 → 防御性失败帧(不让异常冲出流)
|
||||
logger.warning("extension trial stream: executor 异常 %s", exc)
|
||||
log_buf.write_file(log_path)
|
||||
frame = _frame("extension_trial", {
|
||||
"ext_id": ext_id, "status": "failed",
|
||||
"log_tail": log_buf.tail(), "exit_code": -1,
|
||||
"error": f"试运行中断(executor 异常): {exc}",
|
||||
})
|
||||
events.publish("extension_trial", frame["data"])
|
||||
yield frame
|
||||
return
|
||||
finally:
|
||||
# 消费方提前退出时关闭 executor 生成器(子进程 reader 及时回收)
|
||||
with _suppress():
|
||||
await agen.aclose()
|
||||
# 防御:流提前结束(executor 未给 final)
|
||||
log_buf.write_file(log_path)
|
||||
frame = _frame("extension_trial", {
|
||||
"ext_id": ext_id, "status": "failed",
|
||||
"log_tail": log_buf.tail(), "exit_code": -1,
|
||||
"error": "试运行中断(executor 异常)",
|
||||
})
|
||||
events.publish("extension_trial", frame["data"])
|
||||
yield frame
|
||||
|
||||
|
||||
async def run_stream(
|
||||
run_id: str,
|
||||
ext_id: str,
|
||||
executor: ExtensionExecutor,
|
||||
request: ExtensionRunRequest,
|
||||
events: Any,
|
||||
*,
|
||||
log_path: Path | None = None,
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
"""正式运行事件流:extension_run_started → extension_run_log(增量) → done/failed。
|
||||
|
||||
与 trial_stream 同管线(executor.run_stream),事件名按契约 §4 区分。
|
||||
"""
|
||||
log_path = Path(log_path) if log_path else Path(request.output_dir) / "run.log"
|
||||
yield _frame("extension_run_started", {"run_id": run_id, "ext_id": ext_id})
|
||||
events.publish("extension_run_started", {"run_id": run_id, "ext_id": ext_id})
|
||||
|
||||
log_buf = _LogAccumulator()
|
||||
agen = executor.run_stream(request)
|
||||
try:
|
||||
async for item in agen:
|
||||
if "chunk" in item:
|
||||
log_buf.add(item["chunk"])
|
||||
if log_buf.pending_lines() >= log_buf.flush_lines:
|
||||
frame = _frame("extension_run_log", {
|
||||
"run_id": run_id, "log": log_buf.drain(),
|
||||
})
|
||||
events.publish("extension_run_log", frame["data"])
|
||||
yield frame
|
||||
else: # final
|
||||
log_buf.write_file(log_path)
|
||||
exit_code = int(item.get("exit_code") or 0)
|
||||
if item.get("timed_out"):
|
||||
frame = _frame("extension_run_failed", {
|
||||
"run_id": run_id, "exit_code": exit_code,
|
||||
"error": f"运行超时({exit_code} 进程已终止)",
|
||||
"log_tail": log_buf.tail(),
|
||||
})
|
||||
elif exit_code == EXIT_OUTPUT_CONTRACT:
|
||||
frame = _frame("extension_run_failed", {
|
||||
"run_id": run_id, "exit_code": exit_code,
|
||||
"error": "输出契约校验失败:输出目录为空(退出码 3)",
|
||||
"log_tail": log_buf.tail(),
|
||||
})
|
||||
elif exit_code == 0:
|
||||
frame = _frame("extension_run_done", {
|
||||
"run_id": run_id,
|
||||
"outputs": _output_listing(request.output_dir),
|
||||
})
|
||||
else:
|
||||
frame = _frame("extension_run_failed", {
|
||||
"run_id": run_id, "exit_code": exit_code,
|
||||
"error": f"扩展运行失败:命令退出码 {exit_code}",
|
||||
"log_tail": log_buf.tail(),
|
||||
})
|
||||
events.publish(frame["event"], frame["data"])
|
||||
yield frame
|
||||
return
|
||||
except asyncio.CancelledError:
|
||||
log_buf.write_file(log_path)
|
||||
raise
|
||||
except Exception as exc: # executor 异常 → 防御性失败帧
|
||||
logger.warning("extension run stream: executor 异常 %s", exc)
|
||||
log_buf.write_file(log_path)
|
||||
frame = _frame("extension_run_failed", {
|
||||
"run_id": run_id, "exit_code": -1,
|
||||
"error": f"运行中断(executor 异常): {exc}", "log_tail": log_buf.tail(),
|
||||
})
|
||||
events.publish("extension_run_failed", frame["data"])
|
||||
yield frame
|
||||
return
|
||||
finally:
|
||||
with _suppress():
|
||||
await agen.aclose()
|
||||
# 防御:流提前结束(executor 未给 final)
|
||||
log_buf.write_file(log_path)
|
||||
frame = _frame("extension_run_failed", {
|
||||
"run_id": run_id, "exit_code": -1,
|
||||
"error": "运行中断(executor 异常)", "log_tail": log_buf.tail(),
|
||||
})
|
||||
events.publish("extension_run_failed", frame["data"])
|
||||
yield frame
|
||||
|
||||
|
||||
def _output_listing(output_dir: Path) -> list[dict[str, Any]]:
|
||||
"""输出目录清单(契约 run_result.outputs 形态:path/size/type)。"""
|
||||
out: list[dict[str, Any]] = []
|
||||
try:
|
||||
for path in sorted(Path(output_dir).iterdir()):
|
||||
record: dict[str, Any] = {
|
||||
"path": str(path),
|
||||
"type": "dir" if path.is_dir() else "file",
|
||||
}
|
||||
if path.is_file():
|
||||
record["size"] = path.stat().st_size
|
||||
out.append(record)
|
||||
except OSError:
|
||||
pass
|
||||
return out
|
||||
File diff suppressed because it is too large
Load Diff
@@ -358,3 +358,178 @@ def _family_hit(text_lower: str) -> bool:
|
||||
if not tail_stripped.strip(",。!?;:,.!?;:()() "):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 用户自定义扩展分析意图(契约 §5 两层判定;Dev-A 纯函数导出)
|
||||
#
|
||||
# 匹配位置 = 生信服务 chat/stream 入口(先于 LLM 阶段规划);chat/stream 的
|
||||
# 接线由 Dev-C 联调时在主分支整合,本节只导出纯函数。
|
||||
#
|
||||
# A 层创建:动作词 × 目标词同句命中 → 创建流程(上传引导卡)。
|
||||
# B 层运行:仅 published 扩展参与打分(display_name 模糊 + nl_triggers +
|
||||
# analysis_type 加权,动作词加分)。
|
||||
# 歧义消解:目标词命中但无动作词 → 有 B 层命中走运行;无命中回
|
||||
# 「没有匹配的已发布扩展,是否创建?」引导卡。
|
||||
# 漏判安全:无命中且非创建 → None 路径(回落正常生信阶段规划,零影响)。
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# 创建意图目标词(契约 §5 A 层)
|
||||
EXTENSION_TARGET_WORD = ["拓展分析", "分析模板", "分析脚本", "分析工具", "扩展"]
|
||||
|
||||
# 创建意图动作词(契约 §5 A 层;全词边界匹配——「建」不误命中「建议」等)
|
||||
EXTENSION_ACTION_WORD = ["创建", "新建", "上传", "自定义", "做一个", "建", "写", "弄"]
|
||||
|
||||
# 运行意图动作词(契约 §5 B 层加分)
|
||||
EXTENSION_RUN_ACTION_WORD = ["帮我跑", "帮我用", "做一下", "分析一下", "跑一下", "用一下"]
|
||||
|
||||
# 运行打分阈值:高于此分才算 B 层命中(漏判安全)
|
||||
EXTENSION_SCORE_THRESHOLD = 2.0
|
||||
|
||||
# 单字动作词的后随禁用字(防「建议/建设/写作」等复合词误命中;
|
||||
# 契约例「建个分析模板」→「个」不在禁用集 → 命中)
|
||||
_ACTION_STOP_SUFFIX: dict[str, set[str]] = {
|
||||
"建": {"议", "设", "筑", "造", "国"},
|
||||
"写": {"作", "字"},
|
||||
"弄": set(),
|
||||
}
|
||||
|
||||
|
||||
def _create_action_hit(text_lower: str) -> bool:
|
||||
"""创建动作词命中(长词子串 + 单字词禁用后缀检查)。"""
|
||||
for w in sorted(EXTENSION_ACTION_WORD, key=len, reverse=True):
|
||||
if len(w) > 1:
|
||||
if w in text_lower:
|
||||
return True
|
||||
continue
|
||||
start = 0
|
||||
while True:
|
||||
idx = text_lower.find(w, start)
|
||||
if idx < 0:
|
||||
break
|
||||
nxt = text_lower[idx + 1] if idx + 1 < len(text_lower) else ""
|
||||
if nxt not in _ACTION_STOP_SUFFIX.get(w, set()):
|
||||
return True
|
||||
start = idx + 1
|
||||
return False
|
||||
|
||||
|
||||
def extension_create_intent(text: str) -> bool:
|
||||
"""A 层创建意图:动作词 × 目标词同句命中(契约 §5)。"""
|
||||
tl = (text or "").lower()
|
||||
if not _contains_any_longest_first(tl, EXTENSION_TARGET_WORD):
|
||||
return False
|
||||
return _create_action_hit(tl)
|
||||
|
||||
|
||||
def extension_run_actions(text: str) -> bool:
|
||||
"""句子中是否含运行动作词(B 层加分项)。"""
|
||||
return _contains_any_longest_first((text or "").lower(), EXTENSION_RUN_ACTION_WORD)
|
||||
|
||||
|
||||
_ANALYSIS_TYPE_KEYWORDS: dict[str, tuple[str, ...]] = {
|
||||
"blast_search": ("blast", "比对"),
|
||||
"hmmer_search": ("hmm",),
|
||||
"table_filter": ("过滤", "筛选", "表格"),
|
||||
"sequence_extraction": ("提取", "抽取"),
|
||||
"plot_render": ("绘图", "画图", "作图"),
|
||||
"enrichment_or_stats": ("富集", "统计"),
|
||||
"gff_processing": ("gff",),
|
||||
"id_cleaning": ("去重", "清洗"),
|
||||
}
|
||||
|
||||
|
||||
def score_extension(text: str, ext: dict[str, object]) -> float:
|
||||
"""单扩展打分(纯函数):display_name 模糊 + nl_triggers + analysis_type,
|
||||
动作词加分。分项权重与 detect 内注释一致。"""
|
||||
score, _reasons = _score_extension(text, ext)
|
||||
return score
|
||||
|
||||
|
||||
def _score_extension(text: str, ext: dict[str, object]) -> tuple[float, list[str]]:
|
||||
tl = (text or "").lower()
|
||||
score = 0.0
|
||||
reasons: list[str] = []
|
||||
display = str(ext.get("display_name") or "").strip().lower()
|
||||
if display:
|
||||
if display in tl: # display_name 全文命中(最强信号)
|
||||
score += 4.0
|
||||
reasons.append("display_name")
|
||||
else:
|
||||
# 字符 2-gram 覆盖度(2-6 字名称模糊匹配;漏判安全:过半才计分)
|
||||
grams = {display[i:i + 2] for i in range(len(display) - 1)} if len(display) >= 2 else {display}
|
||||
hit = sum(1 for g in grams if g in tl)
|
||||
if hit == len(grams):
|
||||
score += 3.0
|
||||
reasons.append("display_fuzzy")
|
||||
elif hit >= max(1, len(grams) // 2):
|
||||
score += 1.5
|
||||
reasons.append("display_partial")
|
||||
triggers = [str(t).strip().lower() for t in (ext.get("nl_triggers") or []) if str(t).strip()]
|
||||
for t in triggers: # nl_triggers 关键词(只计一次最高)
|
||||
if t and t in tl:
|
||||
score += 2.5
|
||||
reasons.append(f"trigger:{t}")
|
||||
break
|
||||
for kw in _ANALYSIS_TYPE_KEYWORDS.get(str(ext.get("analysis_type") or ""), ()):
|
||||
if kw in tl:
|
||||
score += 1.5
|
||||
reasons.append(f"analysis_type:{kw}")
|
||||
break
|
||||
if extension_run_actions(tl):
|
||||
score += 1.0
|
||||
reasons.append("action_verb")
|
||||
return score, reasons
|
||||
|
||||
|
||||
def match_extensions(
|
||||
text: str,
|
||||
extensions: list[dict[str, object]],
|
||||
*,
|
||||
threshold: float = EXTENSION_SCORE_THRESHOLD,
|
||||
) -> list[dict[str, object]]:
|
||||
"""B 层运行意图打分:仅 published 扩展参与,得分 > 阈值,按分降序。
|
||||
|
||||
返回 [{ext_id, display_name, analysis_type, score, reasons, ext}]。
|
||||
"""
|
||||
out: list[dict[str, object]] = []
|
||||
for ext in extensions or []:
|
||||
if str(ext.get("status") or "") != "published":
|
||||
continue
|
||||
score, reasons = _score_extension(text, ext)
|
||||
if score > threshold:
|
||||
out.append({
|
||||
"ext_id": ext.get("ext_id"),
|
||||
"display_name": ext.get("display_name"),
|
||||
"analysis_type": ext.get("analysis_type"),
|
||||
"score": score,
|
||||
"reasons": reasons,
|
||||
"ext": ext,
|
||||
})
|
||||
out.sort(key=lambda x: float(x.get("score") or 0.0), reverse=True)
|
||||
return out
|
||||
|
||||
|
||||
def detect_extension_intent(
|
||||
text: str,
|
||||
extensions: list[dict[str, object]],
|
||||
) -> dict[str, object]:
|
||||
"""扩展意图两层判定主入口(契约 §5)。
|
||||
|
||||
返回(Dev-C 据此发卡):
|
||||
- {"layer": "create", ...} 创建流程(上传引导卡)
|
||||
- {"layer": "run", "candidates": [...]} 运行意图(高置信单命中/多命中候选)
|
||||
- {"layer": "create_ambiguous", ...} 歧义引导卡(「没有匹配的已发布扩展,是否创建?」)
|
||||
- {"layer": "none", ...} 漏判安全:回落正常生信阶段规划
|
||||
"""
|
||||
has_target = _contains_any_longest_first((text or "").lower(), EXTENSION_TARGET_WORD)
|
||||
has_create_action = _create_action_hit((text or "").lower())
|
||||
if has_target and has_create_action:
|
||||
return {"layer": "create", "create": True, "candidates": [], "ambiguous": False}
|
||||
candidates = match_extensions(text, extensions)
|
||||
if candidates:
|
||||
return {"layer": "run", "create": False, "candidates": candidates, "ambiguous": False}
|
||||
if has_target:
|
||||
# 目标词命中、无创建动作词、无 B 层命中 → 歧义引导卡(不误判为创建)
|
||||
return {"layer": "create_ambiguous", "create": False, "candidates": [], "ambiguous": True}
|
||||
return {"layer": "none", "create": False, "candidates": [], "ambiguous": False}
|
||||
|
||||
@@ -19,14 +19,19 @@ import json
|
||||
import logging
|
||||
import os
|
||||
import shlex
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import Any, Awaitable, Callable
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CONTAINER_CMD_PREFIX = ("python", "-B") # 容器内命令形态(TC-BIO-RUN-037)
|
||||
|
||||
# 日志流式落盘/回调节流常数:每 N 行或每 S 秒 flush+fsync 并推送一次
|
||||
LOG_FLUSH_LINES = 20
|
||||
LOG_FLUSH_SECS = 1.0
|
||||
|
||||
|
||||
class ImageMapError(Exception):
|
||||
"""镜像映射表缺失/非法(服务可用但所有 docker 执行前置拦截)。"""
|
||||
@@ -71,6 +76,9 @@ class KernelConfig:
|
||||
bba_env: dict[str, str] = field(default_factory=dict)
|
||||
cpu_allow: str = "4,8,16,32"
|
||||
memory_allow: str = "16g,32g,64g"
|
||||
extra_mounts: list[str] = field(default_factory=list) # 登记路径只读挂载(src==dst)
|
||||
kernel_dir: str = "" # 内核树宿主路径(运行时挂载进容器)
|
||||
container_kernel_dir: str = "/bioflow-kernel" # 容器内内核树路径
|
||||
|
||||
|
||||
class Kernel:
|
||||
@@ -118,6 +126,8 @@ class Kernel:
|
||||
cpu: str | None = None,
|
||||
memory: str | None = None,
|
||||
module: str | None = None,
|
||||
extra_mounts: list[str] | None = None,
|
||||
container_name: str | None = None,
|
||||
) -> list[str]:
|
||||
"""拼装 docker run 命令(list,无 shell;TC-BIO-KERN-001 基元)。
|
||||
|
||||
@@ -136,10 +146,19 @@ class Kernel:
|
||||
raise UnknownStageError(f"阶段 {stage} 无执行模块")
|
||||
|
||||
cmd: list[str] = ["docker", "run", "--rm", "--user", str(self.cfg.host_uid)]
|
||||
if container_name:
|
||||
cmd += ["--name", container_name]
|
||||
cmd += ["-v", f"{job_input_dir}:{self.cfg.container_input_dir}"]
|
||||
cmd += ["-v", f"{job_workspace_dir}:{self.cfg.container_workspace_dir}"]
|
||||
if self.cfg.data_root:
|
||||
cmd += ["-v", f"{self.cfg.data_root}:{self.cfg.container_data_dir}:ro"]
|
||||
# 登记路径只读挂载(src==dst,容器内原路径直接可见;登记时已过黑名单)
|
||||
for m in (extra_mounts or self.cfg.extra_mounts):
|
||||
cmd += ["-v", f"{m}:{m}:ro"]
|
||||
# 内核树运行时挂载(除 bioflow-gf 外其余镜像未内置内核,统一走挂载)
|
||||
if self.cfg.kernel_dir:
|
||||
cmd += ["-v", f"{self.cfg.kernel_dir}:{self.cfg.container_kernel_dir}:ro"]
|
||||
cmd += ["-w", self.cfg.container_kernel_dir]
|
||||
eff_cpu = cpu or str(params.get("cpu") or "")
|
||||
eff_memory = memory or str(params.get("memory") or "")
|
||||
cmd += ["--cpus", self._safe_cpu(eff_cpu or self.cfg.default_cpus)]
|
||||
@@ -154,9 +173,13 @@ class Kernel:
|
||||
cmd.append(f"{tag}:{self.cfg.docker_tag_suffix}")
|
||||
cmd.extend(CONTAINER_CMD_PREFIX)
|
||||
cmd.extend(["-m", mod])
|
||||
# 通用容器内参数:input-dir/workspace/cpu(各入口 argparse 兼容)
|
||||
cmd += ["--input-dir", self.cfg.container_input_dir]
|
||||
cmd += ["--workspace", self.cfg.container_workspace_dir]
|
||||
# 参数注入:有 cli_contract 的阶段按契约,否则沿用默认 --input-dir/--workspace
|
||||
contract = entry.get("cli_contract")
|
||||
if contract:
|
||||
cmd += self._contract_args(contract, params, job_input_dir)
|
||||
else:
|
||||
cmd += ["--input-dir", self.cfg.container_input_dir]
|
||||
cmd += ["--workspace", self.cfg.container_workspace_dir]
|
||||
for key in sorted(params):
|
||||
value = params[key]
|
||||
if key in ("env", "cpu", "memory"):
|
||||
@@ -174,6 +197,58 @@ class Kernel:
|
||||
cmd.append(str(value))
|
||||
return cmd
|
||||
|
||||
def _contract_args(self, contract: list[Any], params: dict[str, Any],
|
||||
job_input_dir: str | Path) -> list[str]:
|
||||
"""cli_contract 解析:[(flag, source), ...] → argv。
|
||||
|
||||
source 取值:
|
||||
- input_dir/workspace:容器内 /job/input、/job/workspace
|
||||
- output_final:/job/workspace/final(rnaseq 系 -o 语义:产物目录,
|
||||
CLI 内部以 output.parent 为 workspace 并建 output.parent/final——
|
||||
故 -o 必须指向 workspace/final,manifest 也落此处与 runner 检索对齐)
|
||||
- samples_yaml:/job/input/samples.yaml(宿主侧不存在则跳过该 flag,
|
||||
交由 CLI 报缺失,错误信息比服务端更可读)
|
||||
- prev_manifest:<stage>:上游阶段 manifest 路径
|
||||
/job/workspace/final/{prefix}.{stage}.manifest.json(缺失跳过)
|
||||
- prefix/job_id:从 params 取值,缺失跳过
|
||||
- param:<name>:从 params[name] 取值,缺失跳过(用户经确认表单提供)
|
||||
- tool:<name>:容器内工具命令(镜像 PATH 内裸命令,如 gatk/samtools)
|
||||
"""
|
||||
args: list[str] = []
|
||||
for flag, source in contract:
|
||||
if source == "input_dir":
|
||||
args += [flag, self.cfg.container_input_dir]
|
||||
elif source == "workspace":
|
||||
args += [flag, self.cfg.container_workspace_dir]
|
||||
elif source == "output_final":
|
||||
args += [flag, f"{self.cfg.container_workspace_dir}/final"]
|
||||
elif str(source).startswith("prev_manifest:"):
|
||||
stage_name = str(source).split(":", 1)[1]
|
||||
value = str(params.get("prefix") or "").strip()
|
||||
if value:
|
||||
args += [flag,
|
||||
f"{self.cfg.container_workspace_dir}/final/{value}.{stage_name}.manifest.json"]
|
||||
elif str(source).startswith("tool:"):
|
||||
name = str(source).split(":", 1)[1]
|
||||
args += [flag, name]
|
||||
elif source == "samples_yaml":
|
||||
if (Path(job_input_dir) / "samples.yaml").is_file():
|
||||
args += [flag, f"{self.cfg.container_input_dir}/samples.yaml"]
|
||||
elif source == "prefix":
|
||||
value = str(params.get("prefix") or "").strip()
|
||||
if value:
|
||||
args += [flag, value]
|
||||
elif source == "job_id":
|
||||
value = str(params.get("job_id") or "").strip()
|
||||
if value:
|
||||
args += [flag, value]
|
||||
elif str(source).startswith("param:"):
|
||||
key = str(source).split(":", 1)[1]
|
||||
value = params.get(key)
|
||||
if value is not None and str(value).strip():
|
||||
args += [flag, str(value)]
|
||||
return args
|
||||
|
||||
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()
|
||||
@@ -195,6 +270,19 @@ class Kernel:
|
||||
|
||||
# ---- 执行 ----
|
||||
|
||||
@staticmethod
|
||||
async def _terminate_then_kill(proc: asyncio.subprocess.Process) -> None:
|
||||
"""terminate → 3s → kill(TC-BIO-KERN-009;与 extension_exec.run_stream 同构)。"""
|
||||
try:
|
||||
proc.terminate()
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=3.0)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
except ProcessLookupError: # pragma: no cover - 进程已退出
|
||||
pass
|
||||
|
||||
async def run_stage(
|
||||
self,
|
||||
stage: str,
|
||||
@@ -205,18 +293,28 @@ class Kernel:
|
||||
log_path: str | Path | None = None,
|
||||
cpu: str | None = None,
|
||||
memory: str | None = None,
|
||||
extra_mounts: list[str] | None = None,
|
||||
container_name: str | None = None,
|
||||
on_log: Callable[[str], Awaitable[None]] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""执行一个阶段 docker:返回 {exit_code, timed_out, log_tail, cmd}。
|
||||
|
||||
- 超时:terminate → 3s → kill(TC-BIO-KERN-009)
|
||||
- 日志全量写 log_path(任务 logs/),返回尾部 log_tail_lines 行
|
||||
- extra_mounts:本次执行额外只读挂载的宿主绝对路径(登记路径,src==dst)
|
||||
- container_name:docker run --name(暂停时按名 kill 容器用)
|
||||
- on_log:增量日志回调(每 LOG_FLUSH_LINES 行或 LOG_FLUSH_SECS 秒推送
|
||||
一次 chunk,阶段运行期间实时可见;尾块兜底再推一次)。None 时不回调。
|
||||
- 超时:单行等待以 stage_timeout_sec 为限(与旧 communicate 语义一致,
|
||||
保守方案:长时间无输出不误杀)→ terminate → 3s → kill(TC-BIO-KERN-009)
|
||||
- 日志流式追加写 log_path(任务 logs/,执行中可读),返回尾部 log_tail_lines 行
|
||||
- 退出码非 0 / 超时 → 失败语义由 runner 判定(本函数只报事实)
|
||||
"""
|
||||
params = params or {}
|
||||
cpu = cpu or str(params.get("cpu") or "")
|
||||
memory = memory or str(params.get("memory") or "")
|
||||
cmd = self.build_command(stage, job_input_dir, job_workspace_dir, params,
|
||||
cpu=cpu or None, memory=memory or None)
|
||||
cpu=cpu or None, memory=memory or None,
|
||||
extra_mounts=extra_mounts,
|
||||
container_name=container_name)
|
||||
log_path = Path(log_path) if log_path else Path(job_workspace_dir) / "logs" / f"{stage}.log"
|
||||
log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
@@ -226,25 +324,53 @@ class Kernel:
|
||||
stderr=asyncio.subprocess.STDOUT,
|
||||
)
|
||||
timed_out = False
|
||||
text_parts: list[str] = []
|
||||
# 追加模式落盘:阶段执行中前端 logs 端点即可读;fsync 保证跨进程可见
|
||||
fh = log_path.open("a", encoding="utf-8")
|
||||
flush_buf: list[str] = []
|
||||
last_flush = time.monotonic()
|
||||
try:
|
||||
stdout, _ = await asyncio.wait_for(
|
||||
proc.communicate(), timeout=self.cfg.stage_timeout_sec
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
timed_out = True
|
||||
try:
|
||||
proc.terminate()
|
||||
while True:
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=3.0)
|
||||
line = await asyncio.wait_for(
|
||||
proc.stdout.readline(), timeout=self.cfg.stage_timeout_sec
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
stdout = b""
|
||||
timed_out = True
|
||||
await self._terminate_then_kill(proc)
|
||||
break
|
||||
if not line:
|
||||
break
|
||||
decoded = line.decode("utf-8", errors="replace")
|
||||
text_parts.append(decoded)
|
||||
fh.write(decoded)
|
||||
flush_buf.append(decoded)
|
||||
now = time.monotonic()
|
||||
if len(flush_buf) >= LOG_FLUSH_LINES or now - last_flush >= LOG_FLUSH_SECS:
|
||||
fh.flush()
|
||||
os.fsync(fh.fileno())
|
||||
if on_log is not None:
|
||||
await on_log("".join(flush_buf))
|
||||
flush_buf = []
|
||||
last_flush = now
|
||||
finally:
|
||||
# 尾块兜底 + 进程收尾(正常 EOF 或超时 kill 都落此处)
|
||||
if flush_buf:
|
||||
try:
|
||||
fh.flush()
|
||||
os.fsync(fh.fileno())
|
||||
except OSError: # pragma: no cover - fd 异常不阻塞结果
|
||||
pass
|
||||
if on_log is not None:
|
||||
try:
|
||||
await on_log("".join(flush_buf))
|
||||
except Exception: # pragma: no cover - 回调异常不吞结果
|
||||
logger.exception("bioflow kernel: on_log 尾块回调失败")
|
||||
fh.close()
|
||||
if proc.returncode is None:
|
||||
await proc.wait()
|
||||
exit_code = proc.returncode
|
||||
text = (stdout or b"").decode("utf-8", errors="replace")
|
||||
log_path.write_text(text, encoding="utf-8")
|
||||
text = "".join(text_parts)
|
||||
return {
|
||||
"exit_code": exit_code,
|
||||
"timed_out": timed_out,
|
||||
@@ -278,6 +404,8 @@ def _default_kernel_config() -> KernelConfig:
|
||||
stage_timeout_sec=s.stage_timeout_sec,
|
||||
log_tail_lines=s.log_tail_lines,
|
||||
bba_env=s.bba_env,
|
||||
kernel_dir=s.kernel_dir_abs,
|
||||
container_kernel_dir=s.container_kernel_dir,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -60,10 +60,18 @@ class ManifestError(Exception):
|
||||
def locate_manifest(workspace: str | Path, stage: str, prefix: str | None = None) -> Path | None:
|
||||
"""在 workspace final/ 下定位阶段 manifest(prefix + 阶段名双匹配)。
|
||||
|
||||
匹配优先级:
|
||||
1. prefix 匹配 + 阶段名中缀:final/{prefix}.{stage}.manifest.json
|
||||
2. 阶段名中缀:final/{任意}.{stage}.manifest.json
|
||||
3. 无中缀名兜底(BUG-FIX 2026-09-18):final/{stage}.manifest.json
|
||||
(rnaseq 系 CLI 空 prefix 时产出的文件名形态,如 rnaseq_qc.manifest.json;
|
||||
此前无此分支 → 退出码 0 也被判 no_manifest)。
|
||||
|
||||
不存在返回 None(runner 据此判 reason=no_manifest)。
|
||||
注意:无兜底——其他阶段的 manifest 绝不误配给当前阶段
|
||||
(重启恢复按 manifest 判成败时曾因兜底把 identify 的 manifest 判给
|
||||
phylogeny,导致 running 阶段被误判 done 不重跑)。
|
||||
注意:兜底只按「文件名精确等于 {stage}.manifest.json」匹配——其他阶段的
|
||||
manifest 绝不误配给当前阶段(重启恢复按 manifest 判成败时曾因任意
|
||||
*.manifest.json 兜底把 identify 的 manifest 判给 phylogeny,导致 running
|
||||
阶段被误判 done 不重跑)。
|
||||
"""
|
||||
final_dir = Path(workspace) / "final"
|
||||
if not final_dir.is_dir():
|
||||
@@ -78,6 +86,9 @@ def locate_manifest(workspace: str | Path, stage: str, prefix: str | None = None
|
||||
for cand in candidates:
|
||||
if f".{stage}." in cand.name:
|
||||
return cand
|
||||
for cand in candidates:
|
||||
if cand.name == f"{stage}.manifest.json":
|
||||
return cand
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@@ -19,6 +19,8 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
@@ -92,10 +94,523 @@ def _stage_input_satisfied(stage: str, job_input_dir: Path) -> tuple[bool, list[
|
||||
return False, list(DEFAULT_INPUT_HINTS)
|
||||
|
||||
|
||||
def ensure_samples_yaml(job_input_dir: Path, source_roots: list[Path] | None = None) -> list[str]:
|
||||
"""为 rnaseq 系阶段生成/修正 samples.yaml,返回需挂载的宿主目录列表。
|
||||
|
||||
逻辑(用户要求,勿做变体猜测):
|
||||
1. 描述文件发现:扫描登记目录树与任务 input(限深度 4)找 tsv/csv > xlsx >
|
||||
yaml/yml > md,解析样本清单(fastq 路径 + run 标识)。
|
||||
2. 文件实况分析:扫描源目录(登记目录本身)与任务 input,分类原始序列
|
||||
(.fq/.fastq 系)与压缩包(.tar.gz/.tar/.zip);压缩包解压到
|
||||
input/archives/(幂等);序列文件用于配对(不假设命名规律,只按 R1/R2
|
||||
语义:文件名含 f1/r1/1/forward 与 f2/r2/2/reverse)。
|
||||
3. 描述文件路径校正:逐条校验存在性;失效路径按 run 标识在实况序列文件中
|
||||
查找(文件名包含 run),唯一命中替换(按 R1/R2 语义分边);零/多命中保留
|
||||
原路径(CLI 报错)。
|
||||
4. 生成/重写 input/samples.yaml(引用真实存在路径),返回引用 fastq 的父目录
|
||||
(供执行容器挂载)。
|
||||
幂等:重跑时读入已是正确路径,校验通过后重写同内容。
|
||||
"""
|
||||
import re as _re
|
||||
import yaml as _yaml
|
||||
|
||||
# 描述文件解析函数(候选择优用;定义在循环前供嵌套闭包使用)
|
||||
def _parse_tsv_csv(desc: Path, seq_exts_local: tuple) -> list[dict]:
|
||||
import csv as _csv
|
||||
out: list[dict] = []
|
||||
with desc.open(encoding="utf-8", errors="replace", newline="") as fh:
|
||||
reader = _csv.DictReader(
|
||||
fh, delimiter="\t" if desc.name.lower().endswith(".tsv") else ","
|
||||
)
|
||||
for row in reader:
|
||||
r1 = str(row.get("fastq_1") or "").strip()
|
||||
r2 = str(row.get("fastq_2") or "").strip()
|
||||
run = str(row.get("run") or "").strip()
|
||||
sid = str(row.get("sample_id") or "").strip()
|
||||
if not r1 and not r2:
|
||||
for v in row.values():
|
||||
if str(v).strip().lower().endswith(seq_exts_local):
|
||||
if not r1:
|
||||
r1 = str(v).strip()
|
||||
elif not r2:
|
||||
r2 = str(v).strip()
|
||||
if r1 or r2 or run:
|
||||
out.append({
|
||||
"name": sid or run or (r1 and Path(r1).name.split(".")[0]) or "sample",
|
||||
"run": run, "r1": r1, "r2": r2,
|
||||
"meta": {k: v for k, v in row.items()
|
||||
if k not in ("fastq_1", "fastq_2", "sample_id", "run")},
|
||||
})
|
||||
return out
|
||||
|
||||
def _parse_xlsx(desc: Path) -> list[dict]:
|
||||
import openpyxl as _xl
|
||||
out: list[dict] = []
|
||||
wb = _xl.load_workbook(desc, read_only=True, data_only=True)
|
||||
try:
|
||||
ws = wb.active
|
||||
rows = list(ws.iter_rows(values_only=True))
|
||||
if rows:
|
||||
header = [str(h or "").strip() for h in rows[0]]
|
||||
for row in rows[1:]:
|
||||
d = dict(zip(header, row))
|
||||
run = str(d.get("run") or "").strip()
|
||||
if not run:
|
||||
continue
|
||||
out.append({
|
||||
"name": run, "run": run, "r1": "", "r2": "",
|
||||
"meta": {k: v for k, v in d.items() if k != "run"},
|
||||
})
|
||||
finally:
|
||||
wb.close()
|
||||
return out
|
||||
|
||||
def _parse_yaml(desc: Path) -> list[dict]:
|
||||
raw = _yaml.safe_load(desc.read_text(encoding="utf-8", errors="replace"))
|
||||
out: list[dict] = []
|
||||
samples = raw.get("samples", {}) if isinstance(raw, dict) else {}
|
||||
for sname, spaths in samples.items():
|
||||
r1 = r2 = run = ""
|
||||
if isinstance(spaths, dict):
|
||||
r1 = str(spaths.get("r1") or spaths.get("read1") or "")
|
||||
r2 = str(spaths.get("r2") or spaths.get("read2") or "")
|
||||
run = str(spaths.get("run") or "")
|
||||
elif isinstance(spaths, (list, tuple)) and len(spaths) >= 2:
|
||||
r1, r2 = str(spaths[0]), str(spaths[1])
|
||||
if r1 or r2:
|
||||
out.append({
|
||||
"name": str(sname), "run": run, "r1": r1, "r2": r2, "meta": {},
|
||||
})
|
||||
return out
|
||||
|
||||
yaml_path = job_input_dir / "samples.yaml"
|
||||
seq_exts = (".fq.gz", ".fastq.gz", ".fq", ".fastq")
|
||||
roots: list[Path] = [job_input_dir] + [Path(r) for r in (source_roots or []) if Path(r).exists()]
|
||||
|
||||
def _is_seq(p: Path) -> bool:
|
||||
return p.is_file() and p.name.lower().endswith(seq_exts)
|
||||
|
||||
def _r1_sem(name: str) -> bool:
|
||||
n = name.lower()
|
||||
return bool(_re.search(r"[._-](f1|r1|1|forward)([._-]|$)", n))
|
||||
|
||||
def _r2_sem(name: str) -> bool:
|
||||
n = name.lower()
|
||||
return bool(_re.search(r"[._-](f2|r2|2|reverse)([._-]|$)", n))
|
||||
|
||||
def _walk(root: Path):
|
||||
yield from root.rglob("*")
|
||||
|
||||
# ---- 2. 实况扫描:源目录(真实文件)+ 任务 input(跳过 symlink 重复) ----
|
||||
# 同文件去重:符号链接与其目标(samefile)只留一个 —— 优先保留真实文件(非链接),
|
||||
# 避免链接的别名文件名(如 CRR809753_2.fq.gz → _r2.fq.gz)被当成第二个样本。
|
||||
seq_files: list[Path] = []
|
||||
archives: list[Path] = []
|
||||
seen_inodes: set[tuple[int, int]] = set()
|
||||
|
||||
def _seq_dedup(p: Path, pool: list[Path]) -> bool:
|
||||
try:
|
||||
key = (p.stat().st_dev, p.stat().st_ino)
|
||||
except OSError:
|
||||
key = None
|
||||
if key is None or key not in seen_inodes:
|
||||
if key is not None:
|
||||
seen_inodes.add(key)
|
||||
pool.append(p)
|
||||
return True
|
||||
# 已存在同 inode 文件:若新候选是真实文件、旧的是链接,用真实文件替换
|
||||
for i, q in enumerate(pool):
|
||||
if q.samefile(p):
|
||||
if p.is_symlink() and not q.is_symlink():
|
||||
return False
|
||||
if not p.is_symlink() and q.is_symlink():
|
||||
pool[i] = p
|
||||
return True
|
||||
return False
|
||||
return False
|
||||
|
||||
for root in roots:
|
||||
for p in sorted(_walk(root)):
|
||||
if not p.is_file():
|
||||
continue
|
||||
if p.is_symlink() and root == job_input_dir:
|
||||
continue # input 里是登记 symlink,源目录已有真实文件
|
||||
if _is_seq(p):
|
||||
_seq_dedup(p, seq_files)
|
||||
elif p.name.lower().endswith((".tar.gz", ".tar", ".zip", ".tgz")):
|
||||
_seq_dedup(p, archives)
|
||||
seq_files = sorted({str(p): p for p in seq_files}.values(), key=lambda p: str(p))
|
||||
archives = sorted({str(p): p for p in archives}.values(), key=lambda p: str(p))
|
||||
|
||||
# 解压压缩包(幂等:archives/ 已存在同名解出物则跳过)
|
||||
archive_dir = job_input_dir / "archives"
|
||||
for a in archives:
|
||||
target_dir = archive_dir / (a.name.rsplit(".", 1)[0])
|
||||
if target_dir.exists():
|
||||
continue
|
||||
try:
|
||||
import tarfile
|
||||
import zipfile
|
||||
target_dir.mkdir(parents=True, exist_ok=True)
|
||||
if a.name.endswith((".tar.gz", ".tar", ".tgz")):
|
||||
with tarfile.open(a, "r:*") as tf:
|
||||
tf.extractall(target_dir)
|
||||
else:
|
||||
with zipfile.ZipFile(a) as zf:
|
||||
zf.extractall(target_dir)
|
||||
except Exception: # noqa: BLE001
|
||||
pass # 解压失败不阻断;后续扫描自然不含其内容
|
||||
if archive_dir.exists():
|
||||
for p in sorted(archive_dir.rglob("*")):
|
||||
if _is_seq(p):
|
||||
seq_files.append(p)
|
||||
seq_files = sorted({str(p): p for p in seq_files}.values(), key=lambda p: str(p))
|
||||
|
||||
# ---- 1. 描述文件发现:候选全收 → 逐份解析 → 按解析出的样本条目数择优 ----
|
||||
def _depth_ok(p: Path, root: Path) -> bool:
|
||||
try:
|
||||
return len(p.relative_to(root).parts) <= 4
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
candidates: list[Path] = []
|
||||
for root in roots:
|
||||
for p in sorted(root.rglob("*")):
|
||||
if not p.is_file() or p.name == "samples.yaml":
|
||||
continue
|
||||
if not _depth_ok(p, root):
|
||||
continue
|
||||
if p.name.lower().endswith((".tsv", ".csv", ".xlsx", ".xls", ".yaml", ".yml")):
|
||||
candidates.append(p)
|
||||
|
||||
described: list[dict] = []
|
||||
best: tuple[int, list[dict]] | None = None
|
||||
for desc in candidates:
|
||||
try:
|
||||
if desc.name.lower().endswith((".tsv", ".csv")):
|
||||
parsed = _parse_tsv_csv(desc, seq_exts)
|
||||
elif desc.name.lower().endswith((".xlsx", ".xls")):
|
||||
parsed = _parse_xlsx(desc)
|
||||
else:
|
||||
parsed = _parse_yaml(desc)
|
||||
except Exception: # noqa: BLE001
|
||||
parsed = []
|
||||
# 样本条目数最多者胜;平手时 tsv/csv 优先(语义最全),再按路径排序稳定
|
||||
is_table = desc.name.lower().endswith((".tsv", ".csv"))
|
||||
if best is None or len(parsed) > best[0] or (len(parsed) == best[0] and is_table and not best_is_table):
|
||||
best = (len(parsed), parsed)
|
||||
best_is_table = is_table
|
||||
if best and best[0] > 0:
|
||||
described = best[1]
|
||||
|
||||
# ---- 3. 描述文件路径校正(实况为准,按 run 标识,不做变体猜测) ----
|
||||
seq_by_run: dict[str, list[Path]] = {}
|
||||
for p in seq_files:
|
||||
m = _re.search(r"([A-Za-z0-9]{6,})", p.stem)
|
||||
if m:
|
||||
seq_by_run.setdefault(m.group(1), []).append(p)
|
||||
if described:
|
||||
for d in described:
|
||||
for key in ("r1", "r2"):
|
||||
p = Path(d[key]) if d[key] else None
|
||||
if p and p.exists():
|
||||
continue
|
||||
run = d.get("run") or ""
|
||||
cands = [q for q in seq_by_run.get(run, [])] if run else []
|
||||
if not cands:
|
||||
cands = [q for q in seq_files if d["name"] and d["name"] in q.name]
|
||||
if len(cands) == 1:
|
||||
d[key] = str(cands[0])
|
||||
elif len(cands) >= 2:
|
||||
mate = [q for q in cands if (_r2_sem(q.name) if key == "r2" else _r1_sem(q.name))]
|
||||
if len(mate) == 1:
|
||||
d[key] = str(mate[0])
|
||||
# 零命中:保留原路径,CLI 报错
|
||||
# 实况中还有描述未覆盖的序列 → 追加
|
||||
known = {Path(d["r1"]).name for d in described if d["r1"]}
|
||||
known |= {Path(d["r2"]).name for d in described if d["r2"]}
|
||||
for p in seq_files:
|
||||
if p.name not in known:
|
||||
described.append({"name": p.stem, "run": "", "r1": str(p), "r2": "", "meta": {}})
|
||||
|
||||
def _pair_key(name: str) -> str:
|
||||
"""归一配对比较键(小写):_f1/_r2/_1/_2 别名归一到同一样本。"""
|
||||
return _re.sub(
|
||||
r"[._-](f1|r1|1|forward|f2|r2|2|reverse)([._-]|$).*", "", name.lower()
|
||||
)
|
||||
|
||||
def _pair_base(name: str) -> str:
|
||||
"""保留原大小写的样本名:剥离 R1/R2 语义后缀。"""
|
||||
return _re.sub(
|
||||
r"[._-](f1|r1|1|forward|f2|r2|2|reverse)([._-]|$).*", "", name,
|
||||
flags=_re.IGNORECASE,
|
||||
)
|
||||
|
||||
# ---- 无描述文件:实况配对(不假设命名规律,按 R1/R2 语义) ----
|
||||
if not described and seq_files:
|
||||
r1s = [p for p in seq_files if _r1_sem(p.name)]
|
||||
r2s = [p for p in seq_files if _r2_sem(p.name)]
|
||||
used_r2: set[str] = set() # 一对一配对:已被认领的 r2 不再参与
|
||||
for p in r1s:
|
||||
base = _pair_key(p.name)
|
||||
mate_cands = [
|
||||
q for q in r2s if _pair_key(q.name) == base and str(q) not in used_r2
|
||||
]
|
||||
mate: Path | None = None
|
||||
if len(mate_cands) == 1:
|
||||
mate = mate_cands[0]
|
||||
elif len(mate_cands) > 1:
|
||||
# 多个同名 stem 候选:优先真实 f2/r2/reverse 命名,纯数字 _2 别名靠后
|
||||
pref = [
|
||||
q for q in mate_cands
|
||||
if _re.search(r"[._-](f2|r2|reverse)([._-]|$)", q.name.lower())
|
||||
]
|
||||
mate = pref[0] if pref else mate_cands[0]
|
||||
described.append({
|
||||
"name": _pair_base(p.name) or p.stem, "run": "",
|
||||
"r1": str(p), "r2": str(mate) if mate else "", "meta": {},
|
||||
})
|
||||
if mate is not None:
|
||||
used_r2.add(str(mate))
|
||||
for p in r2s:
|
||||
if str(p) not in used_r2:
|
||||
described.append({"name": p.stem, "run": "", "r1": "", "r2": str(p), "meta": {}})
|
||||
for p in seq_files:
|
||||
if not _r1_sem(p.name) and not _r2_sem(p.name):
|
||||
described.append({"name": p.stem, "run": "", "r1": str(p), "r2": "", "meta": {}})
|
||||
|
||||
# ---- 4. 生成/重写 samples.yaml + 返回挂载目录 ----
|
||||
mounts: set[str] = set()
|
||||
samples_out: dict[str, dict] = {}
|
||||
for d in described:
|
||||
if not (d.get("r1") or d.get("r2")):
|
||||
continue
|
||||
entry: dict = {}
|
||||
for key in ("r1", "r2"):
|
||||
p = Path(d[key]) if d[key] else None
|
||||
if p and p.exists():
|
||||
entry[key] = str(p)
|
||||
mounts.add(str(p.parent))
|
||||
if not entry:
|
||||
continue
|
||||
for k, v in d.get("meta", {}).items():
|
||||
if v not in (None, ""):
|
||||
entry[str(k)] = v
|
||||
if d.get("run"):
|
||||
entry.setdefault("run", d["run"])
|
||||
samples_out[d["name"]] = entry
|
||||
if samples_out:
|
||||
yaml_path.write_text(
|
||||
_yaml.safe_dump({"samples": samples_out}, allow_unicode=True), encoding="utf-8"
|
||||
)
|
||||
return sorted(mounts)
|
||||
|
||||
|
||||
def inspect_unknown_files(paths: list[Path], max_files: int = 200) -> list[dict]:
|
||||
"""可观测特征函数:对无法按规则归类的文件,收集结构化特征供 LLM 兜底。
|
||||
|
||||
每项:{path, size, ext, magic(首 64 字节 hex + file 类型), parent,
|
||||
archive_kind(tar|zip|gz|None), archive_members(≤20 内嵌名)}。
|
||||
纯 stdlib + 子进程 file/tar/zcat/unzip,失败静默跳过该文件。
|
||||
"""
|
||||
import subprocess
|
||||
features: list[dict] = []
|
||||
for p in sorted(paths, key=lambda q: (q.stat().st_size if q.is_file() else 0), reverse=True)[:max_files]:
|
||||
if not p.is_file():
|
||||
continue
|
||||
try:
|
||||
size = p.stat().st_size
|
||||
head = p.read_bytes()[:64]
|
||||
magic = head.hex()
|
||||
ext = p.suffix.lower() or ""
|
||||
parent = str(p.parent)
|
||||
archive_kind = None
|
||||
members: list[str] = []
|
||||
file_type = ""
|
||||
try:
|
||||
file_type = subprocess.run(
|
||||
["file", "-b", str(p)], capture_output=True, text=True, timeout=10
|
||||
).stdout.strip()
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
if file_type.startswith("gzip") or magic.startswith("1f8b"):
|
||||
archive_kind = "gz"
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["gzip", "-l", str(p)], capture_output=True, text=True, timeout=10
|
||||
).stdout
|
||||
members = out.strip().splitlines()[1:]
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
elif file_type.startswith("POSIX tar") or magic.startswith(("7573746172",)):
|
||||
archive_kind = "tar"
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["tar", "-tzf", str(p)], capture_output=True, text=True, timeout=15
|
||||
).stdout
|
||||
members = [m for m in out.splitlines() if m.strip()][:20]
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
elif file_type.startswith("Zip archive") or magic.startswith("504b"):
|
||||
archive_kind = "zip"
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["unzip", "-l", str(p)], capture_output=True, text=True, timeout=15
|
||||
).stdout
|
||||
members = [m.strip().split()[-1] for m in out.splitlines()[3:-2] if m.strip()][:20]
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
features.append({
|
||||
"path": str(p), "size": size, "ext": ext, "magic": magic[:64],
|
||||
"file_type": file_type, "parent": parent,
|
||||
"archive_kind": archive_kind, "archive_members": members,
|
||||
})
|
||||
except Exception: # noqa: BLE001
|
||||
continue
|
||||
return features
|
||||
|
||||
|
||||
async def llm_sample_fallback(job_input_dir: Path, source_roots: list[Path]) -> bool:
|
||||
"""LLM 兜底:规则无法归纳时,把文件特征清单交给 DeepSeek 生成样本描述。
|
||||
|
||||
触发条件:描述文件解析 0 条且实况配对 0 样本(调用方判断)。
|
||||
返回是否成功生成 samples.yaml。LLM 失败/未配置/输出无法回验 → False,
|
||||
不影响登记主链。r1/r2 必须回验存在(清单路径或解压后可解析),否则丢弃。
|
||||
"""
|
||||
import yaml as _yaml
|
||||
from .llm import LLMError, complete_json
|
||||
|
||||
# 收集未知文件(排除已识别的序列/描述文件)
|
||||
known_exts = (".tsv", ".csv", ".xlsx", ".xls", ".yaml", ".yml", ".md",
|
||||
".fq.gz", ".fastq.gz", ".fq", ".fastq")
|
||||
unknown: list[Path] = []
|
||||
for root in [job_input_dir, *[Path(r) for r in source_roots if Path(r).exists()]]:
|
||||
for p in root.rglob("*"):
|
||||
if not p.is_file() or p.name == "samples.yaml":
|
||||
continue
|
||||
if p.name.lower().endswith(known_exts):
|
||||
continue
|
||||
unknown.append(p)
|
||||
if not unknown:
|
||||
return False
|
||||
features = inspect_unknown_files(unknown)
|
||||
if not features:
|
||||
return False
|
||||
|
||||
# system prompt(独立文件,可迭代)
|
||||
prompt_path = Path(__file__).resolve().parent.parent / "prompts" / "llm_sample_analysis.txt"
|
||||
try:
|
||||
system = prompt_path.read_text(encoding="utf-8")
|
||||
except OSError:
|
||||
system = "你是生信数据文件分析助手,只输出 JSON。"
|
||||
import json as _json
|
||||
user = (
|
||||
"以下是一个数据目录的文件特征清单,请推断其中的 RNA-seq FASTQ 样本并输出 JSON。\n"
|
||||
f"{_json.dumps(features, ensure_ascii=False, default=str)[:24000]}"
|
||||
)
|
||||
try:
|
||||
result = await complete_json(system, user, temperature=0.2, max_tokens=4000, timeout=30.0)
|
||||
except (LLMError, Exception): # noqa: BLE001
|
||||
return False
|
||||
samples = result.get("samples") if isinstance(result, dict) else None
|
||||
if not isinstance(samples, list) or not samples:
|
||||
return False
|
||||
|
||||
# 回验 r1/r2:必须是清单路径(存在)或压缩包内嵌名(解压后可解析)
|
||||
valid_paths = {f["path"] for f in features if Path(f["path"]).exists()}
|
||||
archive_members: dict[str, str] = {}
|
||||
for f in features:
|
||||
for m in (f.get("archive_members") or []):
|
||||
archive_members.setdefault(m, f["path"])
|
||||
yaml_samples: dict[str, dict] = {}
|
||||
mounts: set[str] = set()
|
||||
for s in samples:
|
||||
if not isinstance(s, dict):
|
||||
continue
|
||||
r1 = str(s.get("r1") or "").strip()
|
||||
r2 = str(s.get("r2") or "").strip()
|
||||
entry: dict = {}
|
||||
for key, val in (("r1", r1), ("r2", r2)):
|
||||
if not val:
|
||||
continue
|
||||
if val in valid_paths:
|
||||
entry[key] = val
|
||||
mounts.add(str(Path(val).parent))
|
||||
elif val in archive_members:
|
||||
# 解压该压缩包(幂等),用内嵌名构造解出路径
|
||||
arch = Path(archive_members[val])
|
||||
target_dir = job_input_dir / "archives" / (arch.name.rsplit(".", 1)[0])
|
||||
try:
|
||||
import tarfile
|
||||
import zipfile
|
||||
if not target_dir.exists():
|
||||
target_dir.mkdir(parents=True, exist_ok=True)
|
||||
if arch.name.endswith((".tar.gz", ".tar", ".tgz")):
|
||||
with tarfile.open(arch, "r:*") as tf:
|
||||
tf.extractall(target_dir)
|
||||
else:
|
||||
with zipfile.ZipFile(arch) as zf:
|
||||
zf.extractall(target_dir)
|
||||
resolved = target_dir / val
|
||||
if resolved.exists():
|
||||
entry[key] = str(resolved)
|
||||
mounts.add(str(resolved.parent))
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
if not entry:
|
||||
continue
|
||||
for k in ("run", "group_role", "condition", "genotype", "replicate"):
|
||||
v = str(s.get(k) or "").strip()
|
||||
if v:
|
||||
entry[k] = v
|
||||
yaml_samples[str(s.get("name") or Path(entry.get("r1") or "sample").stem)] = entry
|
||||
if not yaml_samples:
|
||||
return False
|
||||
(job_input_dir / "samples.yaml").write_text(
|
||||
_yaml.safe_dump({"samples": yaml_samples}, allow_unicode=True), encoding="utf-8"
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
# ---- runner ----
|
||||
|
||||
ExecutorFn = Callable[..., Awaitable[dict[str, Any]]]
|
||||
|
||||
# 增量日志事件节流:每 N 行或每 S 秒发布一次 stage_log(增量)
|
||||
LOG_EMIT_LINES = 20
|
||||
LOG_EMIT_SECS = 0.5
|
||||
# 进度事件节流(秒):stage_progress 最低发布间隔
|
||||
PROGRESS_THROTTLE_SECS = 2.0
|
||||
# 进度解析白名单:只有这些阶段的 CLI 有 [PROGRESS] 标记(其余阶段不编造数字)
|
||||
PROGRESS_STAGES = ("rnaseq_qc", "rnaseq_sample_qc")
|
||||
|
||||
|
||||
# 阶段进度解析:CLI 日志中的 [PROGRESS] 标记 → 进度增量信息。
|
||||
# 保守策略:只识别 rnaseq 系 CLI 的既有标记,非白名单阶段不编造数字。
|
||||
_PROGRESS_SAMPLE_RE = re.compile(
|
||||
r"\[PROGRESS\]\s*Processing:\s*(.+?)\s*$"
|
||||
)
|
||||
_PROGRESS_COPY_RE = re.compile(
|
||||
r"\[PROGRESS\]\s*Copying\s+results"
|
||||
)
|
||||
|
||||
|
||||
def parse_progress_marker(line: str) -> dict[str, Any] | None:
|
||||
"""解析单行日志中的 [PROGRESS] 标记(纯函数,可单测)。
|
||||
|
||||
- `[PROGRESS] Processing: X` → {"kind": "sample", "message": "X"}
|
||||
- `[PROGRESS] Copying results ...` → {"kind": "copy", "percent": 95}
|
||||
- 其余行 → None
|
||||
"""
|
||||
if "[PROGRESS]" not in line:
|
||||
return None
|
||||
m = _PROGRESS_SAMPLE_RE.search(line)
|
||||
if m:
|
||||
return {"kind": "sample", "message": m.group(1).strip()}
|
||||
if _PROGRESS_COPY_RE.search(line):
|
||||
return {"kind": "copy", "percent": 95}
|
||||
return None
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobRecord:
|
||||
@@ -110,6 +625,8 @@ class JobRecord:
|
||||
auto_continue: bool = False
|
||||
archive_id: str | None = None
|
||||
pause_requested: bool = False
|
||||
registered_server_paths: list[str] = field(default_factory=list) # 登记路径(自动挂载)
|
||||
_running_container: str | None = None # 当前执行容器名(暂停时按名 kill)
|
||||
_stage_states: dict[str, dict[str, Any]] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
@@ -346,17 +863,83 @@ class Runner:
|
||||
self.events.publish("stage_started", {
|
||||
"job_id": record.id, "stage": stage, "attempt": attempt,
|
||||
})
|
||||
attempt_no = int(record._stage_states.get(stage, {}).get("attempt") or 0) + 1
|
||||
record._running_container = f"plantai-bio-{record.id}-{stage}-{attempt_no}"
|
||||
|
||||
# 进度/增量日志状态(每次 attempt 重置)
|
||||
emit_buf: list[str] = []
|
||||
emit_last = time.monotonic()
|
||||
prog_state: dict[str, Any] = {"seen": 0, "total": 0, "last": 0.0}
|
||||
track_progress = stage in PROGRESS_STAGES
|
||||
total_re = re.compile(r"Processing\s+(\d+)\s+samples")
|
||||
valid_re = re.compile(r"Valid\s+samples:\s*(\d+)")
|
||||
|
||||
async def _on_log(chunk: str) -> None:
|
||||
nonlocal emit_last
|
||||
emit_buf.append(chunk)
|
||||
now = time.monotonic()
|
||||
if len(emit_buf) >= LOG_EMIT_LINES or now - emit_last >= LOG_EMIT_SECS:
|
||||
self.events.publish("stage_log", {
|
||||
"job_id": record.id, "stage": stage, "log": "".join(emit_buf),
|
||||
})
|
||||
emit_buf.clear()
|
||||
emit_last = now
|
||||
# [PROGRESS] 解析(白名单内才跟踪)
|
||||
if track_progress:
|
||||
for line in chunk.splitlines():
|
||||
marker = parse_progress_marker(line)
|
||||
if not marker:
|
||||
continue
|
||||
if marker["kind"] == "copy":
|
||||
prog_state["copying"] = True
|
||||
elif marker["kind"] == "sample":
|
||||
prog_state["seen"] += 1
|
||||
prog_state["sample"] = marker["message"]
|
||||
tm = total_re.search(chunk)
|
||||
if tm:
|
||||
prog_state["total"] = int(tm.group(1))
|
||||
vm = valid_re.search(chunk)
|
||||
if vm:
|
||||
prog_state["total"] = int(vm.group(1))
|
||||
if (prog_state["seen"] and prog_state["total"]
|
||||
and now - prog_state["last"] >= PROGRESS_THROTTLE_SECS):
|
||||
prog_state["last"] = now
|
||||
percent = min(int(prog_state["seen"] * 100 / prog_state["total"]), 99)
|
||||
self.events.publish("stage_progress", {
|
||||
"job_id": record.id, "stage": stage,
|
||||
"percent": percent,
|
||||
"message": str(prog_state.get("sample") or ""),
|
||||
"sample": str(prog_state.get("sample") or ""),
|
||||
})
|
||||
|
||||
result = await self.executor(stage, record.input_dir, record.workspace_dir,
|
||||
params, log_path=record.logs_dir / f"{stage}.log")
|
||||
params, log_path=record.logs_dir / f"{stage}.log",
|
||||
extra_mounts=self._registered_mounts(record),
|
||||
container_name=record._running_container,
|
||||
on_log=_on_log)
|
||||
record._running_container = None
|
||||
# 尾块兜底(增量 stage_log 与日志文件对齐)
|
||||
if emit_buf:
|
||||
self.events.publish("stage_log", {
|
||||
"job_id": record.id, "stage": stage, "log": "".join(emit_buf),
|
||||
})
|
||||
log_tail = str(result.get("log_tail") or "")
|
||||
self._set_stage(record, stage, log_tail=log_tail)
|
||||
await self._persist_stage(record, stage)
|
||||
self.events.publish("stage_log", {
|
||||
"job_id": record.id, "stage": stage, "log": log_tail,
|
||||
})
|
||||
if record.pause_requested:
|
||||
# 暂停请求在阶段执行期间到达(容器已被 pause_job kill):
|
||||
# 不按 exit 137 判 failed,走暂停分支
|
||||
self._set_stage(record, stage, status=STAGE_PAUSED, error="")
|
||||
await self._persist_stage(record, stage)
|
||||
await self._pause_job(record, stage)
|
||||
return "paused"
|
||||
outcome = await self._judge_stage(record, stage, result)
|
||||
return outcome
|
||||
finally:
|
||||
record._running_container = None
|
||||
self.semaphore.release()
|
||||
|
||||
async def _acquire(self, record: JobRecord, stage: str) -> bool:
|
||||
@@ -409,6 +992,11 @@ class Runner:
|
||||
self._set_stage(record, stage, status=STAGE_DONE, manifest_path=manifest.path,
|
||||
finished_at=_now_iso())
|
||||
await self._persist_stage(record, stage)
|
||||
# 终态强制 100%(白名单阶段;非白名单阶段前端会因 done 状态显示完成)
|
||||
if stage in PROGRESS_STAGES:
|
||||
self.events.publish("stage_progress", {
|
||||
"job_id": record.id, "stage": stage, "percent": 100, "message": "",
|
||||
})
|
||||
self.events.publish("stage_done", {
|
||||
"job_id": record.id, "stage": stage, "exit_code": exit_code,
|
||||
"manifest_ref": manifest.path,
|
||||
@@ -458,6 +1046,10 @@ class Runner:
|
||||
async def confirm_stage(self, record: JobRecord, stage: str, params: dict[str, Any]) -> None:
|
||||
"""阶段确认(FR-BIO-009):校验参数→存表→执行(并发排队)。"""
|
||||
status = self.stage_status(record, stage)
|
||||
# 幂等:已确认且正在执行/排队(如前端重复点击「继续」),直接放行不报 409
|
||||
st_existing = record._stage_states.get(stage, {})
|
||||
if status in (STAGE_RUNNING, STAGE_QUEUED) and st_existing.get("_confirmed"):
|
||||
return
|
||||
if status not in (STAGE_WAITING_CONFIRM, STAGE_QUEUED, STAGE_WAITING):
|
||||
raise RunnerError("STAGE_NOT_CONFIRMABLE", f"阶段 {stage} 当前状态 {status} 不可确认")
|
||||
allowed = self._allowed_param_names(stage)
|
||||
@@ -572,14 +1164,25 @@ class Runner:
|
||||
logger.exception("runner: skip 后推进异常 %s", record.id)
|
||||
|
||||
async def pause_job(self, record: JobRecord) -> None:
|
||||
"""暂停(FR-BIO-010):等待当前阶段完成后停在阶段边界。"""
|
||||
"""暂停(FR-BIO-010):running 阶段有执行容器时立即 kill,任务即刻暂停。"""
|
||||
if record.status not in (JOB_RUNNING, JOB_AWAIT_INPUTS):
|
||||
raise RunnerError("JOB_NOT_RUNNING", f"任务 {record.id} 当前状态 {record.status} 不可暂停")
|
||||
record.pause_requested = True
|
||||
await self._persist_job(record)
|
||||
# 无 running 阶段 → 立即暂停
|
||||
# 有 running 阶段 → kill 执行容器(立即生效;失败仅日志,不阻断暂停语义)
|
||||
running = [s for s in record.stage_chain
|
||||
if self.stage_status(record, s) == STAGE_RUNNING]
|
||||
if running and getattr(record, "_running_container", None):
|
||||
name = record._running_container
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_shell(
|
||||
f"docker kill {name}",
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.STDOUT,
|
||||
)
|
||||
await asyncio.wait_for(proc.communicate(), timeout=30.0)
|
||||
except Exception as exc: # noqa: BLE001 - kill 失败不阻断(仍走阶段边界暂停)
|
||||
logger.warning("pause_job: docker kill %s 失败: %s", name, exc)
|
||||
if not running:
|
||||
await self._pause_job(record, None)
|
||||
|
||||
@@ -611,6 +1214,29 @@ class Runner:
|
||||
await self._persist_job(record)
|
||||
self.spawn(self._resume_loop(record))
|
||||
|
||||
def _registered_mounts(self, record: JobRecord) -> list[str]:
|
||||
"""登记路径 → 挂载清单:取每个登记文件的父目录(去重)。
|
||||
|
||||
另并入 samples.yaml 引用的 fastq 父目录(描述文件路径校正/实况分析产物)。
|
||||
"""
|
||||
parents: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for p in getattr(record, "registered_server_paths", None) or []:
|
||||
parent = str(Path(p).parent)
|
||||
if parent not in seen:
|
||||
seen.add(parent)
|
||||
parents.append(parent)
|
||||
try:
|
||||
for m in ensure_samples_yaml(record.input_dir, source_roots=[
|
||||
Path(p).parent for p in getattr(record, "registered_server_paths", None) or []
|
||||
]):
|
||||
if m not in seen:
|
||||
seen.add(m)
|
||||
parents.append(m)
|
||||
except Exception: # noqa: BLE001 - 挂载补充失败不阻断
|
||||
pass
|
||||
return parents
|
||||
|
||||
async def retry_after_files(self, record: JobRecord) -> None:
|
||||
"""文件上传/路径登记后的自动续跑(BUG-P7-01:文件引导闭环 FR-BIO-005/006)。
|
||||
|
||||
@@ -680,6 +1306,17 @@ class Runner:
|
||||
auto_continue=bool(row["auto_continue"]),
|
||||
archive_id=row["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"], "params": params_dict(s["params"]),
|
||||
|
||||
@@ -15,12 +15,12 @@ HOST="${BIOFLOW_BIND_HOST:-127.0.0.1}"
|
||||
APP_MODULE="app.main:app"
|
||||
PID=""
|
||||
if command -v lsof >/dev/null 2>&1; then
|
||||
PID="$(lsof -ti tcp:"$PORT" 2>/dev/null || true)"
|
||||
PID="$(lsof -ti tcp:"$PORT" -sTCP:LISTEN 2>/dev/null || true)"
|
||||
fi
|
||||
|
||||
if [[ -n "$PID" ]]; then
|
||||
echo "[restart] 停止 8098 旧进程 pid=$PID"
|
||||
kill "$PID" 2>/dev/null || true
|
||||
kill $PID 2>/dev/null || true
|
||||
sleep 2
|
||||
fi
|
||||
|
||||
|
||||
@@ -57,6 +57,7 @@ STAGE_MAP_JSON = {
|
||||
"note": "跨环境: 主 rnaR; rnaseq_quant 矩阵子步骤拆 bioflow-rnaseq",
|
||||
},
|
||||
"rnaseq_quant": {"image": "bioflow-rnaseq", "module": "cli.run_rnaseq_quant", "executable": True},
|
||||
"rnaseq_qc": {"image": "bioflow-rnaseq", "module": "cli.run_rnaseq_qc", "executable": True},
|
||||
"synteny_plot_revise": {"image": "bioflow-gf", "module": None, "executable": True,
|
||||
"note": "入口脚本缺失"},
|
||||
},
|
||||
@@ -244,6 +245,7 @@ def test_settings(tmp_path, monkeypatch):
|
||||
_env_file=None,
|
||||
data_root=str(tmp_path / "data"),
|
||||
jobs_dir=str(tmp_path / "jobs"),
|
||||
extensions_dir=str(tmp_path / "extensions"),
|
||||
db_path=str(tmp_path / "bioflow.db"),
|
||||
archives_dir=str(tmp_path / "archives"),
|
||||
image_map_path=str(tmp_path / "stage_image_map.json"),
|
||||
@@ -254,7 +256,7 @@ def test_settings(tmp_path, monkeypatch):
|
||||
)
|
||||
Path(s.image_map_abs).parent.mkdir(parents=True, exist_ok=True)
|
||||
Path(s.image_map_abs).write_text(json.dumps(STAGE_MAP_JSON), encoding="utf-8")
|
||||
for d in (s.data_root_abs, s.jobs_dir_abs, s.archives_dir_abs, s.flow_docs_dir_abs):
|
||||
for d in (s.data_root_abs, s.jobs_dir_abs, s.archives_dir_abs, s.flow_docs_dir_abs, s.extensions_dir_abs):
|
||||
Path(d).mkdir(parents=True, exist_ok=True)
|
||||
monkeypatch.setenv("BIOFLOW_TEST_UID", "1234")
|
||||
|
||||
@@ -278,6 +280,7 @@ def _isolated_env(tmp_path, monkeypatch):
|
||||
_env_file=None,
|
||||
data_root=str(tmp_path / "data"),
|
||||
jobs_dir=str(tmp_path / "jobs"),
|
||||
extensions_dir=str(tmp_path / "extensions"),
|
||||
db_path=str(tmp_path / "bioflow.db"),
|
||||
archives_dir=str(tmp_path / "archives"),
|
||||
image_map_path=str(tmp_path / "stage_image_map.json"),
|
||||
|
||||
@@ -288,7 +288,7 @@ def test_catalog_unknown_chain(api):
|
||||
|
||||
def test_chat_stream_plan_ready(api_context):
|
||||
"""TC-BIO-API-001/002:chat/stream SSE,plan_ready 下发候选。"""
|
||||
events = _read_sse(_chat_stream_call("帮我跑水稻 SPL 家族鉴定"), max_events=1)
|
||||
events = _read_sse(_chat_stream_call("帮我跑水稻 SPL 家族鉴定"), max_events=2)
|
||||
types = [e["type"] for e in events]
|
||||
assert "plan_ready" in types
|
||||
plan = [e for e in events if e["type"] == "plan_ready"][0]["data"]
|
||||
@@ -305,7 +305,7 @@ def test_chat_stream_unknown_stage_rejected(api_context, monkeypatch):
|
||||
{"purpose": "坏", "stages": ["fake_stage_xyz"]},
|
||||
{"purpose": "好", "stages": ["identify"]},
|
||||
]))
|
||||
events = _read_sse(_chat_stream_call("帮我跑水稻 SPL 家族鉴定"), max_events=1)
|
||||
events = _read_sse(_chat_stream_call("帮我跑水稻 SPL 家族鉴定"), max_events=2)
|
||||
plan = [e for e in events if e["type"] == "plan_ready"][0]["data"]
|
||||
for c in plan["candidates"]:
|
||||
assert "fake_stage_xyz" not in c["stages"]
|
||||
@@ -343,7 +343,7 @@ def test_chat_stream_llm_timeout_degrade(api_context, monkeypatch):
|
||||
from conftest import FakeLLM
|
||||
import app.services.bioflow.planner as planner_mod
|
||||
monkeypatch.setattr(planner_mod, "complete_json", FakeLLM(error="timeout"))
|
||||
events = _read_sse(_chat_stream_call("帮我跑水稻 SPL 家族鉴定"), max_events=1)
|
||||
events = _read_sse(_chat_stream_call("帮我跑水稻 SPL 家族鉴定"), max_events=2)
|
||||
plan = [e for e in events if e["type"] == "plan_ready"][0]["data"]
|
||||
assert plan["fallback"] # 降级提示
|
||||
assert plan["candidates"]
|
||||
@@ -529,7 +529,7 @@ def test_files_upload_no_job(api):
|
||||
|
||||
|
||||
def test_path_register_valid_and_whitelist(api, tmp_path):
|
||||
"""TC-BIO-API-042/045:路径登记合法 + 白名单外 403。"""
|
||||
"""TC-BIO-API-042/045:路径登记合法 + 敏感黑名单 403。"""
|
||||
resp = _create_job(api)
|
||||
job_id = _job_id(resp)
|
||||
# api_context 的 settings 已指向 tmp_path/data(conftest _isolated_env 派生)
|
||||
@@ -542,11 +542,13 @@ def test_path_register_valid_and_whitelist(api, tmp_path):
|
||||
f.write(">x\nACGT")
|
||||
ok = api.post(f"/api/bioflow/jobs/{job_id}/files", json={"server_path": allowed})
|
||||
assert ok.status_code == 200
|
||||
assert ok.json()["data"]["source"] == "server_path"
|
||||
# 白名单外
|
||||
outside = str(tmp_path / "outside.fa")
|
||||
body = ok.json()["data"]
|
||||
assert body["items"][0]["source"] == "server_path"
|
||||
assert body["items"][0]["name"] == "SPL.fa"
|
||||
# 敏感黑名单(.env 文件名形态)
|
||||
outside = str(tmp_path / ".env")
|
||||
with open(outside, "w") as f:
|
||||
f.write("x")
|
||||
f.write("SECRET=1")
|
||||
forbidden = api.post(f"/api/bioflow/jobs/{job_id}/files", json={"server_path": outside})
|
||||
assert forbidden.status_code == 403
|
||||
assert forbidden.json()["error"]["code"] == "PATH_NOT_ALLOWED"
|
||||
@@ -557,6 +559,39 @@ def test_path_register_valid_and_whitelist(api, tmp_path):
|
||||
assert missing.json()["error"]["code"] == "PATH_NOT_EXIST"
|
||||
|
||||
|
||||
def test_path_register_directory_expand(api, tmp_path):
|
||||
"""目录登记:递归展开多文件,逐个 items;挂载清单落 server_registered.json。"""
|
||||
resp = _create_job(api)
|
||||
job_id = _job_id(resp)
|
||||
d = tmp_path / "rnaseq_fastq"
|
||||
(d / "s1").mkdir(parents=True)
|
||||
(d / "s2").mkdir(parents=True)
|
||||
(d / "s1" / "CRR1_1.fastq.gz").write_text("g1", encoding="utf-8")
|
||||
(d / "s1" / "CRR1_2.fastq.gz").write_text("g2", encoding="utf-8")
|
||||
(d / "s2" / "CRR2_1.fastq.gz").write_text("g3", encoding="utf-8")
|
||||
(d / "notes.txt").write_text("n", encoding="utf-8") # 允许扩展名内
|
||||
(d / "skip.exe").write_text("x", encoding="utf-8") # 扩展名白名单外,跳过
|
||||
ok = api.post(f"/api/bioflow/jobs/{job_id}/files", json={"server_path": str(d)})
|
||||
assert ok.status_code == 200
|
||||
body = ok.json()["data"]
|
||||
names = {it["name"] for it in body["items"]}
|
||||
assert names == {"CRR1_1.fastq.gz", "CRR1_2.fastq.gz", "CRR2_1.fastq.gz", "notes.txt"}
|
||||
assert body["total"] == 4
|
||||
# 挂载清单落盘
|
||||
import json as _json
|
||||
from pathlib import Path as _Path
|
||||
from app.config import get_settings as _gs
|
||||
job_dir = _Path(_gs().jobs_dir_abs)
|
||||
reg = _json.loads((job_dir / job_id / "server_registered.json").read_text(encoding="utf-8"))
|
||||
assert len(reg["registered_server_paths"]) >= 4 # 4 文件 + samples.yaml 引用的父目录
|
||||
# 空目录报错
|
||||
empty = tmp_path / "empty_dir"
|
||||
empty.mkdir()
|
||||
err = api.post(f"/api/bioflow/jobs/{job_id}/files", json={"server_path": str(empty)})
|
||||
assert err.status_code == 400
|
||||
assert err.json()["error"]["code"] == "VALIDATION"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stages 控制 + manifest
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -625,6 +660,58 @@ def test_manifest_unknown_stage(api):
|
||||
assert m.status_code == 404
|
||||
|
||||
|
||||
def test_stage_logs_pagination(api, api_context):
|
||||
"""TC-BIO-API-102:logs 分页端点(offset/limit 切片、eof、log_tail、超尾)。"""
|
||||
api_context.executor.log = "line0\nline1\nline2\nline3\n"
|
||||
resp = _create_job(api, auto_continue=True)
|
||||
job_id = _job_id(resp)
|
||||
|
||||
def done():
|
||||
return _job_stage_status(api, job_id) == "done"
|
||||
_wait_until(done, timeout=5)
|
||||
|
||||
r1 = api.get(f"/api/bioflow/jobs/{job_id}/stages/identify/logs?offset=0&limit=1")
|
||||
assert r1.status_code == 200
|
||||
d1 = r1.json()["data"]
|
||||
assert d1["lines"] == ["line0"]
|
||||
assert d1["eof"] is False
|
||||
# 增量衔接:offset=1&limit=2
|
||||
r2 = api.get(f"/api/bioflow/jobs/{job_id}/stages/identify/logs?offset=1&limit=2")
|
||||
d2 = r2.json()["data"]
|
||||
assert d2["lines"] == ["line1", "line2"]
|
||||
assert d2["eof"] is False
|
||||
# 到尾:offset=2&limit=500 → eof:true
|
||||
r3 = api.get(f"/api/bioflow/jobs/{job_id}/stages/identify/logs?offset=2&limit=500")
|
||||
d3 = r3.json()["data"]
|
||||
assert d3["lines"] == ["line2", "line3"]
|
||||
assert d3["eof"] is True
|
||||
# 超尾:offset=1000 → 空 lines + eof:true + log_tail 仍在
|
||||
r4 = api.get(f"/api/bioflow/jobs/{job_id}/stages/identify/logs?offset=1000&limit=500")
|
||||
d4 = r4.json()["data"]
|
||||
assert d4["lines"] == []
|
||||
assert d4["eof"] is True
|
||||
assert d4["log_tail"]
|
||||
|
||||
|
||||
def test_stage_logs_no_file_and_404(api):
|
||||
"""TC-BIO-API-103:log 文件不存在 → 空 lines + eof;未知 job/阶段 404。"""
|
||||
resp = _create_job(api) # 手动模式:identify 停在 waiting_confirm,从未执行
|
||||
job_id = _job_id(resp)
|
||||
|
||||
def waiting():
|
||||
return _job_stage_status(api, job_id) == "waiting_confirm"
|
||||
_wait_until(waiting, timeout=5)
|
||||
|
||||
r = api.get(f"/api/bioflow/jobs/{job_id}/stages/identify/logs")
|
||||
assert r.status_code == 200
|
||||
d = r.json()["data"]
|
||||
assert d["lines"] == [] and d["eof"] is True and d["total_bytes"] == 0
|
||||
# 未知 job 404
|
||||
assert api.get("/api/bioflow/jobs/nopejob/stages/identify/logs").status_code == 404
|
||||
# 未知阶段 404
|
||||
assert api.get(f"/api/bioflow/jobs/{job_id}/stages/nope/logs").status_code == 404
|
||||
|
||||
|
||||
def test_rerun_and_skip_on_failed(api, api_context):
|
||||
"""TC-BIO-API-062/065:rerun/skip 失败阶段。"""
|
||||
api_context.executor.exit_code = 1
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user