212 lines
7.7 KiB
Python
212 lines
7.7 KiB
Python
|
# ============================================================
|
|||
|
|
# 知识文档导入与切块（moduleId: knowledge-ingest, 可重生 ✅）
|
|||
|
|
# plan.md §8.2：PDF/docx/md/txt → chunk → 知识资产（版本化+审批）
|
|||
|
|
# ============================================================
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import re
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
# 切块目标长度（字符）与重叠
|
|||
|
|
_CHUNK_SIZE = 500
|
|||
|
|
_CHUNK_OVERLAP = 50
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _ext(path: str) -> str:
|
|||
|
|
return os.path.splitext(path)[1].lower()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def extract_text_from_bytes(filename: str, raw: bytes) -> list[dict[str, Any]]:
|
|||
|
|
"""按扩展名解析文件为带页/段元数据的文本块列表（尚未切 chunk）。
|
|||
|
|
|
|||
|
|
Returns: [{text, page?, heading?}, ...]
|
|||
|
|
"""
|
|||
|
|
name = filename or "upload.txt"
|
|||
|
|
ext = _ext(name)
|
|||
|
|
if ext in (".md", ".txt", ".markdown"):
|
|||
|
|
text = raw.decode("utf-8", errors="replace")
|
|||
|
|
return _split_markdown_sections(text)
|
|||
|
|
if ext == ".pdf":
|
|||
|
|
return _extract_pdf(raw)
|
|||
|
|
if ext in (".docx",):
|
|||
|
|
return _extract_docx(raw)
|
|||
|
|
if ext in (".xlsx", ".xlsm"):
|
|||
|
|
return _extract_xlsx(raw)
|
|||
|
|
# 兜底当纯文本
|
|||
|
|
try:
|
|||
|
|
return [{"text": raw.decode("utf-8"), "page": 1}]
|
|||
|
|
except UnicodeDecodeError as exc:
|
|||
|
|
raise ValueError(f"不支持的文件类型或编码：{name}") from exc
|
|||
|
|
|
|||
|
|
|
|||
|
|
def extract_text_from_path(path: str) -> list[dict[str, Any]]:
|
|||
|
|
"""从本地路径读取并解析。"""
|
|||
|
|
if not os.path.isfile(path):
|
|||
|
|
raise FileNotFoundError(f"文件不存在：{path}")
|
|||
|
|
with open(path, "rb") as f:
|
|||
|
|
raw = f.read()
|
|||
|
|
return extract_text_from_bytes(os.path.basename(path), raw)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _split_markdown_sections(text: str) -> list[dict[str, Any]]:
|
|||
|
|
parts: list[dict[str, Any]] = []
|
|||
|
|
current_heading = ""
|
|||
|
|
buf: list[str] = []
|
|||
|
|
for line in text.splitlines():
|
|||
|
|
if re.match(r"^#{1,3}\s+", line):
|
|||
|
|
if buf:
|
|||
|
|
parts.append({"text": "\n".join(buf).strip(), "heading": current_heading or None})
|
|||
|
|
buf = []
|
|||
|
|
current_heading = re.sub(r"^#{1,3}\s+", "", line).strip()
|
|||
|
|
buf.append(line)
|
|||
|
|
else:
|
|||
|
|
buf.append(line)
|
|||
|
|
if buf:
|
|||
|
|
parts.append({"text": "\n".join(buf).strip(), "heading": current_heading or None})
|
|||
|
|
return [p for p in parts if p["text"]]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _extract_pdf(raw: bytes) -> list[dict[str, Any]]:
|
|||
|
|
try:
|
|||
|
|
import pdfplumber
|
|||
|
|
except ImportError as exc:
|
|||
|
|
raise ValueError("解析 PDF 需要安装 pdfplumber：pip install pdfplumber") from exc
|
|||
|
|
from io import BytesIO
|
|||
|
|
pages: list[dict[str, Any]] = []
|
|||
|
|
with pdfplumber.open(BytesIO(raw)) as pdf:
|
|||
|
|
for i, page in enumerate(pdf.pages, start=1):
|
|||
|
|
t = (page.extract_text() or "").strip()
|
|||
|
|
if t:
|
|||
|
|
pages.append({"text": t, "page": i})
|
|||
|
|
if not pages:
|
|||
|
|
raise ValueError("PDF 未提取到文本（可能是扫描件）")
|
|||
|
|
return pages
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _extract_docx(raw: bytes) -> list[dict[str, Any]]:
|
|||
|
|
try:
|
|||
|
|
from docx import Document
|
|||
|
|
except ImportError as exc:
|
|||
|
|
raise ValueError("解析 docx 需要安装 python-docx：pip install python-docx") from exc
|
|||
|
|
from io import BytesIO
|
|||
|
|
doc = Document(BytesIO(raw))
|
|||
|
|
paras = [p.text.strip() for p in doc.paragraphs if p.text and p.text.strip()]
|
|||
|
|
if not paras:
|
|||
|
|
raise ValueError("docx 未提取到文本")
|
|||
|
|
return [{"text": "\n".join(paras), "page": 1}]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _extract_xlsx(raw: bytes) -> list[dict[str, Any]]:
|
|||
|
|
from io import BytesIO
|
|||
|
|
from openpyxl import load_workbook
|
|||
|
|
wb = load_workbook(BytesIO(raw), data_only=True, read_only=True)
|
|||
|
|
parts: list[dict[str, Any]] = []
|
|||
|
|
for sheet in wb.worksheets:
|
|||
|
|
rows = []
|
|||
|
|
for row in sheet.iter_rows(values_only=True):
|
|||
|
|
cells = [str(c) for c in row if c is not None and str(c).strip()]
|
|||
|
|
if cells:
|
|||
|
|
rows.append(" | ".join(cells))
|
|||
|
|
if rows:
|
|||
|
|
parts.append({"text": f"# {sheet.title}\n" + "\n".join(rows), "heading": sheet.title})
|
|||
|
|
if not parts:
|
|||
|
|
raise ValueError("xlsx 无有效单元格")
|
|||
|
|
return parts
|
|||
|
|
|
|||
|
|
|
|||
|
|
def chunk_sections(sections: list[dict[str, Any]], *,
|
|||
|
|
size: int = _CHUNK_SIZE, overlap: int = _CHUNK_OVERLAP) -> list[dict[str, Any]]:
|
|||
|
|
"""把段落/页切成固定长度 chunk（带重叠）。"""
|
|||
|
|
chunks: list[dict[str, Any]] = []
|
|||
|
|
idx = 0
|
|||
|
|
for sec in sections:
|
|||
|
|
text = (sec.get("text") or "").strip()
|
|||
|
|
if not text:
|
|||
|
|
continue
|
|||
|
|
heading = sec.get("heading")
|
|||
|
|
page = sec.get("page")
|
|||
|
|
if len(text) <= size:
|
|||
|
|
chunks.append({
|
|||
|
|
"chunkId": f"c{idx:04d}", "text": text,
|
|||
|
|
"heading": heading, "page": page, "seq": idx,
|
|||
|
|
})
|
|||
|
|
idx += 1
|
|||
|
|
continue
|
|||
|
|
start = 0
|
|||
|
|
while start < len(text):
|
|||
|
|
end = min(start + size, len(text))
|
|||
|
|
piece = text[start:end].strip()
|
|||
|
|
if piece:
|
|||
|
|
chunks.append({
|
|||
|
|
"chunkId": f"c{idx:04d}", "text": piece,
|
|||
|
|
"heading": heading, "page": page, "seq": idx,
|
|||
|
|
})
|
|||
|
|
idx += 1
|
|||
|
|
if end >= len(text):
|
|||
|
|
break
|
|||
|
|
start = max(end - overlap, start + 1)
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
|
|||
|
|
def preview_ingest(filename: str, raw: bytes | None = None, path: str | None = None,
|
|||
|
|
*, kind: str = "sop", title: str | None = None) -> dict[str, Any]:
|
|||
|
|
"""预览导入：解析+切块，不写库。"""
|
|||
|
|
if path:
|
|||
|
|
sections = extract_text_from_path(path)
|
|||
|
|
fname = os.path.basename(path)
|
|||
|
|
elif raw is not None:
|
|||
|
|
sections = extract_text_from_bytes(filename, raw)
|
|||
|
|
fname = filename
|
|||
|
|
else:
|
|||
|
|
raise ValueError("需要提供 path 或 raw")
|
|||
|
|
chunks = chunk_sections(sections)
|
|||
|
|
title = title or os.path.splitext(fname)[0]
|
|||
|
|
return {
|
|||
|
|
"filename": fname,
|
|||
|
|
"title": title,
|
|||
|
|
"kind": kind,
|
|||
|
|
"sectionCount": len(sections),
|
|||
|
|
"chunkCount": len(chunks),
|
|||
|
|
"preview": [{"seq": c["seq"], "heading": c.get("heading"), "page": c.get("page"),
|
|||
|
|
"chars": len(c["text"]), "snippet": c["text"][:80]} for c in chunks[:8]],
|
|||
|
|
"chunks": chunks,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def confirmation_for_knowledge_import(preview: dict[str, Any]) -> tuple[str, list[str]]:
|
|||
|
|
"""P2 确认卡文案。"""
|
|||
|
|
title = f"导入知识文档《{preview.get('title')}》"
|
|||
|
|
lines = [
|
|||
|
|
f"文件：{preview.get('filename')}",
|
|||
|
|
f"类型：{preview.get('kind')}",
|
|||
|
|
f"切块：{preview.get('chunkCount')} 段（来自 {preview.get('sectionCount')} 节/页）",
|
|||
|
|
"入库后可检索；默认已审批（approved=true）。",
|
|||
|
|
]
|
|||
|
|
return title, lines
|
|||
|
|
|
|||
|
|
|
|||
|
|
def apply_knowledge_import(preview: dict[str, Any], *, tags: list[str] | None = None,
|
|||
|
|
store: Any = None) -> dict[str, Any]:
|
|||
|
|
"""把预览结果写入知识库（带 chunks）。"""
|
|||
|
|
from server.knowledge.assets import get_knowledge
|
|||
|
|
kb = store or get_knowledge()
|
|||
|
|
chunks = preview.get("chunks") or []
|
|||
|
|
content = "\n\n".join(
|
|||
|
|
(f"【{c.get('heading') or '段落'}】\n" if c.get("heading") else "") + c["text"]
|
|||
|
|
for c in chunks
|
|||
|
|
)
|
|||
|
|
if not content.strip():
|
|||
|
|
raise ValueError("无正文可入库")
|
|||
|
|
asset = kb.add_with_chunks(
|
|||
|
|
kind=str(preview.get("kind") or "sop"),
|
|||
|
|
title=str(preview.get("title") or preview.get("filename") or "导入文档"),
|
|||
|
|
content=content[:8000], # 资产摘要正文（检索仍优先 chunk）
|
|||
|
|
chunks=chunks,
|
|||
|
|
tags=tags or ["导入", str(preview.get("filename") or "")],
|
|||
|
|
source=str(preview.get("filename") or ""),
|
|||
|
|
approved=True,
|
|||
|
|
)
|
|||
|
|
return {"assetId": asset["assetId"], "title": asset["title"],
|
|||
|
|
"version": asset["version"], "chunkCount": len(asset.get("chunks") or [])}
|