151 lines
4.9 KiB
Python
151 lines
4.9 KiB
Python
|
"""
|
||
|
|
一次性批量导入脚本：将 doco/deliverables/ 下的 22 期融合A稿 docx 文件
|
||
|
|
导入知识库（knowledge_items + knowledge_embeddings）。
|
||
|
|
|
||
|
|
运行方式：
|
||
|
|
cd backend && python -m scripts.import_doco_transcripts
|
||
|
|
"""
|
||
|
|
|
||
|
|
import re
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from docx import Document
|
||
|
|
from sqlmodel import Session, select
|
||
|
|
|
||
|
|
from app.db.session import engine
|
||
|
|
from app.models.knowledge import KnowledgeItem
|
||
|
|
from app.services.knowledge_service import KnowledgeService
|
||
|
|
|
||
|
|
# ── 路径 ──────────────────────────────────────────────────────
|
||
|
|
# backend/scripts/ → backend/ → 项目根目录
|
||
|
|
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
|
||
|
|
DOCO_DIR = PROJECT_ROOT / "doco" / "deliverables"
|
||
|
|
|
||
|
|
# ── 文件名正则 ────────────────────────────────────────────────
|
||
|
|
# 第{NN}期_{YYYYMMDD}_{节目名}_{编导名}_融合A稿.docx
|
||
|
|
FILENAME_RE = re.compile(
|
||
|
|
r"^第(\d{2,3})期_(\d{8})_(.+?)_(.+?)_融合A稿\.docx$"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def parse_filename(filename: str) -> dict | None:
|
||
|
|
"""从文件名解析期次号、播出日期、节目名、编导名。不匹配返回 None。"""
|
||
|
|
m = FILENAME_RE.match(filename)
|
||
|
|
if not m:
|
||
|
|
return None
|
||
|
|
issue_num = int(m.group(1))
|
||
|
|
date_raw = m.group(2) # YYYYMMDD
|
||
|
|
title = m.group(3)
|
||
|
|
author = m.group(4)
|
||
|
|
broadcast_date = f"{date_raw[:4]}-{date_raw[4:6]}-{date_raw[6:8]}"
|
||
|
|
return {
|
||
|
|
"issue_num": issue_num,
|
||
|
|
"broadcast_date": broadcast_date,
|
||
|
|
"title": title,
|
||
|
|
"author": author,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def extract_docx_text(file_path: Path) -> str:
|
||
|
|
"""用 python-docx 提取 docx 正文纯文本，段落间用 \\n\\n 分隔。"""
|
||
|
|
doc = Document(str(file_path))
|
||
|
|
paragraphs = [p.text for p in doc.paragraphs if p.text.strip()]
|
||
|
|
return "\n\n".join(paragraphs)
|
||
|
|
|
||
|
|
|
||
|
|
def build_md_bytes(title: str, author: str, broadcast_date: str, body: str) -> bytes:
|
||
|
|
"""组装带 YAML frontmatter 的 md 格式 bytes，供 store_md_file() 使用。"""
|
||
|
|
md_content = (
|
||
|
|
f"---\n"
|
||
|
|
f"类型: 节目文稿\n"
|
||
|
|
f"名称: {title}\n"
|
||
|
|
f"编导: {author}\n"
|
||
|
|
f"播出日期: {broadcast_date}\n"
|
||
|
|
f"---\n"
|
||
|
|
f"{body}"
|
||
|
|
)
|
||
|
|
return md_content.encode("utf-8")
|
||
|
|
|
||
|
|
|
||
|
|
def check_exists(source_file_name: str) -> bool:
|
||
|
|
"""查 knowledge_items 表，判断 source_file_name 是否已存在。"""
|
||
|
|
with Session(engine) as session:
|
||
|
|
stmt = select(KnowledgeItem).where(
|
||
|
|
KnowledgeItem.source_file_name == source_file_name
|
||
|
|
)
|
||
|
|
existing = session.exec(stmt).first()
|
||
|
|
return existing is not None
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
# 确认 docx 目录存在
|
||
|
|
if not DOCO_DIR.is_dir():
|
||
|
|
print(f"✗ 目录不存在: {DOCO_DIR}")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
# 收集并排序 docx 文件
|
||
|
|
docx_files = sorted(DOCO_DIR.glob("*.docx"), key=lambda p: p.name)
|
||
|
|
if not docx_files:
|
||
|
|
print("✗ 未找到任何 .docx 文件")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
print(f"共发现 {len(docx_files)} 个 docx 文件，开始导入...\n")
|
||
|
|
|
||
|
|
service = KnowledgeService()
|
||
|
|
success_count = 0
|
||
|
|
skip_count = 0
|
||
|
|
fail_count = 0
|
||
|
|
total = len(docx_files)
|
||
|
|
|
||
|
|
for idx, file_path in enumerate(docx_files, 1):
|
||
|
|
filename = file_path.name
|
||
|
|
|
||
|
|
# 解析文件名
|
||
|
|
meta = parse_filename(filename)
|
||
|
|
if meta is None:
|
||
|
|
print(f"[{idx}/{total}] ✗ {filename} — 文件名格式不匹配，跳过")
|
||
|
|
fail_count += 1
|
||
|
|
continue
|
||
|
|
|
||
|
|
display_label = f"第{meta['issue_num']:02d}期 {meta['title']}"
|
||
|
|
|
||
|
|
# 防重复：用去掉 .docx 后缀加 .md 的文件名
|
||
|
|
source_file_name = filename.replace(".docx", ".md")
|
||
|
|
|
||
|
|
try:
|
||
|
|
if check_exists(source_file_name):
|
||
|
|
print(f"[{idx}/{total}] ⊘ {display_label} — 已存在，跳过")
|
||
|
|
skip_count += 1
|
||
|
|
continue
|
||
|
|
|
||
|
|
# 提取正文
|
||
|
|
body_text = extract_docx_text(file_path)
|
||
|
|
char_count = len(body_text)
|
||
|
|
|
||
|
|
# 组装 md bytes
|
||
|
|
md_bytes = build_md_bytes(
|
||
|
|
title=meta["title"],
|
||
|
|
author=meta["author"],
|
||
|
|
broadcast_date=meta["broadcast_date"],
|
||
|
|
body=body_text,
|
||
|
|
)
|
||
|
|
|
||
|
|
# 调用入库链路（含 embedding API 调用）
|
||
|
|
service.store_md_file(file_content=md_bytes, file_name=source_file_name)
|
||
|
|
|
||
|
|
print(f"[{idx}/{total}] ✓ {display_label} — 入库成功（{char_count}字）")
|
||
|
|
success_count += 1
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
print(f"[{idx}/{total}] ✗ {display_label} — 失败：{e}")
|
||
|
|
fail_count += 1
|
||
|
|
|
||
|
|
# 汇总
|
||
|
|
print(f"\n{'='*50}")
|
||
|
|
print(f"导入完成：成功 {success_count} 篇 / 跳过 {skip_count} 篇 / 失败 {fail_count} 篇")
|
||
|
|
print(f"{'='*50}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|