app.py

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import requests
import os
import logging
import threading
import time

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# Enable CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allow all origins, or replace with your frontend URL
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Read OpenRouter API key from environment variable (HF Secret)
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
if not OPENROUTER_API_KEY:
    raise RuntimeError("OPENROUTER_API_KEY environment variable not set!")

# ---------------------------------------------------------------------------
# Global rate limiter for OpenRouter calls
#
# Free-tier OpenRouter models throttle/reject requests that arrive too close
# together. Since /classify and /explain can be hit concurrently by the
# frontend (each running in its own thread pool worker), we serialize ALL
# outbound OpenRouter calls behind a single lock and enforce a minimum gap
# between them. Whichever request arrives second simply waits its turn.
# ---------------------------------------------------------------------------
MIN_INTERVAL_SECONDS = float(os.environ.get("OPENROUTER_MIN_INTERVAL", "6.0"))
MAX_RETRIES = int(os.environ.get("OPENROUTER_MAX_RETRIES", "3"))
RETRY_BACKOFF_SECONDS = float(os.environ.get("OPENROUTER_RETRY_BACKOFF", "5.0"))

_or_lock = threading.Lock()
_last_call_time = 0.0

def _is_rate_limited(response: requests.Response) -> bool:
    """Detect a 429 whether it comes as an HTTP status or embedded in a 200 body
    (OpenRouter sometimes returns 200 with {"error": {"code": 429, ...}})."""
    if response.status_code == 429:
        return True
    try:
        body = response.json()
    except ValueError:
        return False
    return isinstance(body, dict) and body.get("error", {}).get("code") == 429


def call_openrouter(payload: dict, headers: dict, timeout: int = 20):
    """Thread-safe, rate-limited call to OpenRouter's chat completions endpoint.

    - Ensures at least MIN_INTERVAL_SECONDS has elapsed since the previous
      OpenRouter call (across ALL endpoints) before firing this one.
    - If the provider still returns a 429 (as a status code or embedded in the
      JSON body), retries with exponential backoff up to MAX_RETRIES times.
    """
    global _last_call_time

    url = "https://openrouter.ai/api/v1/chat/completions"

    with _or_lock:
        now = time.monotonic()
        elapsed = now - _last_call_time
        wait_for = MIN_INTERVAL_SECONDS - elapsed
        if wait_for > 0:
            logger.info(f"Rate limiting: waiting {wait_for:.2f}s before next OpenRouter call")
            time.sleep(wait_for)

        response = requests.post(url, headers=headers, json=payload, timeout=timeout)
        _last_call_time = time.monotonic()

        attempt = 0
        while _is_rate_limited(response) and attempt < MAX_RETRIES:
            backoff = RETRY_BACKOFF_SECONDS * (2 ** attempt)
            logger.warning(f"OpenRouter rate limited (attempt {attempt + 1}/{MAX_RETRIES}); "
                            f"backing off {backoff:.1f}s before retry")
            time.sleep(backoff)
            response = requests.post(url, headers=headers, json=payload, timeout=timeout)
            _last_call_time = time.monotonic()
            attempt += 1

    return response


# Request schema
class ExplainRequest(BaseModel):
    message: str
    label: str = None  # Optional, for explanations
    model_id: str = "openai/gpt-oss-120b:free"  # Default model, can be overridden

@app.get("/")
def root():
    return {
        "status": "healthy",
        "service": "Anti-Phishing Explainer",
        "endpoints": {
            "/explain": "POST - Generate explanation for classification",
            "/classify": "POST - Classify text as Phishing or Safe",
            "/health": "GET - Health check"
        }
    }

@app.get("/health")
def health():
    return {
        "status": "healthy",
        "service": "explainer",
        "openrouter_configured": bool(OPENROUTER_API_KEY)
    }

@app.post("/explain")
def explain(req: ExplainRequest):
    """Generate a human-readable explanation for why a message was classified as Phishing or Safe"""
    user_message = req.message.strip()
    label = req.label.strip() if req.label else None

    if not user_message or not label:
        raise HTTPException(status_code=400, detail="Missing message or label")

    # Updated system prompt with bullet-point format and language adaptation
    system_prompt = (
        "Respond using the same language as the message."
        f"You are a robot that identifies phishing and safe messages. "
        f"The message was classified as '{label}'. "
        "Explain why this decision was made and point out any words or patterns that led to it. "
        "No greetings, introductions, or closing remarks. "
        "Don't restate the message or its classification. "
        "Output only the explanation as bullet points. "
        "Limit each bullet to 1–2 sentences. "
        "Limit the number of bullets to 3-4. "
        f"Message:\n\n{user_message}\n\n"
    )

    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": req.model_id,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message}
        ]
    }

    try:
        logger.info(f"Calling OpenRouter /explain with model: {req.model_id}")
        response = call_openrouter(payload, headers)
        response.raise_for_status()
        result = response.json()

        # Check for OpenRouter error in response body (even with 200 status)
        if "error" in result:
            error_detail = result.get("error", {}).get("message", str(result.get("error")))
            logger.error(f"OpenRouter returned error in /explain: {error_detail}")
            logger.error(f"Full response: {result}")
            raise HTTPException(status_code=500, detail=f"OpenRouter error: {error_detail}")

        reply = result.get("choices", [{}])[0].get("message", {}).get("content", "").strip()
        if not reply:
            logger.error(f"OpenRouter returned empty response in /explain. Full result: {result}")
            reply = "[No explanation returned]"
        logger.info("Explanation generated successfully")
        return {"reply": reply}
    except requests.RequestException as e:
        logger.error(f"OpenRouter network error in /explain: {e}")
        raise HTTPException(status_code=500, detail=f"Error contacting OpenRouter: {e}")
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Unexpected error in /explain: {e}")
        raise HTTPException(status_code=500, detail=f"Unexpected error: {e}")

@app.post("/classify")
def classify(req: ExplainRequest):
    """Classify text as Phishing or Safe via OpenRouter"""
    user_message = req.message.strip()
    if not user_message:
        raise HTTPException(status_code=400, detail="Missing message")

    # Use provided model_id or default
    model_id = req.model_id or "arcee-ai/trinity-large-preview:free"

    system_prompt = (
        'You are a phishing detector. Classify the text as "Phishing" or "Safe". '
        'Respond ONLY with valid JSON: {"label": "Phishing"|"Safe", "confidence": <0-100 float>}. '
        'No other text.'
    )

    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": model_id,  # Use the passed model_id
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message}
        ],
        "temperature": 0.0,
        "max_tokens": 2000  # Increased for models with extended reasoning/thinking
    }

    try:
        logger.info(f"Calling OpenRouter /classify with model: {model_id}")
        response = call_openrouter(payload, headers)
        response.raise_for_status()
        result = response.json()

        # Check for OpenRouter error in response body (even with 200 status)
        if "error" in result:
            error_detail = result.get("error", {}).get("message", str(result.get("error")))
            logger.error(f"OpenRouter returned error in /classify: {error_detail}")
            logger.error(f"Full response: {result}")
            raise HTTPException(status_code=500, detail=f"OpenRouter error: {error_detail}")

        reply = result.get("choices", [{}])[0].get("message", {}).get("content", "").strip()

        if not reply:
            logger.warning(f"Empty content in response. Finish reason: {result.get('choices', [{}])[0].get('finish_reason')}")
            # If content is empty, try to extract from reasoning or log for debugging
            reasoning = result.get("choices", [{}])[0].get("message", {}).get("reasoning", "")
            if reasoning:
                logger.warning(f"Model has reasoning but no content. This may indicate truncation.")
            logger.error(f"OpenRouter returned empty response in /classify. Full result: {result}")
            raise HTTPException(status_code=500, detail="No response from OpenRouter")

        logger.info(f"Classification successful with model {model_id}: {reply}")
        return {"reply": reply, "status": response.status_code, "model": model_id}
    except requests.RequestException as e:
        logger.error(f"OpenRouter network error in /classify: {e}")
        raise HTTPException(status_code=500, detail=f"Error contacting OpenRouter: {e}")
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Unexpected error in /classify: {e}")
        raise HTTPException(status_code=500, detail=f"Unexpected error: {e}")

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