Files
HuiBrain/app/services/import_service.py
EduBrain Dev b17786b57b feat: v1.2.0 - 图片去重与管理、微信机器人优化、搜索设置可配置
主要功能:
- 图片上传时 OCR 内容去重(3个上传端点统一使用公共函数 _check_ocr_duplicate)
- 图片管理 Tab:展示所有图片、手动删除、一键去重
- 搜索结果详情弹窗增加删除按钮(带确认弹窗)
- 图片管理卡片点击查看详情(复用 showOcrDetailModal)
- 搜索限制和 LLM 批量判断数量可通过网站设置
- MiniMax API 调用添加 reasoning_split=True
- 企业微信机器人:WebSocket 长连接、图片搜索、配置化搜索数量
- 版本号升级至 1.2.0
2026-04-13 22:25:08 +08:00

453 lines
14 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
导入服务
负责从 Markdown / 纯文本 / Word 文件导入知识页面,并进行文本分块和向量化
"""
from __future__ import annotations
import logging
import re
from pathlib import Path
from typing import List, Optional
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession
from app.config import settings
from app.database import IS_SQLITE
logger = logging.getLogger(__name__)
def _parse_frontmatter(content: str) -> tuple[dict, str]:
"""
手动解析 YAML frontmatter。
格式:--- 开头和结尾包围的 YAML 块。
"""
import yaml
if not content.startswith("---"):
return {}, content
parts = content.split("---", 2)
if len(parts) < 3:
return {}, content
try:
metadata = yaml.safe_load(parts[1]) or {}
except Exception:
metadata = {}
body = parts[2].strip()
return metadata, body
class ImportService:
"""知识导入服务"""
def __init__(self, db: AsyncSession):
self.db = db
async def import_file(
self,
file_path: str,
course_name: Optional[str] = None,
teacher_name: Optional[str] = None,
live_date: Optional[str] = None,
) -> dict:
"""
从文件导入知识内容。
支持的格式:
- Markdown (.md): 解析 frontmatter 元数据,按标题分页
- 纯文本 (.txt): 整体作为一个页面
- Word 文档 (.docx): 提取文本,按标题分页
Args:
file_path: 文件路径
course_name: 课程名称(覆盖 frontmatter
teacher_name: 讲师名称(覆盖 frontmatter
live_date: 直播日期(覆盖 frontmatter
Returns:
导入结果统计
"""
path = Path(file_path)
if not path.exists():
raise FileNotFoundError(f"文件不存在: {file_path}")
suffix = path.suffix.lower()
if suffix == ".md":
content = path.read_text(encoding="utf-8")
return await self._import_markdown(
content, path.name, course_name, teacher_name, live_date
)
elif suffix == ".txt":
content = path.read_text(encoding="utf-8")
return await self._import_text(
content, path.name, course_name, teacher_name, live_date
)
elif suffix == ".docx":
return await self._import_docx(
str(path), path.name, course_name, teacher_name, live_date
)
else:
raise ValueError(f"不支持的文件格式: {suffix},仅支持 .md、.txt 和 .docx")
async def _import_markdown(
self,
content: str,
filename: str,
course_name: Optional[str] = None,
teacher_name: Optional[str] = None,
live_date: Optional[str] = None,
) -> dict:
"""导入 Markdown 文件,解析 frontmatter 并按标题分页"""
# 解析 frontmatter
fm_data, body = _parse_frontmatter(content)
# frontmatter 中的元数据可作为默认值
fm_course = fm_data.get("course", fm_data.get("course_name"))
fm_teacher = fm_data.get("teacher", fm_data.get("teacher_name"))
fm_date = fm_data.get("date", fm_data.get("live_date"))
final_course = course_name or fm_course
final_teacher = teacher_name or fm_teacher
final_date = live_date or (str(fm_date) if fm_date else None)
# 按二级标题(##)拆分页面
sections = self._split_markdown_sections(body)
if not sections:
# 没有二级标题,整体作为一个页面
sections = [{"title": Path(filename).stem, "content": body}]
imported_pages = 0
imported_chunks = 0
for idx, section in enumerate(sections):
page_number = idx + 1
# 插入知识页面
page_id = await self._insert_page(
title=section["title"],
content=section["content"],
source_file=filename,
course_name=final_course,
teacher_name=final_teacher,
live_date=final_date,
page_number=page_number,
)
imported_pages += 1
# 分块并向量化
chunks = self._chunk_text(section["content"])
chunk_count = await self._insert_chunks(page_id, chunks)
imported_chunks += chunk_count
logger.info(
"Markdown 导入完成: file=%s, pages=%d, chunks=%d",
filename, imported_pages, imported_chunks,
)
return {
"file": filename,
"pages": imported_pages,
"chunks": imported_chunks,
}
async def _import_text(
self,
content: str,
filename: str,
course_name: Optional[str] = None,
teacher_name: Optional[str] = None,
live_date: Optional[str] = None,
) -> dict:
"""导入纯文本文件,整体作为一个页面"""
title = Path(filename).stem
page_id = await self._insert_page(
title=title,
content=content,
source_file=filename,
course_name=course_name,
teacher_name=teacher_name,
live_date=live_date,
page_number=1,
)
chunks = self._chunk_text(content)
chunk_count = await self._insert_chunks(page_id, chunks)
logger.info(
"文本导入完成: file=%s, chunks=%d",
filename, chunk_count,
)
return {
"file": filename,
"pages": 1,
"chunks": chunk_count,
}
async def _import_docx(
self,
file_path: str,
filename: str,
course_name: Optional[str] = None,
teacher_name: Optional[str] = None,
live_date: Optional[str] = None,
) -> dict:
"""
导入 Word 文档(.docx
提取所有段落文本按标题Heading 1/2拆分为多个知识页面。
"""
from docx import Document
doc = Document(file_path)
# 提取所有段落,保留标题层级信息
sections: List[dict] = [] # [{"title": "...", "content": "...", "level": int}]
current_title = Path(filename).stem
current_content: List[str] = []
current_level = 0
for para in doc.paragraphs:
style_name = (para.style.name or "").lower()
# 判断是否为标题段落
if style_name.startswith("heading"):
try:
level = int(style_name.replace("heading", "").strip())
except ValueError:
level = 1
# 保存上一个段落
if current_content:
text = "\n\n".join(current_content).strip()
if text:
sections.append({
"title": current_title,
"content": text,
"level": current_level,
})
# 开始新段落
current_title = para.text.strip() or f"{len(sections) + 1}"
current_content = []
current_level = level
else:
text = para.text.strip()
if text:
current_content.append(text)
# 保存最后一个段落
if current_content:
text = "\n\n".join(current_content).strip()
if text:
sections.append({
"title": current_title,
"content": text,
"level": current_level,
})
if not sections:
raise ValueError(f"Word 文档内容为空: {filename}")
# 如果只有一个段落且没有标题,整体作为一个页面
if len(sections) == 1 and sections[0]["level"] == 0:
sections[0]["title"] = Path(filename).stem
imported_pages = 0
imported_chunks = 0
for idx, section in enumerate(sections):
page_id = await self._insert_page(
title=section["title"],
content=section["content"],
source_file=filename,
course_name=course_name,
teacher_name=teacher_name,
live_date=live_date,
page_number=idx + 1,
)
imported_pages += 1
chunks = self._chunk_text(section["content"])
chunk_count = await self._insert_chunks(page_id, chunks)
imported_chunks += chunk_count
logger.info(
"Word 文档导入完成: file=%s, pages=%d, chunks=%d",
filename, imported_pages, imported_chunks,
)
return {
"file": filename,
"pages": imported_pages,
"chunks": imported_chunks,
}
async def _insert_page(
self,
title: str,
content: str,
source_file: str,
course_name: Optional[str] = None,
teacher_name: Optional[str] = None,
live_date: Optional[str] = None,
page_number: Optional[int] = None,
) -> int:
"""插入知识页面记录,返回页面 ID"""
sql = text("""
INSERT INTO knowledge_pages (title, content, source_file, course_name, teacher_name, live_date, page_number)
VALUES (:title, :content, :source_file, :course_name, :teacher_name, :live_date, :page_number)
RETURNING id
""")
result = await self.db.execute(sql, {
"title": title,
"content": content,
"source_file": source_file,
"course_name": course_name,
"teacher_name": teacher_name,
"live_date": live_date,
"page_number": page_number,
})
row = result.fetchone()
await self.db.flush()
return row.id
async def _insert_chunks(self, page_id: int, chunks: List[str]) -> int:
"""
批量插入分块记录并生成嵌入向量。
Args:
page_id: 关联的知识页面 ID
chunks: 分块文本列表
Returns:
插入的分块数量
"""
if not chunks:
return 0
# 批量生成嵌入向量
try:
from app.services.embedding_service import EmbeddingService
embeddings = await EmbeddingService.embed_batch(chunks)
except Exception as exc:
logger.error("嵌入向量生成失败,分块将不包含向量: %s", exc)
embeddings = [None] * len(chunks)
# 批量插入
for idx, (chunk_text, embedding) in enumerate(zip(chunks, embeddings)):
if IS_SQLITE:
# SQLite: 向量存为 JSON 文本
import json as _json
embedding_str = _json.dumps(embedding) if embedding else None
sql = text("""
INSERT INTO knowledge_chunks (page_id, chunk_index, content, embedding)
VALUES (:page_id, :chunk_index, :content, :embedding)
""")
else:
# PostgreSQL: 使用 pgvector 类型
embedding_str = "[" + ",".join(str(x) for x in embedding) + "]" if embedding else None
sql = text("""
INSERT INTO knowledge_chunks (page_id, chunk_index, content, embedding)
VALUES (:page_id, :chunk_index, :content, :embedding::vector)
""")
await self.db.execute(sql, {
"page_id": page_id,
"chunk_index": idx,
"content": chunk_text,
"embedding": embedding_str,
})
# 更新全文搜索向量(仅 PostgreSQL
if not IS_SQLITE:
await self.db.execute(text("""
UPDATE knowledge_chunks
SET search_vector = to_tsvector('chinese_zh', content)
WHERE page_id = :page_id
"""), {"page_id": page_id})
await self.db.flush()
return len(chunks)
@staticmethod
def _split_markdown_sections(body: str) -> List[dict]:
"""
按二级标题(##)拆分 Markdown 内容为多个段落。
Returns:
[{"title": "...", "content": "..."}, ...]
"""
sections: List[dict] = []
# 匹配 ## 开头的标题
pattern = re.compile(r"^##\s+(.+)$", re.MULTILINE)
matches = list(pattern.finditer(body))
if not matches:
return []
for i, match in enumerate(matches):
title = match.group(1).strip()
start = match.end()
end = matches[i + 1].start() if i + 1 < len(matches) else len(body)
content = body[start:end].strip()
if content:
sections.append({"title": title, "content": content})
return sections
@staticmethod
def _chunk_text(text: str) -> List[str]:
"""
将文本按固定大小分块,保留重叠部分。
使用简单的字符数分块策略,按段落边界切分以保持语义完整性。
"""
chunk_size = settings.CHUNK_SIZE
overlap = settings.CHUNK_OVERLAP
if not text or len(text) <= chunk_size:
return [text] if text.strip() else []
# 按段落分割
paragraphs = re.split(r"\n{2,}", text)
chunks: List[str] = []
current_chunk = ""
for para in paragraphs:
para = para.strip()
if not para:
continue
if len(current_chunk) + len(para) + 2 <= chunk_size:
# 当前块还能容纳这个段落
if current_chunk:
current_chunk += "\n\n" + para
else:
current_chunk = para
else:
# 当前块已满,保存并开始新块
if current_chunk:
chunks.append(current_chunk)
if len(para) > chunk_size:
# 单个段落超过 chunk_size强制切分
for i in range(0, len(para), chunk_size - overlap):
piece = para[i : i + chunk_size]
if piece.strip():
chunks.append(piece)
current_chunk = ""
else:
current_chunk = para
if current_chunk.strip():
chunks.append(current_chunk)
return chunks