103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
|
"""
|
|||
|
|
本地 PDF 转 Markdown 微服务，兼容 langchain-chat 中 file_converter.pdf_to_html 的调用约定：
|
|||
|
|
|
|||
|
|
POST /convert/
|
|||
|
|
Body: {"pdf_path": "<相对于知识库根目录的路径>"}
|
|||
|
|
成功: {"status": "success", "markdown_path": "<生成的 .md 绝对路径>"}
|
|||
|
|
失败: {"status": "error", "message": "..."} + HTTP 4xx/5xx
|
|||
|
|
|
|||
|
|
环境变量:
|
|||
|
|
PDF_CONVERT_KB_ROOT 知识库根目录，须与 configs.kb_config.PDF_CONVERT_KB_ROOT 一致
|
|||
|
|
PDF_CONVERT_PORT 监听端口，默认 6006
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import logging
|
|||
|
|
import os
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
import fitz # PyMuPDF
|
|||
|
|
from fastapi import FastAPI, HTTPException
|
|||
|
|
from pydantic import BaseModel, Field
|
|||
|
|
import uvicorn
|
|||
|
|
|
|||
|
|
logger = logging.getLogger("pdf_convert_service")
|
|||
|
|
logging.basicConfig(level=logging.INFO)
|
|||
|
|
|
|||
|
|
# 与 langchain-chat 默认知识库目录一致（本文件位于 langchain-chat/pdf_convert_service/app.py）
|
|||
|
|
_LC_ROOT = Path(__file__).resolve().parent.parent
|
|||
|
|
_DEFAULT_KB = _LC_ROOT / "knowledge_base"
|
|||
|
|
KB_ROOT = Path(os.environ.get("PDF_CONVERT_KB_ROOT", str(_DEFAULT_KB))).resolve()
|
|||
|
|
|
|||
|
|
app = FastAPI(title="PDF Convert Service", version="1.0.0")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ConvertRequest(BaseModel):
|
|||
|
|
pdf_path: str = Field(..., description="相对于知识库根的路径，如 kb_name/file.pdf")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _under_kb_root(full: Path) -> bool:
|
|||
|
|
try:
|
|||
|
|
full.resolve().relative_to(KB_ROOT)
|
|||
|
|
return True
|
|||
|
|
except ValueError:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _pdf_to_markdown_file(pdf_abs: Path, md_abs: Path) -> None:
|
|||
|
|
doc = fitz.open(pdf_abs)
|
|||
|
|
try:
|
|||
|
|
chunks: list[str] = []
|
|||
|
|
for i in range(len(doc)):
|
|||
|
|
page = doc.load_page(i)
|
|||
|
|
text = page.get_text()
|
|||
|
|
chunks.append(f"\n\n## 第 {i + 1} 页\n\n{text.strip()}\n")
|
|||
|
|
body = "".join(chunks).strip()
|
|||
|
|
if not body:
|
|||
|
|
body = "_（未能从 PDF 提取到文本，可能是扫描件或加密文档）_"
|
|||
|
|
md_abs.parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
md_abs.write_text(body, encoding="utf-8")
|
|||
|
|
finally:
|
|||
|
|
doc.close()
|
|||
|
|
|
|||
|
|
|
|||
|
|
@app.get("/health")
|
|||
|
|
def health():
|
|||
|
|
return {"status": "ok", "kb_root": str(KB_ROOT)}
|
|||
|
|
|
|||
|
|
|
|||
|
|
@app.post("/convert/")
|
|||
|
|
@app.post("/convert")
|
|||
|
|
def convert(req: ConvertRequest):
|
|||
|
|
rel = req.pdf_path.replace("\\", "/").lstrip("/")
|
|||
|
|
if ".." in rel.split("/"):
|
|||
|
|
raise HTTPException(status_code=400, detail="invalid pdf_path")
|
|||
|
|
|
|||
|
|
pdf_abs = (KB_ROOT / rel).resolve()
|
|||
|
|
if not _under_kb_root(pdf_abs):
|
|||
|
|
raise HTTPException(status_code=403, detail="pdf_path outside knowledge base root")
|
|||
|
|
if not pdf_abs.is_file():
|
|||
|
|
raise HTTPException(status_code=404, detail=f"file not found: {pdf_abs}")
|
|||
|
|
if pdf_abs.suffix.lower() != ".pdf":
|
|||
|
|
raise HTTPException(status_code=400, detail="not a pdf file")
|
|||
|
|
|
|||
|
|
md_abs = pdf_abs.with_suffix(".md")
|
|||
|
|
try:
|
|||
|
|
_pdf_to_markdown_file(pdf_abs, md_abs)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.exception("pdf convert failed")
|
|||
|
|
raise HTTPException(status_code=500, detail=str(e)) from e
|
|||
|
|
|
|||
|
|
return {"status": "success", "markdown_path": str(md_abs)}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
port = int(os.environ.get("PDF_CONVERT_PORT", "6006"))
|
|||
|
|
host = os.environ.get("PDF_CONVERT_HOST", "127.0.0.1")
|
|||
|
|
logger.info("KB_ROOT=%s listen=%s:%s", KB_ROOT, host, port)
|
|||
|
|
uvicorn.run(app, host=host, port=port, log_level="info")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|