worker.py

"""
Dolphin AI Worker Service
-------------------------
Standalone FastAPI microservice for hosting the Dolphin-v2 model on Hugging Face Spaces.
Accepts PDF uploads, runs the full Dolphin pipeline, and returns structured data.
"""

import os
import shutil
import tempfile
import uvicorn
from fastapi import FastAPI, UploadFile, File, HTTPException, Depends
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
from typing import List, Dict, Any, Optional

# Import local Dolphin modules (we will copy these into the container)
# The Dockerfile should ensure visique/backend is in PYTHONPATH or copied correctly
try:
    from app.services.ingestion.dolphin.client import DolphinClient, DolphinDocumentResult
    from app.services.ingestion.dolphin.extractor import DolphinExtractor
except ImportError:
    # Fallback for when running relative to worker dir
    import sys
    sys.path.append(os.path.join(os.path.dirname(__file__), "..", "backend"))
    from app.services.ingestion.dolphin.client import DolphinClient, DolphinDocumentResult


# --- Configuration ---
API_KEY = os.getenv("DOLPHIN_API_KEY") 
security = HTTPBearer(auto_error=False)

app = FastAPI(title="Dolphin AI Worker", version="1.0.0")

# --- Global Model Instance ---
# Load model once on startup
dolphin_client = None

@app.on_event("startup")
def load_model():
    global dolphin_client
    # HF Spaces Free Tier: Force CPU + float32 via environment
    device = os.getenv("DOLPHIN_DEVICE", "cpu")
    max_batch = int(os.getenv("DOLPHIN_MAX_BATCH_SIZE", "1"))
    print(f"🚀 Loading Dolphin-v2 model on device={device}, batch={max_batch}...")
    dolphin_client = DolphinClient(device=device, max_batch_size=max_batch)
    print("✅ Model loaded and ready.")


# --- Auth Dependency ---
async def verify_token(credentials: Optional[HTTPAuthorizationCredentials] = Depends(security)):
    if not API_KEY:
        return True # Open access if no key set (default for private spaces)
    
    if not credentials or credentials.credentials != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API Key")
    return True


# --- Endpoints ---

@app.get("/")
def health_check():
    return {"status": "ok", "model_loaded": dolphin_client is not None}

@app.post("/process", dependencies=[Depends(verify_token)])
async def process_pdf(file: UploadFile = File(...)):
    if not dolphin_client:
        raise HTTPException(status_code=503, detail="Model functionality not ready")
        
    if file.content_type != "application/pdf":
        raise HTTPException(status_code=400, detail="File must be a PDF")

    # Save uploaded file to temp
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
        shutil.copyfileobj(file.file, tmp)
        tmp_path = tmp.name

    try:
        print(f"📄 Processing: {file.filename}")
        
        # Run Dolphin Parsing
        result: DolphinDocumentResult = dolphin_client.parse_document(tmp_path)
        
        # Convert to dict for JSON response
        response_data = {
            "total_pages": result.total_pages,
            "full_markdown": result.full_markdown,
            "pages": [],
            "layouts": []
        }
        
        for p in result.pages:
            response_data["pages"].append({
                "page_number": p.page_number,
                "markdown": p.markdown,
                "structured_json": p.structured_json,
                "elements": [
                    {
                        "element_type": e.element_type,
                        "content": e.content,
                        "bbox": e.bbox,
                        "metadata": e.metadata
                    } for e in p.elements
                ]
            })
            
        for l in result.layouts:
            response_data["layouts"].append({
                "page_number": l.page_number,
                "sections": l.sections,
                "reading_order": l.reading_order,
                "doc_type_hint": l.doc_type_hint
            })
            
        return response_data

    except Exception as e:
        print(f"❌ Error processing {file.filename}: {e}")
        raise HTTPException(status_code=500, detail=str(e))
        
    finally:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)


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