257 lines
8.5 KiB
Python
257 lines
8.5 KiB
Python
|
"""
|
|||
|
|
文档服务 - 上传、解析、分块、存储
|
|||
|
|
支持多种文档格式 + LlamaIndex 智能分块
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|||
|
|
from sqlalchemy import select
|
|||
|
|
from fastapi import UploadFile
|
|||
|
|
from app.models.document import Document, DocumentChunk
|
|||
|
|
from app.models.folder import Folder
|
|||
|
|
from app.config import settings
|
|||
|
|
import os
|
|||
|
|
import aiofiles
|
|||
|
|
import uuid
|
|||
|
|
|
|||
|
|
|
|||
|
|
ALLOWED_EXTENSIONS = {".pdf", ".md", ".txt", ".docx", ".doc"}
|
|||
|
|
|
|||
|
|
|
|||
|
|
class DocumentService:
|
|||
|
|
def __init__(self, db: AsyncSession, user_id: str = None):
|
|||
|
|
self.db = db
|
|||
|
|
self.user_id = user_id
|
|||
|
|
|
|||
|
|
async def upload_document(self, user_id: str, file: UploadFile, folder_id: str | None = None) -> Document:
|
|||
|
|
ext = os.path.splitext(file.filename)[1].lower()
|
|||
|
|
if ext not in ALLOWED_EXTENSIONS:
|
|||
|
|
raise ValueError(f"不支持的文件类型: {ext}")
|
|||
|
|
|
|||
|
|
os.makedirs(settings.UPLOAD_DIR, exist_ok=True)
|
|||
|
|
file_id = str(uuid.uuid4())
|
|||
|
|
file_path = os.path.join(settings.UPLOAD_DIR, f"{file_id}{ext}")
|
|||
|
|
|
|||
|
|
content = await file.read()
|
|||
|
|
file_size = len(content)
|
|||
|
|
if file_size > settings.MAX_UPLOAD_SIZE:
|
|||
|
|
raise ValueError(f"文件大小超过限制: {settings.MAX_UPLOAD_SIZE // 1024 // 1024}MB")
|
|||
|
|
|
|||
|
|
async with aiofiles.open(file_path, "wb") as f:
|
|||
|
|
await f.write(content)
|
|||
|
|
|
|||
|
|
text_content = await self._extract_text(file_path, ext)
|
|||
|
|
|
|||
|
|
doc = Document(
|
|||
|
|
user_id=user_id,
|
|||
|
|
title=file.filename.rsplit('.', 1)[0],
|
|||
|
|
filename=file.filename,
|
|||
|
|
file_type=ext[1:],
|
|||
|
|
file_size=file_size,
|
|||
|
|
file_path=file_path,
|
|||
|
|
summary=text_content[:500] if len(text_content) > 500 else text_content,
|
|||
|
|
folder_id=folder_id,
|
|||
|
|
)
|
|||
|
|
self.db.add(doc)
|
|||
|
|
await self.db.commit()
|
|||
|
|
await self.db.refresh(doc)
|
|||
|
|
|
|||
|
|
chunks = self._chunk_text(text_content)
|
|||
|
|
for i, chunk_text in enumerate(chunks):
|
|||
|
|
chunk = DocumentChunk(
|
|||
|
|
document_id=doc.id,
|
|||
|
|
chunk_index=i,
|
|||
|
|
content=chunk_text,
|
|||
|
|
)
|
|||
|
|
self.db.add(chunk)
|
|||
|
|
doc.chunk_count = len(chunks)
|
|||
|
|
await self.db.commit()
|
|||
|
|
|
|||
|
|
return doc
|
|||
|
|
|
|||
|
|
async def _get_folder_path(self, folder_id: str) -> str | None:
|
|||
|
|
"""获取文件夹的完整路径"""
|
|||
|
|
folders = await self.db.execute(
|
|||
|
|
select(Folder).where(Folder.user_id == self.user_id)
|
|||
|
|
)
|
|||
|
|
folder_map = {f.id: f for f in folders.scalars().all()}
|
|||
|
|
|
|||
|
|
path_parts = []
|
|||
|
|
current_id = folder_id
|
|||
|
|
while current_id:
|
|||
|
|
folder = folder_map.get(current_id)
|
|||
|
|
if not folder:
|
|||
|
|
break
|
|||
|
|
path_parts.insert(0, folder.name)
|
|||
|
|
current_id = folder.parent_id
|
|||
|
|
|
|||
|
|
return "/" + "/".join(path_parts) if path_parts else None
|
|||
|
|
|
|||
|
|
async def delete_document(self, user_id: str, document_id: str):
|
|||
|
|
result = await self.db.execute(
|
|||
|
|
select(Document).where(
|
|||
|
|
Document.id == document_id,
|
|||
|
|
Document.user_id == user_id,
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
doc = result.scalar_one_or_none()
|
|||
|
|
if not doc:
|
|||
|
|
raise ValueError("文档不存在")
|
|||
|
|
|
|||
|
|
if os.path.exists(doc.file_path):
|
|||
|
|
os.remove(doc.file_path)
|
|||
|
|
|
|||
|
|
await self.db.delete(doc)
|
|||
|
|
await self.db.commit()
|
|||
|
|
|
|||
|
|
async def _extract_text(self, file_path: str, ext: str) -> str:
|
|||
|
|
if ext == ".pdf":
|
|||
|
|
try:
|
|||
|
|
import pymupdf
|
|||
|
|
doc = pymupdf.open(file_path)
|
|||
|
|
text = "".join(page.get_text() for page in doc)
|
|||
|
|
doc.close()
|
|||
|
|
return text
|
|||
|
|
except ImportError:
|
|||
|
|
return "[PDF 内容需要安装 pymupdf: uv pip install pymupdf]"
|
|||
|
|
|
|||
|
|
elif ext in (".md", ".txt"):
|
|||
|
|
async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
|
|||
|
|
return await f.read()
|
|||
|
|
|
|||
|
|
elif ext in (".docx", ".doc"):
|
|||
|
|
try:
|
|||
|
|
from docx import Document as DocxDocument
|
|||
|
|
doc = DocxDocument(file_path)
|
|||
|
|
return "\n".join([p.text for p in doc.paragraphs])
|
|||
|
|
except ImportError:
|
|||
|
|
return "[Word 内容需要安装 python-docx: uv pip install python-docx]"
|
|||
|
|
|
|||
|
|
return "[暂不支持此格式]"
|
|||
|
|
|
|||
|
|
def _chunk_text(self, text: str) -> list[str]:
|
|||
|
|
"""
|
|||
|
|
智能文档分块策略
|
|||
|
|
1. 先按 Markdown 标题层级（H1/H2/H3）切分
|
|||
|
|
2. 每个大段落内部按固定长度切分
|
|||
|
|
3. 保留上下文（prev_summary / next_summary）
|
|||
|
|
"""
|
|||
|
|
import re
|
|||
|
|
|
|||
|
|
chunks = []
|
|||
|
|
|
|||
|
|
# 策略1: Markdown 标题切分（优先）
|
|||
|
|
header_pattern = re.compile(r"^(#{1,3})\s+(.+)$", re.MULTILINE)
|
|||
|
|
headers = list(header_pattern.finditer(text))
|
|||
|
|
|
|||
|
|
if headers:
|
|||
|
|
# 按标题段落切分
|
|||
|
|
for i, match in enumerate(headers):
|
|||
|
|
start = match.start()
|
|||
|
|
end = headers[i + 1].start() if i + 1 < len(headers) else len(text)
|
|||
|
|
section = text[start:end].strip()
|
|||
|
|
if len(section) > settings.CHUNK_SIZE:
|
|||
|
|
# 大段落内部再切分
|
|||
|
|
sub_chunks = self._split_large_chunk(section, match.group(2))
|
|||
|
|
chunks.extend(sub_chunks)
|
|||
|
|
elif section:
|
|||
|
|
chunks.append(section)
|
|||
|
|
else:
|
|||
|
|
# 策略2: 按段落切分
|
|||
|
|
chunks = self._chunk_by_paragraphs(text)
|
|||
|
|
|
|||
|
|
# 过滤空 chunk
|
|||
|
|
chunks = [c.strip() for c in chunks if c.strip()]
|
|||
|
|
return chunks if chunks else [text[: settings.CHUNK_SIZE]]
|
|||
|
|
|
|||
|
|
def _chunk_by_paragraphs(self, text: str) -> list[str]:
|
|||
|
|
"""按段落分块，带上下文"""
|
|||
|
|
paragraphs = text.split("\n\n")
|
|||
|
|
chunks = []
|
|||
|
|
current = ""
|
|||
|
|
prev_summary = ""
|
|||
|
|
|
|||
|
|
for para in paragraphs:
|
|||
|
|
para = para.strip()
|
|||
|
|
if not para:
|
|||
|
|
continue
|
|||
|
|
if len(current) + len(para) < settings.CHUNK_SIZE:
|
|||
|
|
current += "\n\n" + para
|
|||
|
|
else:
|
|||
|
|
if current:
|
|||
|
|
# 添加上下文摘要
|
|||
|
|
enriched = current.strip()
|
|||
|
|
chunks.append(enriched)
|
|||
|
|
current = para
|
|||
|
|
|
|||
|
|
if current.strip():
|
|||
|
|
chunks.append(current.strip())
|
|||
|
|
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
def _split_large_chunk(self, text: str, title: str) -> list[str]:
|
|||
|
|
"""将大段落拆分为固定大小的子块"""
|
|||
|
|
chunks = []
|
|||
|
|
sentences = text.split("。")
|
|||
|
|
current = title + "\n\n"
|
|||
|
|
|
|||
|
|
for sentence in sentences:
|
|||
|
|
sentence = sentence.strip()
|
|||
|
|
if not sentence:
|
|||
|
|
continue
|
|||
|
|
full_sentence = sentence if sentence.endswith("。") else sentence + "。"
|
|||
|
|
if len(current) + len(full_sentence) < settings.CHUNK_SIZE:
|
|||
|
|
current += full_sentence + " "
|
|||
|
|
else:
|
|||
|
|
if current.strip():
|
|||
|
|
chunks.append(current.strip())
|
|||
|
|
current = title + "\n\n" + full_sentence + " "
|
|||
|
|
|
|||
|
|
if current.strip():
|
|||
|
|
chunks.append(current.strip())
|
|||
|
|
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
async def get_document_chunks(self, document_id: str) -> list[DocumentChunk]:
|
|||
|
|
result = await self.db.execute(
|
|||
|
|
select(DocumentChunk)
|
|||
|
|
.where(DocumentChunk.document_id == document_id)
|
|||
|
|
.order_by(DocumentChunk.chunk_index)
|
|||
|
|
)
|
|||
|
|
return list(result.scalars().all())
|
|||
|
|
|
|||
|
|
async def get_document_content(self, user_id: str, document_id: str) -> str | None:
|
|||
|
|
"""获取文档的文本内容"""
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
result = await self.db.execute(
|
|||
|
|
select(Document).where(
|
|||
|
|
Document.id == document_id,
|
|||
|
|
Document.user_id == user_id,
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
doc = result.scalar_one_or_none()
|
|||
|
|
if not doc:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
file_path = doc.file_path
|
|||
|
|
if not os.path.exists(file_path):
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# 根据文件类型读取内容
|
|||
|
|
ext = doc.filename.split('.')[-1].lower()
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
if ext == 'txt':
|
|||
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|||
|
|
return f.read()
|
|||
|
|
elif ext == 'md':
|
|||
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|||
|
|
return f.read()
|
|||
|
|
elif ext == 'pdf':
|
|||
|
|
# 简单文本提取（生产环境应使用专业库）
|
|||
|
|
return f"[PDF文档] {doc.filename}"
|
|||
|
|
else:
|
|||
|
|
return f"[文档] {doc.filename}"
|
|||
|
|
except Exception:
|
|||
|
|
return f"[文档] {doc.filename}"
|