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 [])}
|