SurLink/Explainer
0
1from fastapi import FastAPI, HTTPException2from fastapi.middleware.cors import CORSMiddleware3from pydantic import BaseModel4import requests5import os6import logging7import threading8import time9 10# Configure logging11logging.basicConfig(level=logging.INFO)12logger = logging.getLogger(__name__)13 14app = FastAPI()15 16# Enable CORS17app.add_middleware(18 CORSMiddleware,19 allow_origins=["*"], # Allow all origins, or replace with your frontend URL20 allow_credentials=True,21 allow_methods=["*"],22 allow_headers=["*"],23)24 25# Read OpenRouter API key from environment variable (HF Secret)26OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")27if not OPENROUTER_API_KEY:28 raise RuntimeError("OPENROUTER_API_KEY environment variable not set!")29 30EXPLAIN_MODEL = os.environ.get("EXPLAIN_MODEL")31if not EXPLAIN_MODEL:32 raise RuntimeError("EXPLAIN_MODEL environment variable not set!")33 34# ---------------------------------------------------------------------------35# Global rate limiter for OpenRouter calls36#37# Free-tier OpenRouter models throttle/reject requests that arrive too close38# together. Since /classify and /explain can be hit concurrently by the39# frontend (each running in its own thread pool worker), we serialize ALL40# outbound OpenRouter calls behind a single lock and enforce a minimum gap41# between them. Whichever request arrives second simply waits its turn.42# ---------------------------------------------------------------------------43MIN_INTERVAL_SECONDS = float(os.environ.get("OPENROUTER_MIN_INTERVAL", "6.0"))44MAX_RETRIES = int(os.environ.get("OPENROUTER_MAX_RETRIES", "3"))45RETRY_BACKOFF_SECONDS = float(os.environ.get("OPENROUTER_RETRY_BACKOFF", "5.0"))46 47_or_lock = threading.Lock()48_last_call_time = 0.049 50def _is_rate_limited(response: requests.Response) -> bool:51 """Detect a 429 whether it comes as an HTTP status or embedded in a 200 body52 (OpenRouter sometimes returns 200 with {"error": {"code": 429, ...}})."""53 if response.status_code == 429:54 return True55 try:56 body = response.json()57 except ValueError:58 return False59 return isinstance(body, dict) and body.get("error", {}).get("code") == 42960 61 62def call_openrouter(payload: dict, headers: dict, timeout: int = 20):63 """Thread-safe, rate-limited call to OpenRouter's chat completions endpoint.64 65 - Ensures at least MIN_INTERVAL_SECONDS has elapsed since the previous66 OpenRouter call (across ALL endpoints) before firing this one.67 - If the provider still returns a 429 (as a status code or embedded in the68 JSON body), retries with exponential backoff up to MAX_RETRIES times.69 """70 global _last_call_time71 72 url = "https://openrouter.ai/api/v1/chat/completions"73 74 with _or_lock:75 now = time.monotonic()76 elapsed = now - _last_call_time77 wait_for = MIN_INTERVAL_SECONDS - elapsed78 if wait_for > 0:79 logger.info(f"Rate limiting: waiting {wait_for:.2f}s before next OpenRouter call")80 time.sleep(wait_for)81 82 response = requests.post(url, headers=headers, json=payload, timeout=timeout)83 _last_call_time = time.monotonic()84 85 attempt = 086 while _is_rate_limited(response) and attempt < MAX_RETRIES:87 backoff = RETRY_BACKOFF_SECONDS * (2 ** attempt)88 logger.warning(f"OpenRouter rate limited (attempt {attempt + 1}/{MAX_RETRIES}); "89 f"backing off {backoff:.1f}s before retry")90 time.sleep(backoff)91 response = requests.post(url, headers=headers, json=payload, timeout=timeout)92 _last_call_time = time.monotonic()93 attempt += 194 95 return response96 97 98# Request schema99class ExplainRequest(BaseModel):100 message: str101 label: str = None # Optional, for explanations102 model_id: str = EXPLAIN_MODEL # Default model, can be overridden103 104@app.get("/")105def root():106 return {107 "status": "healthy",108 "service": "Anti-Phishing Explainer",109 "endpoints": {110 "/explain": "POST - Generate explanation for classification",111 "/classify": "POST - Classify text as Phishing or Safe",112 "/health": "GET - Health check"113 }114 }115 116@app.get("/health")117def health():118 return {119 "status": "healthy",120 "service": "explainer",121 "openrouter_configured": bool(OPENROUTER_API_KEY)122 }123 124@app.post("/explain")125def explain(req: ExplainRequest):126 """Generate a human-readable explanation for why a message was classified as Phishing or Safe"""127 user_message = req.message.strip()128 label = req.label.strip() if req.label else None129 130 if not user_message or not label:131 raise HTTPException(status_code=400, detail="Missing message or label")132 133 # Updated system prompt with bullet-point format and language adaptation134 system_prompt = (135 "Respond using the same language as the message."136 f"You are a robot that identifies phishing and safe messages. "137 f"The message was classified as '{label}'. "138 "Explain why this decision was made and point out any words or patterns that led to it. "139 "No greetings, introductions, or closing remarks. "140 "Don't restate the message or its classification. "141 "Output only the explanation as bullet points. "142 "Limit each bullet to 1–2 sentences. "143 "Limit the number of bullets to 3-4. "144 f"Message:\n\n{user_message}\n\n"145 )146 147 headers = {148 "Authorization": f"Bearer {OPENROUTER_API_KEY}",149 "Content-Type": "application/json"150 }151 payload = {152 "model": req.model_id,153 "messages": [154 {"role": "system", "content": system_prompt},155 {"role": "user", "content": user_message}156 ]157 }158 159 try:160 logger.info(f"Calling OpenRouter /explain with model: {req.model_id}")161 response = call_openrouter(payload, headers)162 response.raise_for_status()163 result = response.json()164 165 # Check for OpenRouter error in response body (even with 200 status)166 if "error" in result:167 error_detail = result.get("error", {}).get("message", str(result.get("error")))168 logger.error(f"OpenRouter returned error in /explain: {error_detail}")169 logger.error(f"Full response: {result}")170 raise HTTPException(status_code=500, detail=f"OpenRouter error: {error_detail}")171 172 reply = result.get("choices", [{}])[0].get("message", {}).get("content", "").strip()173 if not reply:174 logger.error(f"OpenRouter returned empty response in /explain. Full result: {result}")175 reply = "[No explanation returned]"176 logger.info("Explanation generated successfully")177 return {"reply": reply}178 except requests.RequestException as e:179 logger.error(f"OpenRouter network error in /explain: {e}")180 raise HTTPException(status_code=500, detail=f"Error contacting OpenRouter: {e}")181 except HTTPException:182 raise183 except Exception as e:184 logger.error(f"Unexpected error in /explain: {e}")185 raise HTTPException(status_code=500, detail=f"Unexpected error: {e}")186 187@app.post("/classify")188def classify(req: ExplainRequest):189 """Classify text as Phishing or Safe via OpenRouter"""190 user_message = req.message.strip()191 if not user_message:192 raise HTTPException(status_code=400, detail="Missing message")193 194 # Use provided model_id or default195 model_id = req.model_id or "arcee-ai/trinity-large-preview:free"196 197 system_prompt = (198 'You are a phishing detector. Classify the text as "Phishing" or "Safe". '199 'Respond ONLY with valid JSON: {"label": "Phishing"|"Safe", "confidence": <0-100 float>}. '200 'No other text.'201 )202 203 headers = {204 "Authorization": f"Bearer {OPENROUTER_API_KEY}",205 "Content-Type": "application/json"206 }207 payload = {208 "model": model_id, # Use the passed model_id209 "messages": [210 {"role": "system", "content": system_prompt},211 {"role": "user", "content": user_message}212 ],213 "temperature": 0.0,214 "max_tokens": 2000 # Increased for models with extended reasoning/thinking215 }216 217 try:218 logger.info(f"Calling OpenRouter /classify with model: {model_id}")219 response = call_openrouter(payload, headers)220 response.raise_for_status()221 result = response.json()222 223 # Check for OpenRouter error in response body (even with 200 status)224 if "error" in result:225 error_detail = result.get("error", {}).get("message", str(result.get("error")))226 logger.error(f"OpenRouter returned error in /classify: {error_detail}")227 logger.error(f"Full response: {result}")228 raise HTTPException(status_code=500, detail=f"OpenRouter error: {error_detail}")229 230 reply = result.get("choices", [{}])[0].get("message", {}).get("content", "").strip()231 232 if not reply:233 logger.warning(f"Empty content in response. Finish reason: {result.get('choices', [{}])[0].get('finish_reason')}")234 # If content is empty, try to extract from reasoning or log for debugging235 reasoning = result.get("choices", [{}])[0].get("message", {}).get("reasoning", "")236 if reasoning:237 logger.warning(f"Model has reasoning but no content. This may indicate truncation.")238 logger.error(f"OpenRouter returned empty response in /classify. Full result: {result}")239 raise HTTPException(status_code=500, detail="No response from OpenRouter")240 241 logger.info(f"Classification successful with model {model_id}: {reply}")242 return {"reply": reply, "status": response.status_code, "model": model_id}243 except requests.RequestException as e:244 logger.error(f"OpenRouter network error in /classify: {e}")245 raise HTTPException(status_code=500, detail=f"Error contacting OpenRouter: {e}")246 except HTTPException:247 raise248 except Exception as e:249 logger.error(f"Unexpected error in /classify: {e}")250 raise HTTPException(status_code=500, detail=f"Unexpected error: {e}")251 252# Optional: Run with uvicorn253if __name__ == "__main__":254 import uvicorn255 uvicorn.run(app, host="0.0.0.0", port=7860)