File size: 4,852 Bytes
793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 b6c412f 793e953 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """
知识库初始化脚本 - 递归扫描 veterinary_articles 目录下所有 PDF
使用方法：
py -3.12 scripts/init_knowledge_base.py
"""
import os
import sys
import shutil
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from app.database import SessionLocal, init_db
from app.models.knowledge import KnowledgeDocument, KnowledgeChunk
from app.services import knowledge_service as ks
from app.config import settings
from app.utils.logger import logger
# 扫描目录：本地 veterinary_articles（含 site1/site2/site3）
ARTICLES_DIR = Path(r"C:\Users\HP\Documents\666666666666666666\veterinary_articles")
PDF_DEST_DIR = Path(settings.UPLOAD_DIR) / "knowledge_pdfs"
EXCLUDE_PATTERNS = [
"cover", "index", "contents", "front", "preface", "acknowledg",
"abbreviat", "glossary", "appendi", "reference", "bibliograph",
]
def should_include(filename: str) -> bool:
name_lower = filename.lower()
for pattern in EXCLUDE_PATTERNS:
if pattern in name_lower:
return False
return True
def main():
if not ARTICLES_DIR.exists():
print(f"错误：目录不存在: {ARTICLES_DIR}")
return
# 确保目标目录存在
PDF_DEST_DIR.mkdir(parents=True, exist_ok=True)
# 初始化数据库
print("初始化数据库...")
init_db()
# 初始化 FAISS 索引
print("初始化 FAISS 索引...")
ks.init_faiss_index()
# 递归扫描所有 PDF
print(f"\n扫描目录: {ARTICLES_DIR}")
pdf_files = sorted(ARTICLES_DIR.rglob("*.pdf"))
pdf_files = [f for f in pdf_files if not f.name.startswith("_")]
print(f"发现 {len(pdf_files)} 个 PDF 文件")
# 过滤
pdf_files = [f for f in pdf_files if should_include(f.name)]
print(f"过滤后剩余 {len(pdf_files)} 个文件待导入\n")
# 按站点分组统计
from collections import Counter
site_counts = Counter()
for f in pdf_files:
parts = f.parts
for p in parts:
if p.startswith("site"):
site_counts[p] += 1
break
else:
site_counts["other"] += 1
print("按站点分布:")
for site, cnt in site_counts.most_common():
print(f" {site}: {cnt} 个")
# 检查已存在
db = SessionLocal()
try:
existing_hashes = set()
all_docs = db.query(KnowledgeDocument).all()
for d in all_docs:
existing_hashes.add(d.file_hash)
imported = 0
skipped = 0
failed = 0
total = len(pdf_files)
for idx, pdf_path in enumerate(pdf_files):
file_hash = ks._compute_sha256(pdf_path)
if file_hash in existing_hashes:
if (idx + 1) % 100 == 0:
print(f" [{idx+1}/{total}] 已导入={imported} 跳过={skipped} 失败={failed}")
skipped += 1
continue
# 复制到知识库 PDF 目录
dest_name = pdf_path.name
dest_path = PDF_DEST_DIR / dest_name
# 处理重名
counter = 1
while dest_path.exists():
stem = pdf_path.stem
dest_name = f"{stem}_{counter}.pdf"
dest_path = PDF_DEST_DIR / dest_name
counter += 1
shutil.copy2(pdf_path, dest_path)
try:
result = ks.upload_document(dest_path, dest_name)
if result.get("duplicate"):
skipped += 1
continue
doc_id = result["id"]
# 同步处理（大文件可能需要较长时间）
ks.process_document(doc_id)
doc = db.query(KnowledgeDocument).filter(
KnowledgeDocument.id == doc_id
).first()
if doc and doc.status == "ready":
imported += 1
else:
status = doc.status if doc else "unknown"
error = doc.error_message if doc else ""
failed += 1
if idx < 5:
print(f" [失败] {pdf_path.name}: status={status} error={error[:80]}")
except Exception as e:
failed += 1
if idx < 5:
print(f" [异常] {pdf_path.name}: {e}")
if (idx + 1) % 50 == 0 or idx < 5:
print(f" [{idx+1}/{total}] 已导入={imported} 跳过={skipped} 失败={failed}")
print(f"\n{'='*50}")
print(f"导入完成！")
print(f" 成功导入: {imported}")
print(f" 跳过（已存在）: {skipped}")
print(f" 失败: {failed}")
print(f"{'='*50}")
finally:
db.close()
if __name__ == "__main__":
main()
|