218 lines
5.9 KiB
Python
218 lines
5.9 KiB
Python
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
|