easy-rag/utils/file_utils.py

56 lines
1.4 KiB
Python

import PyPDF2
from typing import BinaryIO, List
import os
import asyncio
async def extract_text_from_pdf_async(file: BinaryIO) -> str:
"""从PDF文件中提取文本"""
def _parse_pdf():
try:
pdf_reader = PyPDF2.PdfReader(file)
text = ""
for page in pdf_reader.pages:
text += page.extract_text() + "\n"
return text.strip()
except Exception as e:
raise ValueError(f"PDF解析失败: {str(e)}")
return await asyncio.to_thread(_parse_pdf)
async def delete_file_async(filepath: str) -> None:
"""删除文件"""
def _delete():
if os.path.exists(filepath):
os.remove(filepath)
return await asyncio.to_thread(_delete)
def validate_file_size(file_size: int, max_size: int = 10 * 1024 * 1024) -> bool:
"""验证文件大小"""
return file_size <= max_size
def ensure_directory_exists(directory: str) -> None:
"""确保目录存在"""
if not os.path.exists(directory):
os.makedirs(directory, exist_ok=True)
def get_file_extension(filename: str) -> str:
"""获取文件扩展名"""
return os.path.splitext(filename)[1].lower()
def is_supported_file_type(
filename: str, supported_types: List[str] = [".pdf", ".txt"]
) -> bool:
"""检查是否为支持的文件类型"""
return get_file_extension(filename) in supported_types