"""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")