ai_pipeline/src/pipeline/core/utils.py

190 lines
5.1 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import asyncio
import docx
import fitz # PyMuPDF
import aiofiles
import io
import os
import sys
from pathlib import Path
from PIL import Image
from loguru import logger
from src.pipeline.config import config
# -----------------------------
# 日志
# -----------------------------
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