""" 文档服务 - 上传、解析、分块、存储 支持多种文档格式 + 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}"