Update app.py

app.py

@@ -1,11 +1,21 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from typing import List, Optional, Dict
-from llama_cpp import Llama
 import gradio as gr
 import json
 from enum import Enum
 import re
+import os
+import time
+from huggingface_hub import hf_hub_download
+
+# We'll import llama_cpp in a way that provides better error messages
+try:
+    from llama_cpp import Llama
+    LLAMA_IMPORT_ERROR = None
+except Exception as e:
+    LLAMA_IMPORT_ERROR = str(e)
+    print(f"Warning: Failed to import llama_cpp: {e}")
 
 class ConsultationState(Enum):
     INITIAL = "initial"
@@ -23,7 +33,7 @@ class ChatResponse(BaseModel):
     response: str
     finished: bool
 
-# Standard health assessment questions that Nurse Oge always asks
+# Standard health assessment questions
 HEALTH_ASSESSMENT_QUESTIONS = [
     "What are your current symptoms and how long have you been experiencing them?",
     "Do you have any pre-existing medical conditions or chronic illnesses?",
@@ -32,7 +42,6 @@ HEALTH_ASSESSMENT_QUESTIONS = [
     "Have you had any similar symptoms in the past? If yes, what treatments worked?"
 ]
 
-# Personality prompts for Nurse Oge
 NURSE_OGE_IDENTITY = """
 You are Nurse Oge, a medical AI assistant focused on serving patients in Nigeria. Always be empathetic, 
 professional, and thorough in your assessments. When asked about your identity, explain that you are 
@@ -42,14 +51,32 @@ health information before providing any medical advice.
 
 class NurseOgeAssistant:
     def __init__(self):
-        self.llm = Llama.from_pretrained(
-            repo_id="mradermacher/Llama3-Med42-8B-GGUF",
-            filename="Llama3-Med42-8B.IQ3_M.gguf",
-            verbose=False
-        )
-        self.consultation_states = {}  # Tracks state for each conversation
-        self.gathered_info = {}  # Stores gathered health information
+        if LLAMA_IMPORT_ERROR:
+            raise ImportError(f"Cannot initialize NurseOgeAssistant due to llama_cpp import error: {LLAMA_IMPORT_ERROR}")
+            
+        # Download the model file
+        try:
+            model_path = hf_hub_download(
+                repo_id="mradermacher/Llama3-Med42-8B-GGUF",
+                filename="Llama3-Med42-8B.IQ3_M.gguf",
+                resume_download=True
+            )
+            
+            # Initialize the model with the downloaded file
+            self.llm = Llama(
+                model_path=model_path,
+                n_ctx=2048,  # Context window
+                n_threads=4   # Number of CPU threads to use
+            )
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to initialize the model: {str(e)}")
+            
+        self.consultation_states = {}
+        self.gathered_info = {}
 
+    # ... (rest of the NurseOgeAssistant class methods remain the same)
+    
     def _is_identity_question(self, message: str) -> bool:
         identity_patterns = [
             r"who are you",
@@ -123,16 +150,14 @@ class NurseOgeAssistant:
                 )
             else:
                 self.consultation_states[conversation_id] = ConsultationState.DIAGNOSIS
-                # Prepare complete context for final response
                 context = "\n".join([
                     f"Q: {q}\nA: {a}" for q, a in 
                     zip(HEALTH_ASSESSMENT_QUESTIONS, self.gathered_info[conversation_id])
                 ])
                 
-                # Generate final response using the model
                 messages = [
                     {"role": "system", "content": NURSE_OGE_IDENTITY},
-                    {"role": "user", "content": f"Based on the following patient information, provide a thorough assessment, diagnosis and recommendations:\n\n{context}\n\nOriginal query: {message}"}
+                    {"role": "user", "content": f"Based on the following patient information, provide a thorough assessment and recommendations:\n\n{context}\n\nOriginal query: {message}"}
                 ]
                 
                 response = self.llm.create_chat_completion(
@@ -141,7 +166,6 @@ class NurseOgeAssistant:
                     temperature=0.7
                 )
                 
-                # Reset state for next consultation
                 self.consultation_states[conversation_id] = ConsultationState.INITIAL
                 self.gathered_info[conversation_id] = []
                 
@@ -150,22 +174,36 @@ class NurseOgeAssistant:
                     finished=True
                 )
 
-# Initialize FastAPI and Nurse Oge
+# Initialize FastAPI
 app = FastAPI()
-nurse_oge = NurseOgeAssistant()
+
+# Create a global variable for our assistant
+nurse_oge = None
+
+@app.on_event("startup")
+async def startup_event():
+    global nurse_oge
+    try:
+        nurse_oge = NurseOgeAssistant()
+    except Exception as e:
+        print(f"Failed to initialize NurseOgeAssistant: {e}")
+        # We'll continue running but the /chat endpoint will return errors
 
 @app.post("/chat")
 async def chat_endpoint(request: ChatRequest):
-    # Generate a conversation ID (in a real app, you'd want to manage these better)
+    if nurse_oge is None:
+        raise HTTPException(
+            status_code=503,
+            detail="The medical assistant is not available at the moment. Please try again later."
+        )
+    
     conversation_id = "default"
     
-    # Extract the latest message
     if not request.messages:
         raise HTTPException(status_code=400, detail="No messages provided")
     
     latest_message = request.messages[-1].content
     
-    # Process the message
     response = await nurse_oge.process_message(
         conversation_id=conversation_id,
         message=latest_message,
@@ -174,8 +212,11 @@ async def chat_endpoint(request: ChatRequest):
     
     return response
 
-# Initialize Gradio interface (optional, for testing)
+# Gradio interface
 def gradio_chat(message, history):
+    if nurse_oge is None:
+        return "The medical assistant is not available at the moment. Please try again later."
+    
     response = nurse_oge.process_message("gradio_user", message, history)
     return response.response