309 lines
13 KiB
Python
309 lines
13 KiB
Python
"""kb 入库后台任务:sources/ 待入库原件 → docs/*.md + 摘要 + INDEX 行。
|
||
|
||
流水线(每份原件独立,失败不连坐):
|
||
1. markitdown Python API 抽文本(pdf/docx/pptx/xlsx/html/txt...)
|
||
2. PDF 文本近零(扫描件)→ 方舟 seed-2.0-lite 文档理解 OCR 兜底
|
||
(照 tools/read_document.py 同款请求体;ARK_API_KEY 未配则该件报错留待)
|
||
3. deepseek flash 单次 chat 写 标题/摘要/关键词(照 core/context_fold.py 范式;
|
||
LLM 失败降级为文件名 + 正文开头,入库不阻塞)
|
||
4. 追加 INDEX 行(core.kb.format_index_line)——写入即"已入库"
|
||
5. record_chat_usage(kind="kb_ingest", task_id=None, units 带 {"kb","source"})
|
||
|
||
编排照定时执行器范式:web 层 asyncio.create_task(asyncio.to_thread(run_ingest, ...)),
|
||
本模块全同步;per-(user,库) threading.Lock 非阻塞抢占 —— 抢不到 = 已在入库,直接返回
|
||
(上传即触发 + 手动触发天然去重)。进度状态存内存 dict 供前端轮询;崩溃丢状态无妨,
|
||
"待入库"判据在文件系统(INDEX 缺口),重触发即续跑。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import hashlib
|
||
import re
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
from uuid import UUID
|
||
|
||
from core.ark_client import ArkClient, ArkError
|
||
from core.capabilities import ModelCapabilities
|
||
from core.kb import format_index_line, kb_dir, parse_index, pending_sources
|
||
from core.file_store import atomic_write_text
|
||
from core.kb_lock import KbBusyError, kb_is_locked, kb_mutation_lock
|
||
from core.llm import LLM
|
||
from core.storage.usage import record_chat_usage, record_vision_usage
|
||
|
||
# 摘要模型固定走最便宜档(与 web FALLBACK_MODEL_PROFILE 同值;单次几千 token,不随任务模型)
|
||
SUMMARY_PROFILE = "deepseek_v4.flash"
|
||
# 判定"文本近零"(扫描件)的阈值:markitdown 抽出的字符数低于此值即认为无文本层
|
||
_MIN_TEXT_CHARS = 120
|
||
# 喂给摘要模型的正文截断长度(开头段落通常含标题/摘要/目录,足够定性)
|
||
_SUMMARY_INPUT_CHARS = 5000
|
||
|
||
_OCR_QUESTION = (
|
||
"这是一份多页 PDF 文档。请逐页把其中的文字完整 OCR 成 markdown:"
|
||
"每页以「== 第N页 ==」开头;表格转成 markdown 表格;保留标题层级与段落换行;"
|
||
"公式尽量用 LaTeX;不要总结、不要遗漏、不要自行补充原文没有的内容。"
|
||
)
|
||
|
||
_SUMMARY_PROMPT = """\
|
||
下面是文档「{filename}」的正文开头(可能被截断)。请为它写入库索引条目,严格按以下三行格式输出,不要输出任何其他文字:
|
||
标题：<文档标题,没有明确标题就概括一个,不超过 40 字>
|
||
摘要：<两三句话说清这份文档讲什么、有什么用,不超过 120 字>
|
||
关键词：<3-6 个检索关键词,用逗号分隔>
|
||
|
||
正文开头:
|
||
{content}"""
|
||
|
||
# ── 进度状态(内存,供 API 轮询) ─────────────────────────────────────
|
||
_status: Dict[Tuple[str, str], Dict[str, Any]] = {}
|
||
|
||
|
||
def ingest_status(
|
||
user_id: UUID, name: str, workspace_dir: Optional[Path] = None
|
||
) -> Dict[str, Any]:
|
||
"""当前/最近一次入库进度(无记录返回 idle);前端轮询 + 库详情附带。"""
|
||
st = _status.get((str(user_id), name))
|
||
out = dict(st) if st else {"running": False, "total": 0, "done": 0, "errors": []}
|
||
if workspace_dir is not None and not out.get("running"):
|
||
out["running"] = kb_is_locked(workspace_dir, user_id, name)
|
||
return out
|
||
|
||
|
||
# ── 转换 ─────────────────────────────────────────────────────────────
|
||
|
||
def _extract_text(src: Path) -> str:
|
||
"""markitdown 抽文本;不支持的格式 / 解析失败抛 ValueError(留待重试)。"""
|
||
from markitdown import MarkItDown
|
||
try:
|
||
result = MarkItDown(enable_plugins=False).convert(str(src))
|
||
except Exception as e:
|
||
raise ValueError(f"markitdown 转换失败: {type(e).__name__}: {e}")
|
||
return (result.text_content or "").strip()
|
||
|
||
|
||
def _ocr_pdf(src: Path, *, user_id: UUID, kb_name: str) -> str:
|
||
"""扫描件 PDF 走方舟文档理解 OCR(照 tools/read_document.py 请求体)。
|
||
|
||
失败抛 ValueError。记账 kind=vision、task_id=None、units 带 kb/source 溯源。
|
||
"""
|
||
from core.ark_client import ArkConfig
|
||
ark_cfg = ArkConfig.load()
|
||
vision_cfg = (ark_cfg.raw.get("vision") or {}) if ark_cfg else {}
|
||
vis_key, cfg = "", None
|
||
for k, v in vision_cfg.items():
|
||
if isinstance(v, dict):
|
||
vis_key, cfg = k, v
|
||
break
|
||
if ark_cfg is None or cfg is None:
|
||
raise ValueError("扫描件 PDF 需 OCR,但方舟(ARK_API_KEY / vision 段)未配置")
|
||
|
||
max_bytes = int(float(cfg.get("max_pdf_mb", 30)) * 1024 * 1024)
|
||
size = src.stat().st_size
|
||
if size > max_bytes:
|
||
raise ValueError(f"PDF {size / 1e6:.1f}MB 超过 OCR 上限 {max_bytes / 1e6:.0f}MB")
|
||
|
||
data_url = "data:application/pdf;base64," + base64.b64encode(src.read_bytes()).decode()
|
||
body = {
|
||
"model": cfg["model_id"],
|
||
"messages": [{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": _OCR_QUESTION},
|
||
{"type": "file", "file": {"filename": src.name, "file_data": data_url}},
|
||
],
|
||
}],
|
||
}
|
||
timeout_s = float(cfg.get("doc_request_timeout_s", 600))
|
||
try:
|
||
with ArkClient(ark_cfg, timeout_s=timeout_s) as client:
|
||
resp = client.post_json(cfg.get("endpoint", "/chat/completions"), body, timeout_s=timeout_s)
|
||
except ArkError as e:
|
||
raise ValueError(f"OCR 调用失败: {e}")
|
||
|
||
choices = resp.get("choices") or []
|
||
msg = (choices[0].get("message") if choices and isinstance(choices[0], dict) else None) or {}
|
||
content = msg.get("content")
|
||
if isinstance(content, list):
|
||
content = "\n".join(
|
||
c.get("text", "") for c in content if isinstance(c, dict) and c.get("type") == "text"
|
||
)
|
||
text = (content or "").strip() if isinstance(content, (str,)) else ""
|
||
if not text:
|
||
raise ValueError("OCR 响应无文本(PDF 损坏或页面超像素上限)")
|
||
|
||
usage = resp.get("usage") or {}
|
||
try:
|
||
record_vision_usage(
|
||
task_id=None,
|
||
user_id=user_id,
|
||
model_profile=f"doubao.{vis_key}",
|
||
prompt_tokens=int(usage.get("prompt_tokens", 0) or 0),
|
||
completion_tokens=int(usage.get("completion_tokens", 0) or 0),
|
||
input_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_input", 0)),
|
||
output_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_output", 0)),
|
||
extra_units={"kb": kb_name, "source": src.name},
|
||
)
|
||
except Exception:
|
||
pass # 记账失败不阻塞入库(与 loop 同纪律)
|
||
return text
|
||
|
||
|
||
# ── 摘要 ─────────────────────────────────────────────────────────────
|
||
|
||
def _parse_summary(raw: str) -> Tuple[str, str, str]:
|
||
"""解析三行格式;缺行返回空串由调用方兜底。"""
|
||
title = summary = keywords = ""
|
||
for line in raw.splitlines():
|
||
line = line.strip().lstrip("*# ")
|
||
m = re.match(r"^(标题|摘要|关键词)[:：]\s*(.*)$", line)
|
||
if not m:
|
||
continue
|
||
key, val = m.group(1), m.group(2).strip()
|
||
if key == "标题" and not title:
|
||
title = val
|
||
elif key == "摘要" and not summary:
|
||
summary = val
|
||
elif key == "关键词" and not keywords:
|
||
keywords = val
|
||
return title, summary, keywords
|
||
|
||
|
||
def _summarize(
|
||
llm: LLM, caps: ModelCapabilities, *, text: str, filename: str,
|
||
user_id: UUID, kb_name: str,
|
||
) -> Tuple[str, str, str]:
|
||
"""单次 chat 产 (标题, 摘要, 关键词);失败降级文件名 + 正文开头,不阻塞入库。"""
|
||
fallback = (Path(filename).stem, text[:100].replace("\n", " "), "")
|
||
try:
|
||
response = llm.chat(messages=[{
|
||
"role": "user",
|
||
"content": _SUMMARY_PROMPT.format(filename=filename, content=text[:_SUMMARY_INPUT_CHARS]),
|
||
}], tools=None)
|
||
choices = getattr(response, "choices", None) or []
|
||
raw = ((choices[0].message.content if choices else "") or "").strip()
|
||
except Exception as e:
|
||
print(f"[kb_ingest] summary llm failed ({filename}): {type(e).__name__}: {e}", flush=True)
|
||
return fallback
|
||
usage = getattr(response, "usage", None)
|
||
try:
|
||
record_chat_usage(
|
||
task_id=None,
|
||
user_id=user_id,
|
||
message_id=None,
|
||
model_profile=f"{caps.family}.{caps.variant}",
|
||
prompt_tokens=getattr(usage, "prompt_tokens", 0) or 0,
|
||
completion_tokens=getattr(usage, "completion_tokens", 0) or 0,
|
||
input_cny_per_mtoken=caps.input_cny_per_mtoken,
|
||
output_cny_per_mtoken=caps.output_cny_per_mtoken,
|
||
response=response,
|
||
kind="kb_ingest",
|
||
extra_units={"kb": kb_name, "source": filename},
|
||
)
|
||
except Exception:
|
||
pass
|
||
title, summary, keywords = _parse_summary(raw)
|
||
return (title or fallback[0], summary or fallback[1], keywords)
|
||
|
||
|
||
# ── 主流程 ───────────────────────────────────────────────────────────
|
||
|
||
def _doc_name_for(d: Path, source_name: str) -> str:
|
||
"""source 原件名 → docs/ 下的 md 文件名;与既有他源文档撞名时加短 hash 后缀。"""
|
||
stem = Path(source_name).stem
|
||
stem = re.sub(r'[\\/:*?"<>|｜\[\]()]+', "_", stem).strip("._ ") or "doc"
|
||
stem = stem[:80]
|
||
entries = parse_index((d / "INDEX.md").read_text(encoding="utf-8")) if (d / "INDEX.md").is_file() else []
|
||
doc_rel = f"docs/{stem}.md"
|
||
for e in entries:
|
||
if e["doc"] == doc_rel and e["source"] != f"sources/{source_name}":
|
||
suffix = hashlib.md5(source_name.encode("utf-8")).hexdigest()[:6]
|
||
return f"{stem}_{suffix}.md"
|
||
return f"{stem}.md"
|
||
|
||
|
||
def _append_index(d: Path, line: str) -> None:
|
||
idx = d / "INDEX.md"
|
||
text = idx.read_text(encoding="utf-8") if idx.is_file() else ""
|
||
if text and not text.endswith("\n"):
|
||
text += "\n"
|
||
atomic_write_text(idx, text + line + "\n")
|
||
|
||
|
||
def run_ingest(workspace_dir: Path, user_id: UUID, kb_name: str, models_dir: Path) -> bool:
|
||
"""同步跑一轮入库(web 层用 asyncio.to_thread 下沉)。
|
||
|
||
返回 False = 没跑(锁被占,已有入库在进行 / 库不存在);True = 跑完(含空转)。
|
||
幂等:只处理 INDEX 缺口,单件失败记入 status.errors 并保持待入库,下次重触发续跑。
|
||
"""
|
||
d = kb_dir(workspace_dir, user_id, kb_name)
|
||
if d is None or not d.is_dir():
|
||
return False
|
||
key = (str(user_id), kb_name)
|
||
try:
|
||
with kb_mutation_lock(workspace_dir, user_id, kb_name):
|
||
return _run_ingest_locked(d, key, user_id, kb_name, models_dir)
|
||
except KbBusyError:
|
||
return False
|
||
|
||
|
||
def _run_ingest_locked(
|
||
d: Path,
|
||
key: Tuple[str, str],
|
||
user_id: UUID,
|
||
kb_name: str,
|
||
models_dir: Path,
|
||
) -> bool:
|
||
try:
|
||
pending = pending_sources(d)
|
||
st: Dict[str, Any] = {
|
||
"running": True, "total": len(pending), "done": 0,
|
||
"current": None, "errors": [],
|
||
"started_at": datetime.now().isoformat(timespec="seconds"),
|
||
}
|
||
_status[key] = st
|
||
if not pending:
|
||
return True
|
||
|
||
llm: Optional[LLM] = None
|
||
caps: Optional[ModelCapabilities] = None
|
||
try:
|
||
caps = ModelCapabilities.load(SUMMARY_PROFILE, models_dir)
|
||
llm = LLM(caps)
|
||
except Exception as e:
|
||
print(f"[kb_ingest] summary model load failed: {type(e).__name__}: {e}", flush=True)
|
||
|
||
for name in pending:
|
||
st["current"] = name
|
||
src = d / "sources" / name
|
||
try:
|
||
text = _extract_text(src)
|
||
if len(text) < _MIN_TEXT_CHARS and src.suffix.lower() == ".pdf":
|
||
text = _ocr_pdf(src, user_id=user_id, kb_name=kb_name)
|
||
if len(text) < _MIN_TEXT_CHARS:
|
||
raise ValueError(f"抽出文本过少({len(text)} 字符),无法入库")
|
||
doc_name = _doc_name_for(d, name)
|
||
(d / "docs").mkdir(parents=True, exist_ok=True)
|
||
atomic_write_text(d / "docs" / doc_name, text)
|
||
if llm is not None and caps is not None:
|
||
title, summary, keywords = _summarize(
|
||
llm, caps, text=text, filename=name, user_id=user_id, kb_name=kb_name
|
||
)
|
||
else:
|
||
title, summary, keywords = Path(name).stem, text[:100].replace("\n", " "), ""
|
||
_append_index(d, format_index_line(
|
||
title=title, doc=f"docs/{doc_name}", source=f"sources/{name}",
|
||
summary=summary, keywords=keywords,
|
||
))
|
||
st["done"] += 1
|
||
except Exception as e:
|
||
msg = str(e) if isinstance(e, ValueError) else f"{type(e).__name__}: {e}"
|
||
st["errors"].append({"source": name, "error": msg})
|
||
print(f"[kb_ingest] {kb_name}/{name} failed: {msg}", flush=True)
|
||
return True
|
||
finally:
|
||
st = _status.get(key)
|
||
if st is not None:
|
||
st["running"] = False
|
||
st["current"] = None
|
||
st["finished_at"] = datetime.now().isoformat(timespec="seconds")
|