Spaces:

tomo2chin2
/

HTML2PDF_API

Paused

File size: 14,780 Bytes

import os
import uuid
import tempfile
import json
import re
from datetime import datetime
from pathlib import Path
from typing import Optional

import httpx
from fastapi import FastAPI, HTTPException, File, UploadFile
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from playwright.async_api import async_playwright
from huggingface_hub import HfApi, upload_file


class HTMLRequest(BaseModel):
    html_content: str


class YAMLRequest(BaseModel):
    yaml_content: str


class PDFResponse(BaseModel):
    pdf_url: str
    filename: str
    message: str
    repository_url: str


async def call_gemini_api(yaml_content: str) -> str:
    """
    Gemini APIを呼び出してYAMLコンテンツからHTMLを生成
    """
    try:
        # 環境変数から設定を取得
        gemini_api_key = os.getenv("GEMINI_API_KEY")
        model_id = os.getenv("MODEL", "gemini-2.5-flash")
        system_instruction = os.getenv("SYSTEM", "YAMLデータを基にHTMLを生成してください。")
        
        if not gemini_api_key:
            raise HTTPException(status_code=500, detail="GEMINI_API_KEY環境変数が設定されていません")
        
        # Gemini APIリクエストボディを構築（公式ドキュメント準拠のシンプルな形式）
        # システムインストラクションとYAMLデータを結合
        combined_prompt = f"{system_instruction}\n\n以下のYAMLデータを基にして、美しいHTMLドキュメントを生成してください。\n\n{yaml_content}"
        
        request_body = {
            "contents": [
                {
                    "parts": [
                        {
                            "text": combined_prompt
                        }
                    ]
                }
            ],
            "generationConfig": {
                "temperature": 0.75,
                "responseMimeType": "text/plain"
            }
        }
        
        # Gemini APIを呼び出し
        url = f"https://generativelanguage.googleapis.com/v1beta/models/{model_id}:generateContent?key={gemini_api_key}"
        
        async with httpx.AsyncClient(timeout=60.0) as client:
            response = await client.post(url, json=request_body)
            response.raise_for_status()
            
            # レスポンスからHTMLを抽出
            result = response.json()
            
            # レスポンス構造を解析してHTMLを抽出
            if "candidates" in result and len(result["candidates"]) > 0:
                candidate = result["candidates"][0]
                if "content" in candidate and "parts" in candidate["content"]:
                    parts = candidate["content"]["parts"]
                    for part in parts:
                        if "text" in part:
                            text = part["text"]
                            # HTMLタグを含む部分を抽出
                            html_match = re.search(r'<html.*?>.*?</html>', text, re.DOTALL | re.IGNORECASE)
                            if html_match:
                                return html_match.group(0)
                            # HTML全体がない場合はbodyタグを探す
                            body_match = re.search(r'<body.*?>.*?</body>', text, re.DOTALL | re.IGNORECASE)
                            if body_match:
                                return f"<html><head><meta charset='utf-8'></head>{body_match.group(0)}</html>"
                            # それでもない場合は全体をHTMLとして扱う
                            if "<" in text and ">" in text:
                                return text
            
            raise HTTPException(status_code=500, detail="Gemini APIレスポンスからHTMLを抽出できませんでした")
            
    except httpx.HTTPError as e:
        raise HTTPException(status_code=500, detail=f"Gemini API呼び出しエラー: {str(e)}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Gemini API処理中にエラーが発生しました: {str(e)}")


async def html_to_pdf_api(html_content: str) -> tuple[str, str]:
    """
    HTMLコンテンツをPDFに変換してHugging Faceデータセットリポジトリにアップロード
    knowledge.txtのPlaywright手法を使用（Async版）
    """
    try:
        # 環境変数からリポジトリ情報を取得
        hf_repo_id = os.getenv("HF_DATASET_REPO_ID")
        hf_token = os.getenv("HF_TOKEN")
        
        if not hf_repo_id:
            raise HTTPException(status_code=500, detail="HF_DATASET_REPO_ID環境変数が設定されていません")
        if not hf_token:
            raise HTTPException(status_code=500, detail="HF_TOKEN環境変数が設定されていません")
        
        # 一意のファイル名を生成
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        unique_id = str(uuid.uuid4())[:8]
        filename = f"document_{timestamp}_{unique_id}.pdf"
        
        # 一時ファイルでPDFを生成
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_file:
            temp_path = temp_file.name
            
            # Playwrightでヘッドレスブラウザを起動（Async版）
            async with async_playwright() as pw:
                browser = await pw.chromium.launch(headless=True)
                page = await browser.new_page()
                
                # HTMLコンテンツを設定（外部リソース読み込み待機）
                await page.set_content(html_content, wait_until="networkidle")
                
                # 印刷メディアを有効にする
                await page.emulate_media(media="print")
                
                # PDFを生成（test.htmlの設定に準拠）
                await page.pdf(
                    path=temp_path,
                    format="A4",
                    print_background=True,
                    margin={"top":"15mm","bottom":"15mm","left":"15mm","right":"15mm"},
                    scale=0.88  # 90%に縮小してA4 2ページに収める
                )
                
                await browser.close()
            
            # Hugging Face リポジトリにアップロード
            api = HfApi(token=hf_token)
            
            upload_file(
                path_or_fileobj=temp_path,
                path_in_repo=f"pdfs/{filename}",
                repo_id=hf_repo_id,
                repo_type="dataset",
                token=hf_token,
                commit_message=f"Add PDF: {filename}"
            )
            
            # 一時ファイルを削除
            os.unlink(temp_path)
            
            # ダウンロードURLを生成
            pdf_url = f"https://huggingface.co/datasets/{hf_repo_id}/resolve/main/pdfs/{filename}"
            
            return filename, pdf_url
        
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"PDF生成中にエラーが発生しました: {str(e)}")


# FastAPIアプリケーションの初期化
app = FastAPI(
    title="HTML to PDF Converter API",
    description="日本語対応のHTML→PDF変換API。複雑なレイアウトやWebフォントを含むHTMLコンテンツを、正確にA4サイズのPDFに変換します。",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc"
)

# CORS設定
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/")
async def root():
    """
    API情報を返すルートエンドポイント
    """
    hf_repo_id = os.getenv("HF_DATASET_REPO_ID", "未設定")
    return {
        "message": "HTML to PDF Converter API",
        "description": "HTMLコンテンツをA4サイズのPDFに変換し、Hugging Faceデータセットリポジトリに保存するAPI",
        "version": "1.0.0",
        "storage": f"Hugging Face Dataset Repository: {hf_repo_id}",
        "endpoints": {
            "convert": "/convert - HTMLをPDFに変換してHFリポジトリに保存",
            "convert-yaml": "/convert-yaml - YAMLをGemini APIでHTMLに変換してPDF化",
            "files": "/files - HFリポジトリ内のPDFファイル一覧",
            "docs": "/docs - API仕様書",
            "health": "/health - ヘルスチェック"
        }
    }


@app.post("/convert", response_model=PDFResponse)
async def convert_html_to_pdf(request: HTMLRequest):
    """
    HTMLコンテンツをPDFに変換してHugging Faceデータセットリポジトリに保存するエンドポイント
    
    - **html_content**: 変換するHTMLコンテンツ
    
    Returns:
    - **pdf_url**: 生成されたPDFのダウンロードURL（Hugging Face上）
    - **filename**: PDFファイル名
    - **message**: 処理結果メッセージ
    - **repository_url**: リポジトリURL
    """
    if not request.html_content or not request.html_content.strip():
        raise HTTPException(status_code=400, detail="HTMLコンテンツが空です")
    
    try:
        filename, pdf_url = await html_to_pdf_api(request.html_content)
        hf_repo_id = os.getenv("HF_DATASET_REPO_ID")
        repository_url = f"https://huggingface.co/datasets/{hf_repo_id}"
        
        return PDFResponse(
            pdf_url=pdf_url,
            filename=filename,
            message=f"✅ PDFが正常に生成され、Hugging Faceリポジトリに保存されました: {filename}",
            repository_url=repository_url
        )
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"処理中にエラーが発生しました: {str(e)}")


@app.get("/download/{filename}")
async def download_pdf(filename: str):
    """
    Hugging FaceリポジトリのPDFファイルへリダイレクトするエンドポイント
    
    - **filename**: ダウンロードするPDFファイル名
    """
    hf_repo_id = os.getenv("HF_DATASET_REPO_ID")
    if not hf_repo_id:
        raise HTTPException(status_code=500, detail="HF_DATASET_REPO_ID環境変数が設定されていません")
    
    # Hugging Face上のファイルURLにリダイレクト
    pdf_url = f"https://huggingface.co/datasets/{hf_repo_id}/resolve/main/pdfs/{filename}"
    return RedirectResponse(url=pdf_url)


@app.get("/health")
async def health_check():
    """
    ヘルスチェックエンドポイント
    """
    return {"status": "healthy", "timestamp": datetime.now().isoformat()}


@app.post("/convert-yaml", response_model=PDFResponse)
async def convert_yaml_to_pdf(request: YAMLRequest):
    """
    YAMLコンテンツをGemini APIでHTMLに変換し、PDFを生成してHugging Faceデータセットリポジトリに保存するエンドポイント
    
    - **yaml_content**: 変換するYAMLコンテンツ
    
    Returns:
    - **pdf_url**: 生成されたPDFのダウンロードURL（Hugging Face上）
    - **filename**: PDFファイル名
    - **message**: 処理結果メッセージ
    - **repository_url**: リポジトリURL
    """
    if not request.yaml_content or not request.yaml_content.strip():
        raise HTTPException(status_code=400, detail="YAMLコンテンツが空です")
    
    try:
        # YAMLからGemini APIでHTMLを生成
        html_content = await call_gemini_api(request.yaml_content)
        
        # 生成されたHTMLをPDFに変換
        filename, pdf_url = await html_to_pdf_api(html_content)
        hf_repo_id = os.getenv("HF_DATASET_REPO_ID")
        repository_url = f"https://huggingface.co/datasets/{hf_repo_id}"
        
        return PDFResponse(
            pdf_url=pdf_url,
            filename=filename,
            message=f"✅ YAMLからPDFが正常に生成され、Hugging Faceリポジトリに保存されました: {filename}",
            repository_url=repository_url
        )
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"処理中にエラーが発生しました: {str(e)}")


@app.get("/files")
async def list_files():
    """
    Hugging Faceリポジトリ内のPDFファイル一覧を取得するエンドポイント
    """
    try:
        hf_repo_id = os.getenv("HF_DATASET_REPO_ID")
        hf_token = os.getenv("HF_TOKEN")
        
        if not hf_repo_id:
            raise HTTPException(status_code=500, detail="HF_DATASET_REPO_ID環境変数が設定されていません")
        if not hf_token:
            raise HTTPException(status_code=500, detail="HF_TOKEN環境変数が設定されていません")
        
        api = HfApi(token=hf_token)
        
        # リポジトリ内のファイル一覧を取得
        try:
            repo_files = api.list_repo_files(repo_id=hf_repo_id, repo_type="dataset")
            pdf_files = [f for f in repo_files if f.startswith("pdfs/") and f.endswith(".pdf")]
            
            files = []
            for file_path in pdf_files:
                filename = Path(file_path).name
                file_info = api.get_paths_info(repo_id=hf_repo_id, paths=[file_path], repo_type="dataset")[0]
                
                files.append({
                    "filename": filename,
                    "path": file_path,
                    "size": file_info.size if hasattr(file_info, 'size') else 0,
                    "last_modified": file_info.last_commit.date.isoformat() if hasattr(file_info, 'last_commit') and file_info.last_commit else None,
                    "download_url": f"https://huggingface.co/datasets/{hf_repo_id}/resolve/main/{file_path}",
                    "api_download_url": f"/download/{filename}"
                })
            
            return {
                "repository": hf_repo_id,
                "total_files": len(files),
                "files": files
            }
            
        except Exception as e:
            # リポジトリが空の場合やpdfsフォルダが存在しない場合
            return {
                "repository": hf_repo_id,
                "total_files": 0,
                "files": [],
                "note": "No PDF files found or repository is empty"
            }
            
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"ファイル一覧取得中にエラーが発生しました: {str(e)}")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)