app.py

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from gradio_client import Client
import re
import time

app = FastAPI(title="PrivBERT Proxy")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

PRIVBERT_SPACE = "shikhasoneji8/privacywhisper"


class AnalyzeRequest(BaseModel):
    text: str
    threshold: float = 0.75


def extract_choices(result_value) -> list[str]:
    """Extract topic strings from a Gradio Dropdown result (handles dict or list)."""
    if isinstance(result_value, dict):
        choices = result_value.get("choices", [])
        return [c[0] if isinstance(c, list) else c for c in choices if c and (c[0] if isinstance(c, list) else c) != "—"]
    if isinstance(result_value, list):
        return [c for c in result_value if c and c != "—"]
    return []


def count_clauses(html: str) -> int:
    return len(re.findall(r"class='card ", html))


@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/analyze")
def analyze(req: AnalyzeRequest):
    if not req.text or len(req.text.strip()) < 50:
        raise HTTPException(status_code=400, detail="Text too short — provide at least 50 characters of policy text")

    # Each request gets its own client so sessions don't bleed across requests.
    # Retry up to 3 times — HF free-tier spaces can take a moment to wake up.
    client = None
    last_error = ""
    for attempt in range(3):
        try:
            client = Client(PRIVBERT_SPACE, verbose=False)
            break
        except Exception as e:
            last_error = str(e)
            if attempt < 2:
                time.sleep(5)

    if client is None:
        raise HTTPException(status_code=503, detail=f"Could not connect to PrivBERT model after 3 attempts: {last_error}")

    # Step 1: Run annotation — classifies text into privacy topics
    try:
        result1 = client.predict(
            text_input=req.text,
            threshold=req.threshold,
            api_name="/_run_annotation",
        )
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"_run_annotation failed: {str(e)}")

    # result1 is a tuple: (topics_result, granular_result, markdown_placeholder)
    topics = extract_choices(result1[0]) if len(result1) > 0 else []
    granular_topics = extract_choices(result1[1]) if len(result1) > 1 else []

    # Step 2: Render annotated HTML — same client = same session = server has state from step 1
    annotated_html = ""
    try:
        result2 = client.predict(
            selected_topics=topics,
            selected_fine=granular_topics,
            selected_ratings=[],
            selected_compliance=[],
            view_mode="Cards",
            fsize=15,
            enable_expl=False,
            api_name="/_apply_filters_and_render",
        )
        annotated_html = result2[0] if result2 and isinstance(result2[0], str) else ""
    except Exception as e:
        # Annotated HTML is best-effort — still return topics if this fails
        print(f"Warning: _apply_filters_and_render failed: {e}")

    return {
        "success": True,
        "topics": topics,
        "granular_topics": granular_topics,
        "annotated_html": annotated_html,
        "clause_count": count_clauses(annotated_html),
    }


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