177 lines
6.2 KiB
Python
177 lines
6.2 KiB
Python
from fastapi import UploadFile
|
||
import os
|
||
from typing import Optional
|
||
import tempfile
|
||
import requests
|
||
import time
|
||
from tqdm import tqdm
|
||
from pathlib import Path
|
||
|
||
class ProgressFileReader:
|
||
def __init__(self, file_path, progress_bar):
|
||
self.file = open(file_path, 'rb')
|
||
self.progress_bar = progress_bar
|
||
|
||
def read(self, size=-1):
|
||
data = self.file.read(size)
|
||
if data:
|
||
self.progress_bar.update(len(data))
|
||
return data
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
self.file.close()
|
||
async def convert_doc_to_pdf(api_url: str, file: UploadFile, timeout: int = 300) -> Optional[str]:
|
||
"""
|
||
将DOC文件转换为pdf
|
||
|
||
参数:
|
||
api_url: 转换服务的URL (如 "http://172.18.127.124:8000")
|
||
file: 要转换的UploadFile对象
|
||
timeout: 请求超时时间(秒)
|
||
"""
|
||
try:
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=".doc") as tmp_doc:
|
||
content = await file.read()
|
||
tmp_doc.write(content)
|
||
tmp_path = tmp_doc.name
|
||
|
||
file_name = file.filename
|
||
file_size = os.path.getsize(tmp_path)
|
||
output_name = Path(file_name).stem + ".pdf"
|
||
|
||
with tqdm(total=file_size, unit='B', unit_scale=True, desc="上传进度") as pbar:
|
||
with ProgressFileReader(tmp_path, pbar) as file_reader:
|
||
files = {'file': (file_name, file_reader)}
|
||
params = {'src': 'doc', 'to': 'pdf'}
|
||
|
||
response = requests.post(
|
||
f"{api_url}/transfer",
|
||
files=files,
|
||
params=params,
|
||
timeout=timeout
|
||
)
|
||
|
||
if response.status_code == 200:
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_docx:
|
||
tmp_docx.write(response.content)
|
||
output_path = tmp_docx.name
|
||
|
||
os.unlink(tmp_path)
|
||
|
||
return output_path
|
||
else:
|
||
raise Exception(f"转换失败: {response.text}")
|
||
|
||
except requests.exceptions.Timeout:
|
||
raise Exception(f"请求超时(超过{timeout}秒)")
|
||
except requests.exceptions.RequestException as e:
|
||
raise Exception(f"请求错误: {str(e)}")
|
||
except Exception as e:
|
||
# 确保临时文件被清理
|
||
if 'tmp_path' in locals() and os.path.exists(tmp_path):
|
||
os.unlink(tmp_path)
|
||
if 'output_path' in locals() and os.path.exists(output_path):
|
||
os.unlink(output_path)
|
||
raise Exception(f"转换过程中出错: {str(e)}")
|
||
|
||
async def convert_docx_to_pdf(api_url: str, file: UploadFile, timeout: int = 300) -> Optional[str]:
|
||
"""
|
||
将DOCX文件转换为PDF（解决中文缺失问题）
|
||
|
||
参数:
|
||
api_url: 转换服务的URL (如 "http://172.18.127.124:8000")
|
||
file: 要转换的UploadFile对象
|
||
timeout: 请求超时时间(秒)
|
||
"""
|
||
try:
|
||
# 创建临时文件保存上传的DOCX
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=".docx") as tmp_docx:
|
||
content = await file.read()
|
||
tmp_docx.write(content)
|
||
tmp_path = tmp_docx.name
|
||
|
||
# 获取文件信息
|
||
file_name = file.filename
|
||
file_size = os.path.getsize(tmp_path)
|
||
output_name = Path(file_name).stem + ".pdf"
|
||
|
||
print(f"正在转换: {file_name} ({file_size/1024:.1f} KB) -> PDF")
|
||
start_time = time.time()
|
||
|
||
# 使用带进度条的文件读取器
|
||
with tqdm(total=file_size, unit='B', unit_scale=True, desc="上传进度") as pbar:
|
||
with ProgressFileReader(tmp_path, pbar) as file_reader:
|
||
files = {'file': (file_name, file_reader)}
|
||
params = {
|
||
'src': 'docx',
|
||
'to': 'pdf',
|
||
'font': 'SimSun',
|
||
'encoding': 'utf-8'
|
||
}
|
||
|
||
response = requests.post(
|
||
f"{api_url}/transfer",
|
||
files=files,
|
||
params=params,
|
||
timeout=timeout
|
||
)
|
||
|
||
# 处理响应
|
||
if response.status_code == 200:
|
||
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_pdf:
|
||
tmp_pdf.write(response.content)
|
||
output_path = tmp_pdf.name
|
||
|
||
|
||
if os.path.getsize(output_path) < 1024: # 如果PDF文件过小
|
||
raise Exception("生成的PDF可能不包含有效内容")
|
||
|
||
elapsed = time.time() - start_time
|
||
print(f"\n转换成功! 保存为: {output_path}")
|
||
print(f"总耗时: {elapsed:.2f}秒")
|
||
|
||
os.unlink(tmp_path)
|
||
|
||
return output_path
|
||
else:
|
||
raise Exception(f"转换失败: {response.text}")
|
||
|
||
except requests.exceptions.Timeout:
|
||
raise Exception(f"请求超时(超过{timeout}秒)")
|
||
except requests.exceptions.RequestException as e:
|
||
raise Exception(f"请求错误: {str(e)}")
|
||
except Exception as e:
|
||
# 确保临时文件被清理
|
||
if 'tmp_path' in locals() and os.path.exists(tmp_path):
|
||
os.unlink(tmp_path)
|
||
if 'output_path' in locals() and os.path.exists(output_path):
|
||
os.unlink(output_path)
|
||
raise Exception(f"转换过程中出错: {str(e)}")
|
||
|
||
|
||
|
||
#调用示例：
|
||
async def split_result(file: UploadFile = File(...),
|
||
chunk_type: str = Form("01"),
|
||
chunk_size: int = Form(256),
|
||
chunk_overlap: int = Form(50),
|
||
max_level: Optional[int]=Form(3)):
|
||
"""
|
||
拆分结果
|
||
"""
|
||
filename = file.filename
|
||
file_ext = os.path.splitext(filename)[1]
|
||
|
||
if file_ext == ".docx":
|
||
# 处理DOCX文件
|
||
converted_path = await convert_docx_to_pdf(os.getenv("CONVERT_DOC_URL","http://172.18.127.124:8000"), file)
|
||
|
||
# 读取转换后的DOCX文件
|
||
with open(converted_path, 'rb') as f:
|
||
converted_file = UploadFile(filename=f"{PATH(filename).stem}.pdf", file=f,headers=headers)
|
||
#调用Mineru
|
||
md_content = await process_pdf_logic(converted_file) |