Update app.py

app.py

@@ -1,11 +1,83 @@
 import spaces  # MUST be imported before torch
 
+import time
+import json
+import logging
+import uuid
+from datetime import datetime, timezone
+
 import torch
 import gradio as gr
+from fastapi import FastAPI
+from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 
 
+# ============================================================
+# Logging setup (monitoring)
+# ============================================================
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s | %(levelname)s | %(message)s",
+)
+logger = logging.getLogger("mediguide")
+
+REQUEST_LOG_FILE = "request_log.jsonl"  # note: ephemeral storage on free Spaces, resets on restart
+
+
+def log_request(request_id, message, response, latency_s, flagged, error=None):
+    entry = {
+        "request_id": request_id,
+        "timestamp": datetime.now(timezone.utc).isoformat(),
+        "message_len": len(message),
+        "response_len": len(response) if response else 0,
+        "latency_s": round(latency_s, 3),
+        "flagged": flagged,
+        "error": error,
+    }
+    logger.info(json.dumps(entry))
+    try:
+        with open(REQUEST_LOG_FILE, "a") as f:
+            f.write(json.dumps(entry) + "\n")
+    except Exception as e:
+        logger.warning(f"Could not write to request log file: {e}")
+
+
+# ============================================================
+# Guardrails
+# ============================================================
+
+EMERGENCY_KEYWORDS = [
+    "suicide", "kill myself", "want to die", "end my life",
+    "chest pain", "can't breathe", "cannot breathe", "heart attack",
+    "overdose", "severe bleeding", "unconscious", "stroke symptoms",
+]
+
+DISCLAIMER = (
+    "\n\n⚠️ This is not a substitute for professional medical advice. "
+    "If this is urgent, please contact a doctor or emergency services."
+)
+
+EMERGENCY_RESPONSE = (
+    "This sounds like it could be a medical emergency. Please contact "
+    "your local emergency services or go to the nearest emergency room "
+    "immediately. I'm not able to provide emergency medical care."
+)
+
+
+def check_emergency(message: str) -> bool:
+    msg_lower = message.lower()
+    return any(kw in msg_lower for kw in EMERGENCY_KEYWORDS)
+
+
+def apply_output_guardrails(response: str) -> str:
+    if "not a substitute" not in response.lower() and "professional medical" not in response.lower():
+        response = response.strip() + DISCLAIMER
+    return response
+
+
 # ============================================================
 # Configuration
 # ============================================================
@@ -24,71 +96,44 @@ appropriate professional medical care."""
 # Tokenizer & Model – loaded on CPU (no GPU yet)
 # ============================================================
 
-print("Loading tokenizer...")
+logger.info("Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
-print("Loading base model on CPU...")
+logger.info("Loading base model on CPU...")
 model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
     torch_dtype=torch.float16,
-    device_map="cpu",          # force CPU – GPU will be used later inside the decorated function
+    device_map="cpu",
 )
-print("Base model loaded on CPU.")
+logger.info("Base model loaded on CPU.")
 
-print("Loading MediGuide QLoRA adapter on CPU...")
+logger.info("Loading MediGuide QLoRA adapter on CPU...")
 model = PeftModel.from_pretrained(
     model,
     ADAPTER_MODEL,
-    torch_device="cpu",        # load adapter on CPU
+    torch_device="cpu",
 )
-model = model.merge_and_unload()   # merge for faster inference
+model = model.merge_and_unload()
 model.eval()
-print("MediGuide loaded successfully (on CPU).")
+logger.info("MediGuide loaded successfully (on CPU).")
 
 
 # ============================================================
-# Generation – GPU is requested via @spaces.GPU
+# Core generation (unchanged logic, guardrails + logging wrapped around it)
 # ============================================================
 
 @spaces.GPU(duration=60)
-def generate_response(message, history):
-    """
-    Generate a response from the model.
-    history: list of [user_msg, assistant_msg] tuples from Gradio.
-    """
-    # 1. Sanitise inputs
-    if not isinstance(message, str):
-        message = str(message) if message is not None else ""
-
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-
-    for turn in history:
-        if not isinstance(turn, (list, tuple)) or len(turn) < 2:
-            continue
-        user_msg, assistant_msg = turn[0], turn[1]
-        if user_msg is not None and str(user_msg).strip():
-            messages.append({"role": "user", "content": str(user_msg)})
-        if assistant_msg is not None and str(assistant_msg).strip():
-            messages.append({"role": "assistant", "content": str(assistant_msg)})
-
-    if message.strip():
-        messages.append({"role": "user", "content": message.strip()})
-    else:
-        return "Please enter a valid question."
-
-    # 2. Move model to GPU (now available thanks to @spaces.GPU)
+def _run_model(messages):
     model.to("cuda")
 
-    # 3. Apply chat template
     prompt = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True,
     )
 
-    # 4. Tokenize
     inputs = tokenizer(
         prompt,
         return_tensors="pt",
@@ -97,7 +142,6 @@ def generate_response(message, history):
     )
     inputs = {k: v.to("cuda") for k, v in inputs.items()}
 
-    # 5. Generate
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -109,17 +153,60 @@ def generate_response(message, history):
             pad_token_id=tokenizer.eos_token_id,
         )
 
-    # 6. Decode only new tokens
     generated_tokens = outputs[0][inputs["input_ids"].shape[-1]:]
     response = tokenizer.decode(generated_tokens, skip_special_tokens=True)
 
-    # 7. Move model back to CPU to free GPU memory for other requests
     model.to("cpu")
-    torch.cuda.empty_cache()   # optional, helps release memory
+    torch.cuda.empty_cache()
 
     return response.strip()
 
 
+def generate_response(message, history):
+    """
+    history: list of [user_msg, assistant_msg] tuples from Gradio.
+    Wraps model call with guardrails + logging.
+    """
+    request_id = str(uuid.uuid4())[:8]
+    start = time.time()
+
+    if not isinstance(message, str):
+        message = str(message) if message is not None else ""
+    message = message.strip()
+
+    if not message:
+        return "Please enter a valid question."
+
+    # Input guardrail: emergency detection short-circuits the model entirely
+    if check_emergency(message):
+        latency = time.time() - start
+        log_request(request_id, message, EMERGENCY_RESPONSE, latency, flagged="emergency_keyword")
+        return EMERGENCY_RESPONSE
+
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    for turn in history:
+        if not isinstance(turn, (list, tuple)) or len(turn) < 2:
+            continue
+        user_msg, assistant_msg = turn[0], turn[1]
+        if user_msg is not None and str(user_msg).strip():
+            messages.append({"role": "user", "content": str(user_msg)})
+        if assistant_msg is not None and str(assistant_msg).strip():
+            messages.append({"role": "assistant", "content": str(assistant_msg)})
+    messages.append({"role": "user", "content": message})
+
+    try:
+        response = _run_model(messages)
+        response = apply_output_guardrails(response)
+        latency = time.time() - start
+        log_request(request_id, message, response, latency, flagged=None)
+        return response
+    except Exception as e:
+        latency = time.time() - start
+        log_request(request_id, message, "", latency, flagged=None, error=str(e))
+        logger.error(f"[{request_id}] generation failed: {e}")
+        return "Sorry, something went wrong while generating a response. Please try again."
+
+
 # ============================================================
 # Gradio UI
 # ============================================================
@@ -160,8 +247,40 @@ with gr.Blocks(title="MediGuide") as demo:
     clear_btn.click(lambda: [], outputs=chatbot)
 
 
+# ============================================================
+# FastAPI wrapper — gives you a real REST endpoint alongside the UI
+# ============================================================
+
+app = FastAPI(title="MediGuide API")
+
+
+class ChatRequest(BaseModel):
+    message: str
+    history: list = []
+
+
+class ChatResponse(BaseModel):
+    response: str
+
+
+@app.get("/health")
+def health():
+    return {"status": "ok"}
+
+
+@app.post("/chat", response_model=ChatResponse)
+def chat_endpoint(req: ChatRequest):
+    response = generate_response(req.message, req.history)
+    return ChatResponse(response=response)
+
+
+app = gr.mount_gradio_app(app, demo.queue(), path="/")
+
+
 # ============================================================
 # Launch
 # ============================================================
 
-demo.queue().launch()
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)