forked from admin/PlantAI-rag-new
Merge branch 'debug/backend2' into feature/bioflow-conversation-orchestration
This commit is contained in:
Executable → Regular
+69
-1
@@ -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
|
||||
Executable → Regular
+68
@@ -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
|
||||
@@ -3,7 +3,8 @@
|
||||
BFF 直连 DeepSeek 生成 300 字内带 [资料n] 引用的 markdown 短答:
|
||||
- ask_private(question, *, top_k=8, min_score=0.30) -> dict
|
||||
检索 chunk 索引 → 空则 {hit:false, mode:"empty", answer:"", items:[]}(不调 LLM)
|
||||
→ 同 doc 按 score 序去重(每 doc ≤3 块)再取 top_k
|
||||
→ 先按 doc_id 聚合(doc 级 max score 排序取 top 文档)再取各 doc 内
|
||||
最相关块(每 doc ≤3,缺陷 #52c)
|
||||
→ 上下文拼 `[资料n]《name》(相关度 x.xx)\nsnippet`(12k 预算截断)
|
||||
→ complete_text(system 约束, temperature=0.2, max_tokens=800, timeout=40)
|
||||
→ {hit:true, mode:"vector", answer, items:[...带 ref], elapsed_ms}
|
||||
@@ -65,24 +66,73 @@ def _safe_float(value: Any) -> bool:
|
||||
|
||||
|
||||
def _pick_items(items: list[dict[str, Any]], top_k: int) -> list[dict[str, Any]]:
|
||||
"""同 doc 按 score 序去重(每 doc ≤ PRIVATE_QA_PER_DOC_MAX 块),再取 top_k。
|
||||
"""先按 doc_id 聚合再取块(缺陷 #52c):单一文档不再占满全部名额。
|
||||
|
||||
输入已按 score 降序(vstore faiss 顺序);同 doc 靠前块优先。
|
||||
原实现按输入 score 序逐条取(每 doc ≤ PRIVATE_QA_PER_DOC_MAX 块)直到
|
||||
top_k——重复上传的作者页块(得分几乎相同)会占满 top-8、正文块进不了
|
||||
上下文。新语义:
|
||||
1. 按 doc_id 分组,doc 级得分 = 组内 max score;
|
||||
2. doc 级得分降序取前 top_k 个文档;
|
||||
3. 每选中文档取组内得分最高的 ≤ PRIVATE_QA_PER_DOC_MAX 个块,
|
||||
且组内近重复块(内容指纹相同,重复上传残留产物)只取 1 条——
|
||||
否则「同内容 200 条作者页块」仍会把该 doc 的 3 个名额占满、
|
||||
正文块永远进不了上下文;
|
||||
4. 全部选中块按 score 降序稳定排序、截断到 top_k 返回(_build_context
|
||||
的 [资料n] 编号与 items 顺序不变,回归语义保留)。
|
||||
"""
|
||||
picked: list[dict[str, Any]] = []
|
||||
per_doc: dict[str, int] = {}
|
||||
groups: dict[str, list[dict[str, Any]]] = {}
|
||||
doc_score: dict[str, float] = {}
|
||||
doc_order: list[str] = []
|
||||
for item in items:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
doc_id = str(item.get("doc_id") or "")
|
||||
used = per_doc.get(doc_id, 0)
|
||||
if used >= PRIVATE_QA_PER_DOC_MAX:
|
||||
continue
|
||||
per_doc[doc_id] = used + 1
|
||||
picked.append(dict(item))
|
||||
if len(picked) >= top_k:
|
||||
break
|
||||
return picked
|
||||
if doc_id not in doc_order:
|
||||
doc_order.append(doc_id)
|
||||
groups.setdefault(doc_id, []).append(item)
|
||||
try:
|
||||
score = float(item.get("score"))
|
||||
except (TypeError, ValueError):
|
||||
score = 0.0
|
||||
if score > doc_score.get(doc_id, float("-inf")):
|
||||
doc_score[doc_id] = score
|
||||
# 先聚合后排序:得分相同时按输入首现顺序(稳定)
|
||||
top_docs = sorted(doc_order, key=lambda d: doc_score.get(d, 0.0), reverse=True)[:top_k]
|
||||
for doc_id in top_docs:
|
||||
group = sorted(groups.get(doc_id, []), key=_score_of, reverse=True)
|
||||
taken_fingerprints: set[str] = set()
|
||||
for item in group:
|
||||
fingerprint = _chunk_fingerprint(item)
|
||||
if fingerprint in taken_fingerprints:
|
||||
continue # 组内近重复块(重复上传残留)不重复占名额
|
||||
taken_fingerprints.add(fingerprint)
|
||||
picked.append(dict(item))
|
||||
if len(taken_fingerprints) >= PRIVATE_QA_PER_DOC_MAX:
|
||||
break
|
||||
picked.sort(key=_score_of, reverse=True)
|
||||
return picked[:top_k]
|
||||
|
||||
|
||||
def _chunk_fingerprint(item: dict[str, Any]) -> str:
|
||||
"""组内近重复指纹:归一化(去空白)后的 snippet 前 48 字符。
|
||||
|
||||
重复上传残留的块内容逐字相同(同内容多次入库切块结果一致);
|
||||
正常相邻块(重叠 1 段)不会整段相同,不受影响。
|
||||
"""
|
||||
import re
|
||||
|
||||
snippet = str(item.get("snippet") or "")
|
||||
normalized = re.sub(r"\s+", "", snippet)
|
||||
return normalized[:48]
|
||||
|
||||
|
||||
def _score_of(item: dict[str, Any]) -> float:
|
||||
"""item 得分(无/非法 → 0.0)。"""
|
||||
try:
|
||||
return float(item.get("score"))
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def _build_context(picked: list[dict[str, Any]]) -> tuple[str, list[dict[str, Any]]]:
|
||||
@@ -119,7 +169,11 @@ async def ask_private(
|
||||
异常(DocgenUnavailable / Exception / LLMError)原样上抛。
|
||||
"""
|
||||
started = time.perf_counter()
|
||||
payload = await search_chunks(question, top_k=top_k, min_score=min_score, timeout=12.0)
|
||||
# 缺陷 #52c:先聚合再取块需要更宽的候选池——只取 top_k 条时重复文档
|
||||
# 的块会把池占满、聚合阶段无块可选。取 top_k×4(vstore 上限 20 clamp),
|
||||
# 最终 _pick_items 仍收敛到 top_k 条,响应量与 [资料n] 编号语义不变。
|
||||
pool_k = min(max(top_k * 4, top_k), 20)
|
||||
payload = await search_chunks(question, top_k=pool_k, min_score=min_score, timeout=12.0)
|
||||
items = [i for i in payload.get("items") or [] if isinstance(i, dict)]
|
||||
if not items:
|
||||
return {"hit": False, "mode": "empty", "answer": "", "items": []}
|
||||
|
||||
@@ -0,0 +1,183 @@
|
||||
"""test_privateqa_pick.py —— _pick_items 先按 doc_id 聚合再取块(缺陷 #52c,验收 §6 第 3 条)。
|
||||
|
||||
运行:cd sdk && /home/xzq/5090/teamWorkspace/zh/plantai-rag-new-venv/bin/python -m pytest app/services/test_privateqa_pick.py -v
|
||||
(venv = plantai-rag-new-venv;sdk 的 app/ 布局与 pyproject testpaths 一致)
|
||||
|
||||
构造合成数据(不真跑 FAISS):
|
||||
- 4 个 doc_id 各若干「作者页块」(高分、分接近)+ 1 条低分「正文块」。
|
||||
- 旧实现会把高分的作者页块占满 top-8、正文块进不了上下文;
|
||||
新实现按 doc 级 max score 取 top 文档 → 各 doc 只占 ≤3 块,
|
||||
正文块(本组内最高低分块)能进入 picked。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from app.services.privateqa import (
|
||||
PRIVATE_QA_PER_DOC_MAX,
|
||||
_build_context,
|
||||
_pick_items,
|
||||
)
|
||||
|
||||
|
||||
def make_items(doc_specs):
|
||||
"""doc_specs: [(doc_id, [(score, snippet), ...]), ...] → item 列表(输入序打散,
|
||||
模拟 vstore 全局 score 序)。"""
|
||||
items = []
|
||||
for doc_id, chunks in doc_specs:
|
||||
for score, snippet in chunks:
|
||||
items.append({"doc_id": doc_id, "score": score, "snippet": snippet, "doc_name": f"资料{doc_id}"})
|
||||
return sorted(items, key=lambda it: -it["score"])
|
||||
|
||||
|
||||
def test_dup_docs_no_longer_monopolize_topk():
|
||||
"""#52c 主场景:4 个 doc_id 各 200 条近同分作者页块(同 doc 内内容相同,
|
||||
模拟重复上传残留——同一页面被反复切块/复传)+ 正文块 →
|
||||
选出的 items 覆盖多个 doc 且含正文块。"""
|
||||
specs = []
|
||||
for i in range(4):
|
||||
author = [(0.94 - i * 0.001, f"作者页块 {i}") for _ in range(200)]
|
||||
body = [(0.30, f"正文块 {i} 方法与结论")]
|
||||
specs.append((f"doc-{i}", author + body))
|
||||
items = make_items(specs)
|
||||
picked = _pick_items(items, top_k=8)
|
||||
|
||||
doc_ids = [it["doc_id"] for it in picked]
|
||||
assert len(set(doc_ids)) >= 2, f"重复文档仍占满名额:{doc_ids}"
|
||||
for doc_id in set(doc_ids):
|
||||
assert doc_ids.count(doc_id) <= PRIVATE_QA_PER_DOC_MAX
|
||||
snippets = [it["snippet"] for it in picked]
|
||||
assert any("正文块" in s for s in snippets), "正文块未进入上下文"
|
||||
# 总分 ≤ top_k;块按分数降序
|
||||
assert len(picked) <= 8
|
||||
scores = [it["score"] for it in picked]
|
||||
assert scores == sorted(scores, reverse=True)
|
||||
|
||||
|
||||
def test_normal_case_semantics_preserved():
|
||||
"""回归:正常场景 top_k/每 doc 上限语义不变。"""
|
||||
items = [
|
||||
{"doc_id": "A", "score": 0.95, "snippet": "a1"},
|
||||
{"doc_id": "A", "score": 0.90, "snippet": "a2"},
|
||||
{"doc_id": "B", "score": 0.85, "snippet": "b1"},
|
||||
{"doc_id": "C", "score": 0.80, "snippet": "c1"},
|
||||
{"doc_id": "D", "score": 0.75, "snippet": "d1"},
|
||||
]
|
||||
picked = _pick_items(items, top_k=5)
|
||||
assert len(picked) == 5
|
||||
for doc_id in {it["doc_id"] for it in picked}:
|
||||
assert sum(1 for it in picked if it["doc_id"] == doc_id) <= PRIVATE_QA_PER_DOC_MAX
|
||||
# 每 doc 内按分数降序取 ≤3 块
|
||||
a = [it for it in picked if it["doc_id"] == "A"]
|
||||
assert [it["snippet"] for it in a] == ["a1", "a2"]
|
||||
|
||||
|
||||
def test_single_doc_caps_at_per_doc_max():
|
||||
"""回归:单一 doc 大量块时仍 ≤ per-doc 上限。"""
|
||||
items = [{"doc_id": "X", "score": 0.9 - i * 0.001, "snippet": f"x{i}"} for i in range(20)]
|
||||
picked = _pick_items(items, top_k=8)
|
||||
assert 1 <= len(picked) <= PRIVATE_QA_PER_DOC_MAX
|
||||
|
||||
|
||||
def test_ties_are_stable():
|
||||
"""回归:同分并列时文档选择稳定(按输入首现顺序),结果确定。"""
|
||||
items = make_items(
|
||||
[
|
||||
("A", [(0.9, "a1"), (0.9, "a2")]),
|
||||
("B", [(0.9, "b1")]),
|
||||
("C", [(0.9, "c1")]),
|
||||
]
|
||||
)
|
||||
p1 = _pick_items(items, top_k=4)
|
||||
p2 = _pick_items(items, top_k=4)
|
||||
assert [it["doc_id"] for it in p1] == [it["doc_id"] for it in p2]
|
||||
assert len(p1) == 4
|
||||
|
||||
|
||||
def test_malformed_items_ignored():
|
||||
"""回归:非 dict / 缺 doc_id 条目不炸。"""
|
||||
picked = _pick_items(
|
||||
[None, "bad", {"score": 0.9, "snippet": "s"}, {"doc_id": "Z", "score": 0.5, "snippet": "z"}],
|
||||
top_k=4,
|
||||
)
|
||||
assert picked
|
||||
|
||||
|
||||
def test_build_context_integration():
|
||||
"""回归:_build_context 12k 预算截断 + [资料n] 引用编号与 items 顺序一致。"""
|
||||
picked = _pick_items(
|
||||
[
|
||||
{"doc_id": "A", "score": 0.95, "snippet": "正文一"},
|
||||
{"doc_id": "A", "score": 0.90, "snippet": "正文二"},
|
||||
{"doc_id": "B", "score": 0.80, "snippet": "正文三"},
|
||||
],
|
||||
top_k=3,
|
||||
)
|
||||
context, out_items = _build_context(picked)
|
||||
assert "[资料1]" in context
|
||||
assert out_items[0]["ref"] == 1
|
||||
assert [it["snippet"] for it in out_items] == [it["snippet"] for it in picked]
|
||||
|
||||
|
||||
async def test_ask_private_fetches_wider_pool(monkeypatch):
|
||||
"""#52c:ask_private 以 top_k×4 向 vstore 取候选池,再收敛到 top_k。"""
|
||||
import app.services.privateqa as pq
|
||||
|
||||
calls = {}
|
||||
|
||||
async def fake_search(query, *, top_k, min_score, timeout):
|
||||
calls["requested_top_k"] = top_k
|
||||
items = []
|
||||
for doc_i in range(4):
|
||||
for j in range(200):
|
||||
items.append(
|
||||
{
|
||||
"doc_id": f"doc-{doc_i}",
|
||||
"score": 0.94 - doc_i * 0.001,
|
||||
# 同 doc 内作者页块内容相同(重复上传残留的真实形态)
|
||||
"snippet": f"作者页块 {doc_i}",
|
||||
"doc_name": f"资料{doc_i}",
|
||||
}
|
||||
)
|
||||
items.append(
|
||||
{
|
||||
"doc_id": f"doc-{doc_i}",
|
||||
"score": 0.30,
|
||||
"snippet": f"正文块 {doc_i} 方法与结论",
|
||||
"doc_name": f"资料{doc_i}",
|
||||
}
|
||||
)
|
||||
return {"ok": True, "mode": "vector", "items": items}
|
||||
|
||||
async def fake_llm(system, user, **kwargs):
|
||||
return "方法与结论见[资料1]。"
|
||||
|
||||
monkeypatch.setattr(pq, "search_chunks", fake_search)
|
||||
monkeypatch.setattr(pq, "complete_text", fake_llm)
|
||||
|
||||
result = await pq.ask_private("总结这篇论文的主要方法和结论", top_k=8)
|
||||
assert calls["requested_top_k"] == 20, "候选池应取 min(top_k×4, vstore 上限 20)"
|
||||
assert result["hit"] is True
|
||||
assert result["mode"] == "vector"
|
||||
doc_ids = [it["doc_id"] for it in result["items"]]
|
||||
assert len(set(doc_ids)) >= 2
|
||||
assert any("正文块" in (it["snippet"] or "") for it in result["items"])
|
||||
assert len(result["items"]) <= 8
|
||||
assert all(it["ref"] >= 1 for it in result["items"])
|
||||
|
||||
|
||||
async def test_ask_private_pool_capped_at_20(monkeypatch):
|
||||
"""#52c:候选池 = min(top_k×4, 20)(vstore CHUNK_TOP_K_MAX clamp)。"""
|
||||
import app.services.privateqa as pq
|
||||
|
||||
calls = {}
|
||||
|
||||
async def fake_search(query, *, top_k, min_score, timeout):
|
||||
calls["requested_top_k"] = top_k
|
||||
return {"ok": True, "mode": "vector", "items": []}
|
||||
|
||||
monkeypatch.setattr(pq, "search_chunks", fake_search)
|
||||
await pq.ask_private("问题", top_k=8)
|
||||
assert calls["requested_top_k"] == 20
|
||||
await pq.ask_private("问题", top_k=5)
|
||||
assert calls["requested_top_k"] == 20 # 5×4=20
|
||||
await pq.ask_private("问题", top_k=3)
|
||||
assert calls["requested_top_k"] == 12 # 3×4=12(低于 20)
|
||||
@@ -0,0 +1,123 @@
|
||||
"""test_vault_dedup.py —— 内容 SHA256 去重(缺陷 #52b,验收 §6 第 2 条 + 回归 G3)。
|
||||
|
||||
运行:cd vstore && /home/xzq/5090/teamWorkspace/zh/plantai-rag-new-venv/bin/python -m pytest tests/test_vault_dedup.py -v
|
||||
(或 vstore/pyproject.toml testpaths 下全量 pytest;venv = plantai-rag-new-venv)
|
||||
|
||||
G3 断言:
|
||||
- 同内容 add 两次 → 库中该内容仅 1 个 doc_id、chunk 总数不翻倍;
|
||||
- 不同内容 add → 2 个 doc_id 互不覆盖;
|
||||
- 同名不同内容 → 按内容哈希不误判(保留 2 个 doc_id)。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
|
||||
import pytest
|
||||
|
||||
from vstore.vault import Vault
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def vault(tmp_path, embed_fn, monkeypatch):
|
||||
from vstore.settings import VstorePaths, set_paths
|
||||
|
||||
set_paths(VstorePaths(root=str(tmp_path), vault_path="", index_path="", chunk_path=""))
|
||||
# 桩掉 chunk 批量嵌入(vault._chunked_embed 动态查找),不触网真 Ollama
|
||||
dim = 1024
|
||||
monkeypatch.setattr(
|
||||
"vstore.embed.embed_batch_chunked",
|
||||
lambda texts, *a, **k: [[1.0 / dim] * dim for _ in texts],
|
||||
)
|
||||
return Vault(tmp_path / "vault", tmp_path / "index")
|
||||
|
||||
|
||||
def _doc_ids_for_content(vault, content: bytes) -> list[str]:
|
||||
"""库中 content_sha == sha256(content) 的 doc 行 id 列表(磁盘事实)。"""
|
||||
sha = hashlib.sha256(content).hexdigest()
|
||||
import json
|
||||
|
||||
rows = []
|
||||
try:
|
||||
data = json.loads(vault.index.json_path.read_text(encoding="utf-8"))
|
||||
items = data.get("items", []) if isinstance(data, dict) else []
|
||||
except (ValueError, OSError):
|
||||
items = []
|
||||
for item in items:
|
||||
if isinstance(item, dict) and str(item.get("content_sha") or "") == sha:
|
||||
rows.append(item)
|
||||
return [r.get("id") for r in rows]
|
||||
|
||||
|
||||
def _make_pdf(text: str) -> bytes:
|
||||
"""真实最小 PDF 字节(pymupdf 可解析、含目标文本)。"""
|
||||
import pymupdf
|
||||
|
||||
doc = pymupdf.open()
|
||||
page = doc.new_page()
|
||||
page.insert_text((72, 72), text)
|
||||
return doc.tobytes()
|
||||
|
||||
|
||||
def test_same_content_added_twice_keeps_one_doc_id(vault):
|
||||
"""#52b:同内容 add 两次 → 1 个 doc_id、chunk 数不翻倍、旧密文被清。"""
|
||||
content = _make_pdf("同一份论文内容。")
|
||||
first = vault.add(content, "paper.pdf", {"name": "论文", "category": "资料"})
|
||||
assert first["content_sha"] == hashlib.sha256(content).hexdigest()
|
||||
chunks_after_first = vault.chunks.count()
|
||||
assert chunks_after_first >= 1
|
||||
|
||||
second = vault.add(content, "paper.pdf", {"name": "论文", "category": "资料"})
|
||||
|
||||
ids = _doc_ids_for_content(vault, content)
|
||||
assert ids == [second["document_id"]], f"库中该内容 doc_id 应为仅新 doc:{ids}"
|
||||
assert second["document_id"] != first["document_id"]
|
||||
# 旧密文被清理(覆盖语义不留两份)
|
||||
assert not (vault.vault_dir / f"{first['document_id']}.enc").exists()
|
||||
assert (vault.vault_dir / f"{second['document_id']}.enc").exists()
|
||||
# chunk 总数不翻倍(旧 doc_id 的 chunk 已清理)
|
||||
assert vault.chunks.count() == chunks_after_first
|
||||
# chunk 行全部属于新 doc_id(旧 doc_id 不再出现在检索结果)
|
||||
assert {r["doc_id"] for r in vault.chunks.rows()} == {second["document_id"]}
|
||||
|
||||
|
||||
def test_duplicate_reuses_latest_meta(vault):
|
||||
"""#52b:覆盖语义——重复上传 meta 以本次为准(created_at 刷新、name 覆盖)。"""
|
||||
content = "内容甲。".encode("utf-8")
|
||||
first = vault.add(content, "a.txt", {"name": "旧名"})
|
||||
second = vault.add(content, "a.txt", {"name": "新名"})
|
||||
rows = vault.list_rows()
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["name"] == "新名"
|
||||
assert rows[0]["id"] == second["document_id"]
|
||||
assert first["document_id"] not in [r["id"] for r in rows]
|
||||
|
||||
|
||||
def test_different_content_keeps_two_doc_ids(vault):
|
||||
"""G3-1:不同内容上传 → 2 个 doc_id 互不覆盖。"""
|
||||
a = vault.add("内容甲。".encode("utf-8"), "a.txt", {"name": "甲"})
|
||||
b = vault.add("内容乙。".encode("utf-8"), "b.txt", {"name": "乙"})
|
||||
assert a["document_id"] != b["document_id"]
|
||||
assert len(vault.list_rows()) == 2
|
||||
assert set(_doc_ids_for_content(vault, "内容甲。".encode("utf-8"))) == {a["document_id"]}
|
||||
assert set(_doc_ids_for_content(vault, "内容乙。".encode("utf-8"))) == {b["document_id"]}
|
||||
|
||||
|
||||
def test_same_name_different_content_not_deduped(vault):
|
||||
"""G3-2:同名不同内容 → 按 SHA256 不误判(保留 2 个 doc_id)。"""
|
||||
a = vault.add(_make_pdf("内容甲。"), "same.pdf", {"name": "同名"})
|
||||
b = vault.add(_make_pdf("内容乙。"), "same.pdf", {"name": "同名"})
|
||||
assert a["document_id"] != b["document_id"]
|
||||
assert len(vault.list_rows()) == 2
|
||||
assert vault.chunks.count() >= 2
|
||||
|
||||
|
||||
def test_duplicate_dedup_persists_across_reload(vault, tmp_path, monkeypatch):
|
||||
"""#52b:去重依据(content_sha)落盘持久——重启后(新实例 load)仍不判重。"""
|
||||
content = "持久化内容。".encode("utf-8")
|
||||
vault.add(content, "a.txt", {"name": "甲"})
|
||||
vault2 = Vault(tmp_path / "vault", tmp_path / "index")
|
||||
vault2.index.load()
|
||||
vault2.chunks.load()
|
||||
item = vault2.add(content, "a.txt", {"name": "甲"})
|
||||
assert _doc_ids_for_content(vault2, content) == [item["document_id"]]
|
||||
assert vault2.chunks.count() >= 1
|
||||
+89
-11
@@ -130,22 +130,22 @@ class Vault:
|
||||
|
||||
# ---- 增删改查 ----
|
||||
|
||||
def add(
|
||||
def _add_file(
|
||||
self,
|
||||
doc_id: str,
|
||||
file_bytes: bytes,
|
||||
filename: str,
|
||||
meta: dict[str, Any],
|
||||
as_template: bool = False,
|
||||
name: str,
|
||||
category: str,
|
||||
tags: str,
|
||||
description: str,
|
||||
as_template: bool,
|
||||
) -> dict[str, Any]:
|
||||
"""入库一条。原文只在函数内流转,末尾丢弃引用。"""
|
||||
doc_id = str(uuid.uuid4())
|
||||
name = str(meta.get("name") or filename or "未命名")
|
||||
category = str(meta.get("category") or "")
|
||||
tags = str(meta.get("tags") or "")
|
||||
description = str(meta.get("description") or "")
|
||||
"""入库一条(覆盖语义已由调用方保证:同内容 doc 已被移除)。
|
||||
|
||||
if not file_bytes:
|
||||
raise VaultError("上传文件为空")
|
||||
与旧 add() 唯一差异是 doc_id 由调用方传入;函数内失败照常抛 VaultError,
|
||||
由调用方决定失败回滚语义。
|
||||
"""
|
||||
try:
|
||||
extracted = extract(file_bytes, filename)
|
||||
except ParsingError as exc:
|
||||
@@ -162,6 +162,7 @@ class Vault:
|
||||
"description": description,
|
||||
"as_template": bool(as_template),
|
||||
"size": len(file_bytes),
|
||||
"content_sha": hashlib.sha256(bytes(file_bytes)).hexdigest(),
|
||||
"created_at": _now_iso(),
|
||||
"embed_chunks": 0,
|
||||
}
|
||||
@@ -252,6 +253,7 @@ class Vault:
|
||||
"description": description,
|
||||
"as_template": bool(as_template),
|
||||
"enc_path": str(enc_path),
|
||||
"content_sha": doc_meta["content_sha"],
|
||||
"embed_sha": embed_sha,
|
||||
"embed_chunks": embed_chunks,
|
||||
# FR-PRIV-002:chunk 级索引结果(失败不阻断 doc 级)
|
||||
@@ -260,6 +262,75 @@ class Vault:
|
||||
"chunk_error": chunk_error,
|
||||
}
|
||||
|
||||
def _find_content_duplicate(self, content_sha: str) -> dict[str, Any] | None:
|
||||
"""库中同内容哈希的 doc 行(无则 None)。
|
||||
|
||||
内容哈希存于 index.json 每行的 content_sha 字段(#52b 新增);历史行
|
||||
无该字段(字段缺失/不一致时绝不判重)。直接读 index.json 而非内存
|
||||
self.index._rows——覆盖语义要求与磁盘一致(同内容密文先删后建)。
|
||||
"""
|
||||
try:
|
||||
data = json.loads(self.index.json_path.read_text(encoding="utf-8"))
|
||||
items = data.get("items", []) if isinstance(data, dict) else []
|
||||
except (ValueError, OSError):
|
||||
return None
|
||||
for item in items:
|
||||
if isinstance(item, dict) and str(item.get("content_sha") or "") == content_sha:
|
||||
return item
|
||||
return None
|
||||
|
||||
def add(
|
||||
self,
|
||||
file_bytes: bytes,
|
||||
filename: str,
|
||||
meta: dict[str, Any],
|
||||
as_template: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""入库一条(内容 SHA256 去重,缺陷 #52b)。原文只在函数内流转,末尾丢弃引用。
|
||||
|
||||
去重语义(验收 §6 第 2 条 + 回归 G3):
|
||||
- 内容哈希 = SHA256(原文件字节)。同内容重复上传 → 移除旧 doc 的密文/
|
||||
检索行/chunk 行 → 落一条新 doc;该内容在库中始终仅 1 个 doc_id,
|
||||
chunk 总数不翻倍。
|
||||
- 旧行清理先于新 doc 落盘;新 doc 建库失败时旧 doc 已删(覆盖语义,
|
||||
失败由调用方重传即可,绝无同内容双份)。
|
||||
- 新 doc 的 meta 以本次上传为准(name/category/tags/description 覆盖,
|
||||
created_at 刷新);as_template 以本次为准。
|
||||
"""
|
||||
doc_id = str(uuid.uuid4())
|
||||
name = str(meta.get("name") or filename or "未命名")
|
||||
category = str(meta.get("category") or "")
|
||||
tags = str(meta.get("tags") or "")
|
||||
description = str(meta.get("description") or "")
|
||||
|
||||
if not file_bytes:
|
||||
raise VaultError("上传文件为空")
|
||||
|
||||
# ---- 内容去重(#52b):同内容哈希 → 先清旧 doc 再落新 doc ----
|
||||
content_sha = hashlib.sha256(bytes(file_bytes)).hexdigest()
|
||||
dup = self._find_content_duplicate(content_sha)
|
||||
if dup:
|
||||
old_doc_id = str(dup.get("id") or "")
|
||||
if old_doc_id and old_doc_id != doc_id:
|
||||
try:
|
||||
self.index.remove_item(old_doc_id)
|
||||
except Exception as exc:
|
||||
logger.warning("vault: 去重移除旧检索行失败(%s):%s", old_doc_id, exc)
|
||||
old_enc = self.vault_dir / f"{old_doc_id}.enc"
|
||||
if old_enc.exists():
|
||||
try:
|
||||
secure_delete(old_enc)
|
||||
except OSError as exc:
|
||||
logger.warning("vault: 去重移除旧密文失败(%s):%s", old_enc, exc)
|
||||
try:
|
||||
self.chunks.remove_document(old_doc_id)
|
||||
except Exception as exc:
|
||||
logger.warning("vault: 去重移除旧 chunk 行失败(%s):%s", old_doc_id, exc)
|
||||
|
||||
return self._add_file(
|
||||
doc_id, file_bytes, filename, name, category, tags, description, as_template
|
||||
)
|
||||
|
||||
def get_decrypted(self, doc_id: str) -> dict[str, Any] | None:
|
||||
"""取解密后的文档结构(meta + extracted)。不存在返回 None。"""
|
||||
enc_path = self.vault_dir / f"{doc_id}.enc"
|
||||
@@ -362,6 +433,13 @@ class Vault:
|
||||
row = dict(meta)
|
||||
# keywords 只用 meta,不含正文派生的 title(明文纪律)
|
||||
row["keywords"] = _meta_to_keywords(meta)
|
||||
# content_sha(#52b):历史密文 meta 若缺失则以密文内容补算(与入库时一致),
|
||||
# 保证重建后去重依据不丢
|
||||
if not row.get("content_sha"):
|
||||
try:
|
||||
row["content_sha"] = hashlib.sha256(load_encrypted(enc)).hexdigest()
|
||||
except (CryptoError, VaultError, OSError):
|
||||
row["content_sha"] = ""
|
||||
items.append((row, vectors[0]))
|
||||
try:
|
||||
count = self.index.rebuild(items)
|
||||
|
||||
Reference in New Issue
Block a user