app.py

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

# 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!")

# Request schema
class ExplainRequest(BaseModel):
    message: str
    label: str = None  # Optional, for explanations
    model_id: str = "arcee-ai/trinity-large-preview: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 = (
        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"
        "Respond using the same language as the message."
    )

    url = "https://openrouter.ai/api/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": "arcee-ai/trinity-large-preview:free",
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message}
        ]
    }

    try:
        logger.info(f"Calling OpenRouter /explain with model: arcee-ai/trinity-large-preview:free")
        response = requests.post(url, headers=headers, json=payload, timeout=20)
        response.raise_for_status()
        result = response.json()
        reply = result.get("choices", [{}])[0].get("message", {}).get("content", "").strip()
        if not reply:
            reply = "[No explanation returned]"
        logger.info("Explanation generated successfully")
        return {"reply": reply}
    except requests.RequestException as e:
        logger.error(f"OpenRouter error in /explain: {e}")
        raise HTTPException(status_code=500, detail=f"Error contacting OpenRouter: {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.'
    )

    url = "https://openrouter.ai/api/v1/chat/completions"
    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": 256
    }

    try:
        logger.info(f"Calling OpenRouter /classify with model: {model_id}")
        response = requests.post(url, headers=headers, json=payload, timeout=20)
        response.raise_for_status()
        result = response.json()
        reply = result.get("choices", [{}])[0].get("message", {}).get("content", "").strip()
        
        if not reply:
            logger.error("OpenRouter returned empty response")
            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 error in /classify: {e}")
        raise HTTPException(status_code=500, detail=f"Error contacting OpenRouter: {e}")

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