import asyncio import docx import fitz # PyMuPDF import aiofiles import io import re import sys from pathlib import Path from PIL import Image from loguru import logger from src.pipeline.config import config from baidusearch.baidusearch import search # ----------------------------- # 日志 # ----------------------------- logger.remove() # 清除已有 handler(包括 pytest 导致的重复加载 logger.add(sys.stdout, level=config["logger_level"], colorize=True) logger.add( "logs/pipeline.log", level=config["logger_level"], rotation="10 MB", # 自动分割 retention="7 days", # 保留时间 compression="zip", # 自动压缩 enqueue=True, # 多线程安全 colorize=False, # 文件不需要颜色 backtrace=True, diagnose=True, ) # ----------------------------- # 文本加载 # ----------------------------- # 初始化 OCR(只初始化一次),注意这是 CPU 版本,如果需要 GPU 需要额外配置 ocr = {} async def _ocr_image_bytes(img_bytes: bytes) -> str: """对图片字节流做 OCR(线程池避免阻塞 asyncio)""" def _ocr(): img = Image.open(io.BytesIO(img_bytes)) result = ocr.ocr(img, cls=True) # result: [[ ['text', 'score'], ... ]] if not result or not result[0]: return "" lines = [] for line in result: txt = line[1][0] # OCR 文字 confidence = line[1][1] lines.append(txt) return "\n".join(lines) return await asyncio.to_thread(_ocr) async def _load_txt(path: str) -> str: async with aiofiles.open(path, "r", encoding="utf-8") as f: return await f.read() async def _load_md(path: str) -> str: async with aiofiles.open(path, "r", encoding="utf-8") as f: return await f.read() async def _load_docx(path: str) -> str: # docx 读取是阻塞 I/O → 放入线程池 def _read(): doc = docx.Document(path) return "\n".join(p.text for p in doc.paragraphs) return await asyncio.to_thread(_read) async def _load_pdf(path: str) -> str: def _read_pdf(): doc = fitz.open(path) pages = [] for page in doc: page_text = page.get_text() # 获取图片 image_bytes_list = [] # TODO: ocr 识别 # for img in page.get_images(full=True): # xref = img[0] # pix = fitz.Pixmap(doc, xref) # img_bytes = pix.tobytes("png") # image_bytes_list.append(img_bytes) pages.append((page_text, image_bytes_list)) return pages # PDF 解析在子线程执行 pages = await asyncio.to_thread(_read_pdf) final_text = [] # 分页处理 + OCR 并发 for text, image_bytes_list in pages: final_text.append(text) # 并发 OCR ocr_tasks = [asyncio.create_task(_ocr_image_bytes(b)) for b in image_bytes_list] if ocr_tasks: ocr_results = await asyncio.gather(*ocr_tasks) final_text.extend(ocr_results) return "\n".join(final_text) async def load_document(path: str) -> str: suffix = Path(path).suffix.lower() if suffix == ".txt": return await _load_txt(path) if suffix in (".md", ".markdown"): return await _load_md(path) if suffix == ".docx": return await _load_docx(path) if suffix == ".pdf": return await _load_pdf(path) raise ValueError(f"Unsupported file type: {suffix}") # ----------------------------- # 文本切割,目前只有简单长度切割 # ----------------------------- def fixed_size_chunk(text, chunk_size=500, overlap=50): chunks = [] start = 0 while start < len(text): end = start + chunk_size chunks.append(text[max(0, start - overlap) : min(len(text), end + overlap)]) start += chunk_size return chunks def rag_system_prompt() -> str: return """ 你是一个基于知识库的 AI 助手,严格基于提供的文档回答用户问题。 - 不允许编造事实。 - 只使用提供的内容。 - 输出清晰、简洁、准确。 - 如果文档中没有明确答案,请回复 "文档未提供相关信息"。 """ def rag_user_prompt(query: str, documents: list[dict]) -> str: """ 生成基于 RAG 的用户 Prompt。 :param query: 用户问题 :param documents: 检索到的文档列表,每个文档 dict 至少包含 title, content :return: Prompt 字符串 """ # 按相关度排序(如果文档里有 score,可以用) documents = sorted(documents, key=lambda x: x.get("score", 0), reverse=True) # 拼接文档内容 context_lines = [] for i, doc in enumerate(documents, start=1): content = doc.get("content", "").replace("\n", " ").strip() title = doc.get("title", f"文档{i}") context_lines.append(f"文档 {i}:\n标题: {title}\n内容: {content}\n") context_text = "\n".join(context_lines) # 构建最终 Prompt prompt = f""" 以下是从知识库中检索到的内容(按相关度排序): <<< {context_text} >>> 用户问题: {query} 回答要求: 一到两句话回复 """ logger.debug(prompt) return prompt # ----------------------------- # 其他工具 # ----------------------------- async def baidu_search_async(query: str, max_results: int = 5): """ 异步调用 baidusearch(内部用 asyncio.to_thread 封装同步函数) 返回结构化搜索结果列表 """ def sync_search(): return list(search(query, num_results=max_results)) results = await asyncio.to_thread(sync_search) docs = [] for r in results: docs.append( { "title": r.get("title"), "content": re.sub(r"\s+", " ", r.get("abstract", "")), "url": r.get("url"), "type": "web", "score": 0.5, } ) return docs