298 lines
11 KiB
Python
298 lines
11 KiB
Python
from fastapi import UploadFile
|
||
import httpx
|
||
import os
|
||
from typing import Optional
|
||
import tempfile
|
||
import requests
|
||
import time
|
||
from tqdm import tqdm
|
||
from pathlib import Path
|
||
import json
|
||
from io import BytesIO
|
||
from urllib.parse import urljoin
|
||
|
||
|
||
async def process_pdf_logic(pdf_file):
|
||
"""
|
||
接收 PDF 文件并转发到解析服务
|
||
"""
|
||
# 验证文件类型
|
||
if not pdf_file.content_type.startswith("application/pdf"):
|
||
raise Exception(400, "仅支持 PDF 文件")
|
||
|
||
try:
|
||
# 读取文件内容
|
||
# file_content = await pdf_file.read()
|
||
# 配置目标接口参数
|
||
target_url = os.getenv("Mineru_URL", "http://172.18.30.122:9988")
|
||
target_url = target_url + "/analyze-pdf"
|
||
# target_url = os.getenv("Mineru_URL","http://172.18.127.124:9088/pdf_parse")
|
||
# target_url = "http://172.18.30.122:8998/pdf_parse"
|
||
# params = {
|
||
# "parse_method": "auto",
|
||
# "is_json_md_dump": "true",
|
||
# "output_dir": "output"
|
||
# }
|
||
|
||
# 构建文件上传格式
|
||
files = {
|
||
"file": (pdf_file.filename, pdf_file.file, "application/pdf")
|
||
# "pdf_file": (pdf_file.filename, pdf_file.file, "application/pdf")
|
||
}
|
||
|
||
# 发送请求到解析服务
|
||
async with httpx.AsyncClient(timeout=1200.0) as client:
|
||
response = await client.post(
|
||
target_url,
|
||
# params=params,
|
||
files=files,
|
||
headers={"Accept": "application/json"}
|
||
)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
|
||
content_chunk_list = []
|
||
full_content = result.get("data", {}).get("full_content", "")
|
||
content_list = result.get("data", {}).get("content_list", [])
|
||
content_chunk_list = result.get("data", {}).get("content_middle", [])
|
||
print(content_chunk_list)
|
||
# combined_json = ''.join(content_list)
|
||
# content_data = json.loads(combined_json)
|
||
# for idx, item in enumerate(content_data):
|
||
# if isinstance(item, dict): # 确保是字典类型
|
||
# text_type = item.get("type", "")
|
||
# text_level = item.get("text_level", "")
|
||
# text = item.get("text", "")
|
||
# page_idx = item.get("page_idx", -1) # 默认值 -1 表示缺失
|
||
# table_caption = item.get("table_caption", "")
|
||
# table_body = item.get("table_body", "")
|
||
# image_path = item.get("img_path", "")
|
||
# full_url = ""
|
||
# if image_path:
|
||
# base_url = os.getenv("HOST_IP", "http://172.18.30.122:9988") # 基础URL（不带/outputs）
|
||
# outputs_url = urljoin(base_url + "/", "output/")
|
||
# full_url = urljoin(outputs_url, image_path) # 最终拼接image_path
|
||
# if not text and not table_body and not image_path:
|
||
# continue
|
||
# image_caption = item.get("img_caption", "")
|
||
# content_chunk_list.append(
|
||
# {"id": idx, "type": text_type, "text_level": text_level, "text": text, "page_idx": page_idx,
|
||
# "table_caption": table_caption, "table_body": table_body, "image_path": full_url,
|
||
# "image_caption": image_caption})
|
||
# markdown_content =""
|
||
# for content in content_list:
|
||
# text = content["text"]
|
||
|
||
return full_content, content_chunk_list
|
||
|
||
except httpx.RequestError as e:
|
||
raise Exception(503, f"无法连接到解析服务: {str(e)}")
|
||
except Exception as e:
|
||
raise Exception(500, f"服务器内部错误: {str(e)}")
|
||
|
||
|
||
async def process_doc_logic(doc_file):
|
||
"""
|
||
接收Word文档(doc/docx)并转发到解析服务
|
||
"""
|
||
# 验证文件类型
|
||
allowed_extensions = {'.doc', '.docx'}
|
||
file_ext = os.path.splitext(doc_file.filename)[1].lower()
|
||
|
||
if file_ext not in allowed_extensions:
|
||
raise Exception(400, "仅支持doc或docx格式文件")
|
||
|
||
try:
|
||
# 配置目标接口参数
|
||
target_url = os.getenv("MINERU_URL", "http://172.18.30.122:9988")
|
||
target_url = target_url + "/analyze-doc"
|
||
# 根据文件类型设置正确的MIME类型
|
||
mime_type = "application/msword" if file_ext == '.doc' else "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
||
|
||
# 构建文件上传格式
|
||
files = {
|
||
"file": (doc_file.filename, doc_file.file, mime_type)
|
||
}
|
||
|
||
# 发送请求到解析服务
|
||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||
response = await client.post(
|
||
target_url,
|
||
files=files,
|
||
headers={"Accept": "application/json"}
|
||
)
|
||
response.raise_for_status()
|
||
result = response.json()
|
||
|
||
# 假设返回数据结构与PDF不同，需要相应调整
|
||
# 这里根据实际API返回结构调整
|
||
|
||
return result["data"]["full_content"]
|
||
|
||
|
||
except httpx.HTTPStatusError as e:
|
||
raise Exception(e.response.status_code, f"文档解析服务返回错误: {str(e)}")
|
||
except httpx.RequestError as e:
|
||
raise Exception(503, f"无法连接到解析服务: {str(e)}")
|
||
except Exception as e:
|
||
raise Exception(500, f"服务器内部错误: {str(e)}")
|
||
|
||
|
||
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:
|
||
# 创建临时文件保存上传的DOC
|
||
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:
|
||
# 创建临时文件保存转换后的DOCX
|
||
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:
|
||
# 创建临时文件保存转换后的PDF
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_pdf:
|
||
tmp_pdf.write(response.content)
|
||
output_path = tmp_pdf.name
|
||
|
||
# 验证PDF是否包含文字
|
||
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)}")
|
||
|