496e11e26e
- config.py: DATABASE_URL 默认值改为 SQLite - database.py: 移除 PostgreSQL 分支，简化为纯 SQLite - models/base.py: 移除 pgvector 导入和条件分支 - search_service.py: 移除 _search_postgres 方法 - import_service.py: 移除 pgvector 相关代码 - requirements.txt: 移除 asyncpg/alembic/pgvector 依赖 - pyproject.toml: 同步移除相关依赖 - docker-compose.yml: 移除 db 服务，- 删除 alembic.ini/Dockerfile.db/sql 目录 - README.md: 更新文档，移除 PostgreSQL 相关内容 适合 NAS 等资源受限环境的轻量级部署
436 lines
14 KiB
Python
436 lines
14 KiB
Python
"""
|
||
导入服务
|
||
负责从 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
|
||
|
||
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)):
|
||
# 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)
|
||
""")
|
||
await self.db.execute(sql, {
|
||
"page_id": page_id,
|
||
"chunk_index": idx,
|
||
"content": chunk_text,
|
||
"embedding": embedding_str,
|
||
})
|
||
|
||
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
|