Update app.py

app.py

@@ -2,7 +2,12 @@ from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 import requests
-import os  # for reading environment variables
+import os
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
@@ -23,12 +28,34 @@ if not OPENROUTER_API_KEY:
 # Request schema
 class ExplainRequest(BaseModel):
     message: str
-    label: str  # "Phishing" or "Safe"
+    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()
+    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")
@@ -42,8 +69,8 @@ def explain(req: ExplainRequest):
         "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}"
+        "Limit the number of bullets to 3-4. "
+        f"Message:\n\n{user_message}\n\n"
         "Respond using the same language as the message."
     )
 
@@ -61,12 +88,68 @@ def explain(req: ExplainRequest):
     }
 
     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)