313 lines
10 KiB
Python
313 lines
10 KiB
Python
"""PDF → Markdown 转换服务（FastAPI + LangGraph）。
|
||
|
||
接口：
|
||
GET / 网页
|
||
GET /api/health 健康检查
|
||
POST /api/upload 上传 PDF 并启动解析（multipart file）
|
||
GET /api/tasks 任务列表
|
||
GET /api/task/{id} 任务详情
|
||
GET /api/task/{id}/markdown Markdown 原文
|
||
GET /api/task/{id}/page/{n} 页面渲染图
|
||
GET /api/task/{id}/images/{f} 图片资源
|
||
GET /api/task/{id}/pdf 原始 PDF
|
||
GET /api/task/{id}/download 打包下载（output.md + images）
|
||
POST /api/task/{id}/rerun 重新解析
|
||
DELETE /api/task/{id} 删除任务
|
||
"""
|
||
import io
|
||
import json
|
||
import os
|
||
import shutil
|
||
import threading
|
||
import time
|
||
import uuid
|
||
import zipfile
|
||
from datetime import datetime
|
||
|
||
from fastapi import FastAPI, File, HTTPException, UploadFile
|
||
from fastapi.responses import FileResponse, JSONResponse, Response
|
||
from fastapi.staticfiles import StaticFiles
|
||
|
||
import config
|
||
from pipeline.graph import compiled_graph
|
||
|
||
app = FastAPI(title="PDF → Markdown", docs_url="/docs")
|
||
config.ensure_dirs()
|
||
|
||
REGISTRY = os.path.join(config.DATA_DIR, "tasks.json")
|
||
_lock = threading.Lock()
|
||
|
||
|
||
# ---------------- 任务注册表 ----------------
|
||
def _load_registry():
|
||
if os.path.exists(REGISTRY):
|
||
try:
|
||
with open(REGISTRY, encoding="utf-8") as f:
|
||
return json.load(f)
|
||
except Exception:
|
||
return {}
|
||
return {}
|
||
|
||
|
||
def _save_registry(reg):
|
||
with open(REGISTRY, "w", encoding="utf-8") as f:
|
||
json.dump(reg, f, ensure_ascii=False, indent=2)
|
||
|
||
|
||
def _task_dir(task_id):
|
||
return os.path.join(config.TASKS_DIR, task_id)
|
||
|
||
|
||
def _safe_name(name: str) -> str:
|
||
return "".join(c for c in (name or "upload.pdf") if c.isalnum() or c in "._- ").strip() or "upload.pdf"
|
||
|
||
|
||
# ---------------- 解析执行 ----------------
|
||
def _run_parse(task_id, pdf_path, name):
|
||
out_dir = _task_dir(task_id)
|
||
images_dir = os.path.join(out_dir, "images")
|
||
os.makedirs(images_dir, exist_ok=True)
|
||
initial = {
|
||
"task_id": task_id,
|
||
"pdf_path": pdf_path,
|
||
"out_dir": out_dir,
|
||
"images_dir": images_dir,
|
||
"status": "running",
|
||
"logs": [f"开始解析：{name}"],
|
||
}
|
||
t0 = time.time()
|
||
try:
|
||
result = compiled_graph.invoke(initial)
|
||
dt = time.time() - t0
|
||
if result.get("status") == "error":
|
||
with _lock:
|
||
reg = _load_registry()
|
||
if task_id in reg:
|
||
reg[task_id]["status"] = "failed"
|
||
reg[task_id]["error"] = result.get("error")
|
||
reg[task_id]["updated"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
_save_registry(reg)
|
||
return
|
||
# 更新清单耗时
|
||
mp = os.path.join(out_dir, "manifest.json")
|
||
if os.path.exists(mp):
|
||
try:
|
||
with open(mp, encoding="utf-8") as f:
|
||
man = json.load(f)
|
||
man["elapsed_sec"] = round(dt, 1)
|
||
man["status"] = "done"
|
||
with open(mp, "w", encoding="utf-8") as f:
|
||
json.dump(man, f, ensure_ascii=False, indent=2)
|
||
except Exception:
|
||
pass
|
||
with _lock:
|
||
reg = _load_registry()
|
||
if task_id in reg:
|
||
reg[task_id]["status"] = "done"
|
||
reg[task_id]["elapsed_sec"] = round(dt, 1)
|
||
reg[task_id]["updated"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
_save_registry(reg)
|
||
except Exception as e:
|
||
with _lock:
|
||
reg = _load_registry()
|
||
if task_id in reg:
|
||
reg[task_id]["status"] = "failed"
|
||
reg[task_id]["error"] = str(e)
|
||
reg[task_id]["updated"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
_save_registry(reg)
|
||
|
||
|
||
def _start_parse(task_id, pdf_path, name):
|
||
t = threading.Thread(target=_run_parse, args=(task_id, pdf_path, name), daemon=True)
|
||
t.start()
|
||
|
||
|
||
# ---------------- 路由 ----------------
|
||
@app.get("/api/health")
|
||
def health():
|
||
return {"ok": True, "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
|
||
|
||
|
||
@app.post("/api/upload")
|
||
async def upload(file: UploadFile = File(...)):
|
||
name = _safe_name(file.filename)
|
||
if not name.lower().endswith(".pdf"):
|
||
raise HTTPException(400, "仅支持 PDF 文件")
|
||
task_id = uuid.uuid4().hex[:12]
|
||
out_dir = _task_dir(task_id)
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
pdf_path = os.path.join(out_dir, "input.pdf")
|
||
content = await file.read()
|
||
if len(content) < 50 or not content.startswith(b"%PDF"):
|
||
shutil.rmtree(out_dir, ignore_errors=True)
|
||
raise HTTPException(400, "文件不是有效的 PDF")
|
||
with open(pdf_path, "wb") as f:
|
||
f.write(content)
|
||
|
||
with _lock:
|
||
reg = _load_registry()
|
||
reg[task_id] = {
|
||
"task_id": task_id,
|
||
"name": name,
|
||
"status": "running",
|
||
"created": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||
"updated": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||
"size_kb": round(len(content) / 1024, 1),
|
||
"error": None,
|
||
}
|
||
_save_registry(reg)
|
||
_start_parse(task_id, pdf_path, name)
|
||
return {"task_id": task_id, "name": name, "status": "running"}
|
||
|
||
|
||
@app.get("/api/tasks")
|
||
def tasks():
|
||
reg = _load_registry()
|
||
items = []
|
||
for tid, info in sorted(reg.items(), key=lambda x: x[1].get("created", ""), reverse=True):
|
||
d = dict(info)
|
||
out = _task_dir(tid)
|
||
d["has_markdown"] = os.path.exists(os.path.join(out, "output.md"))
|
||
n_imgs = 0
|
||
img_dir = os.path.join(out, "images")
|
||
if os.path.isdir(img_dir):
|
||
n_imgs = len([x for x in os.listdir(img_dir) if x.lower().endswith((".png", ".jpg", ".jpeg"))])
|
||
d["images"] = n_imgs
|
||
if os.path.exists(os.path.join(out, "manifest.json")):
|
||
try:
|
||
with open(os.path.join(out, "manifest.json"), encoding="utf-8") as f:
|
||
man = json.load(f)
|
||
d["total_pages"] = man.get("total_pages")
|
||
d["page_types"] = man.get("page_types")
|
||
except Exception:
|
||
pass
|
||
items.append(d)
|
||
return {"tasks": items}
|
||
|
||
|
||
@app.get("/api/task/{task_id}")
|
||
def task_detail(task_id: str):
|
||
out = _task_dir(task_id)
|
||
if not os.path.isdir(out):
|
||
raise HTTPException(404, "任务不存在")
|
||
reg = _load_registry()
|
||
info = dict(reg.get(task_id, {}))
|
||
info["task_id"] = task_id
|
||
info["out_dir"] = out
|
||
man_path = os.path.join(out, "manifest.json")
|
||
if os.path.exists(man_path):
|
||
try:
|
||
with open(man_path, encoding="utf-8") as f:
|
||
info["manifest"] = json.load(f)
|
||
except Exception:
|
||
info["manifest"] = None
|
||
if not os.path.exists(man_path) and info.get("status") == "running":
|
||
info["status"] = "running"
|
||
info["logs"] = (info.get("manifest") or {}).get("logs") or []
|
||
return info
|
||
|
||
|
||
@app.get("/api/task/{task_id}/markdown")
|
||
def task_markdown(task_id: str):
|
||
p = os.path.join(_task_dir(task_id), "output.md")
|
||
if not os.path.exists(p):
|
||
raise HTTPException(404, "还没有生成 Markdown")
|
||
with open(p, encoding="utf-8") as f:
|
||
return Response(f.read(), media_type="text/markdown; charset=utf-8")
|
||
|
||
|
||
@app.get("/api/task/{task_id}/page/{n}")
|
||
def task_page(task_id: str, n: int):
|
||
p = os.path.join(_task_dir(task_id), "pages", f"page_{n:03d}.png")
|
||
if not os.path.exists(p):
|
||
raise HTTPException(404, "页面不存在")
|
||
return FileResponse(p, media_type="image/png")
|
||
|
||
|
||
@app.get("/api/task/{task_id}/images/{fname}")
|
||
def task_image(task_id: str, fname: str):
|
||
p = os.path.join(_task_dir(task_id), "images", os.path.basename(fname))
|
||
if not os.path.exists(p):
|
||
raise HTTPException(404, "图片不存在")
|
||
return FileResponse(p)
|
||
|
||
|
||
@app.get("/api/task/{task_id}/pdf")
|
||
def task_pdf(task_id: str):
|
||
p = os.path.join(_task_dir(task_id), "input.pdf")
|
||
if not os.path.exists(p):
|
||
raise HTTPException(404, "PDF 不存在")
|
||
reg = _load_registry()
|
||
name = reg.get(task_id, {}).get("name", "input.pdf")
|
||
return FileResponse(p, media_type="application/pdf", filename=name)
|
||
|
||
|
||
@app.get("/api/task/{task_id}/download")
|
||
def task_download(task_id: str):
|
||
out = _task_dir(task_id)
|
||
if not os.path.isdir(out):
|
||
raise HTTPException(404, "任务不存在")
|
||
buf = io.BytesIO()
|
||
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as z:
|
||
for root, _, files in os.walk(out):
|
||
for fn in files:
|
||
fp = os.path.join(root, fn)
|
||
if fn == "input.pdf":
|
||
continue
|
||
arc = os.path.relpath(fp, out)
|
||
z.write(fp, arc)
|
||
buf.seek(0)
|
||
reg = _load_registry()
|
||
name = (reg.get(task_id, {}).get("name") or "task").replace(".pdf", "")
|
||
return Response(buf.getvalue(), media_type="application/zip",
|
||
headers={"Content-Disposition": f'attachment; filename="{name}_markdown.zip"'})
|
||
|
||
|
||
@app.post("/api/task/{task_id}/rerun")
|
||
def task_rerun(task_id: str):
|
||
out = _task_dir(task_id)
|
||
pdf_path = os.path.join(out, "input.pdf")
|
||
if not os.path.exists(pdf_path):
|
||
raise HTTPException(404, "PDF 不存在")
|
||
# 清掉旧产出
|
||
for d in ("images", "pages"):
|
||
shutil.rmtree(os.path.join(out, d), ignore_errors=True)
|
||
for fn in os.listdir(out):
|
||
if fn in ("input.pdf",):
|
||
continue
|
||
fp = os.path.join(out, fn)
|
||
if os.path.isfile(fp):
|
||
os.remove(fp)
|
||
else:
|
||
shutil.rmtree(fp, ignore_errors=True)
|
||
reg = _load_registry()
|
||
if task_id in reg:
|
||
reg[task_id]["status"] = "running"
|
||
reg[task_id]["error"] = None
|
||
reg[task_id]["updated"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
_save_registry(reg)
|
||
_start_parse(task_id, pdf_path, reg.get(task_id, {}).get("name", "input.pdf"))
|
||
return {"task_id": task_id, "status": "running"}
|
||
|
||
|
||
@app.delete("/api/task/{task_id}")
|
||
def task_delete(task_id: str):
|
||
out = _task_dir(task_id)
|
||
if not os.path.isdir(out):
|
||
raise HTTPException(404, "任务不存在")
|
||
shutil.rmtree(out, ignore_errors=True)
|
||
with _lock:
|
||
reg = _load_registry()
|
||
reg.pop(task_id, None)
|
||
_save_registry(reg)
|
||
return {"ok": True}
|
||
|
||
|
||
# 静态前端（放在最后，避免覆盖 API）
|
||
app.mount("/", StaticFiles(directory=os.path.join(config.BASE_DIR, "static"), html=True), name="static")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import uvicorn
|
||
uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "16101")), log_level="info")
|