Update app.py

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel
 from typing import List, Optional, Dict
 import gradio as gr
@@ -7,9 +7,15 @@ from enum import Enum
 import re
 import os
 import time
+import gc
 from huggingface_hub import hf_hub_download
 
-# We'll import llama_cpp in a way that provides better error messages
+# Environment variables for configuration
+MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "mradermacher/Llama3-Med42-8B-GGUF")
+MODEL_FILENAME = os.getenv("MODEL_FILENAME", "Llama3-Med42-8B.Q4_K_M.gguf")
+N_THREADS = int(os.getenv("N_THREADS", "4"))
+
+# Import llama_cpp with error handling for better debugging
 try:
     from llama_cpp import Llama
     LLAMA_IMPORT_ERROR = None
@@ -33,7 +39,7 @@ class ChatResponse(BaseModel):
     response: str
     finished: bool
 
-# Standard health assessment questions
+# Standard health assessment questions for thorough patient evaluation
 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?",
@@ -42,6 +48,7 @@ HEALTH_ASSESSMENT_QUESTIONS = [
     "Have you had any similar symptoms in the past? If yes, what treatments worked?"
 ]
 
+# Define the AI assistant's identity and role
 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 
@@ -54,19 +61,14 @@ class NurseOgeAssistant:
         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
+            # Initialize the model using from_pretrained for better compatibility with free tier
+            self.llm = Llama.from_pretrained(
+                repo_id=MODEL_REPO_ID,
+                filename=MODEL_FILENAME,
+                n_ctx=2048,      # Context window size
+                n_threads=N_THREADS,  # Adjust based on available CPU resources
+                n_gpu_layers=0   # CPU-only inference for free tier
             )
             
         except Exception as e:
@@ -75,8 +77,6 @@ class NurseOgeAssistant:
         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",
@@ -107,72 +107,94 @@ class NurseOgeAssistant:
         return None
 
     async def process_message(self, conversation_id: str, message: str, history: List[Dict]) -> ChatResponse:
-        # Initialize state if new conversation
-        if conversation_id not in self.consultation_states:
-            self.consultation_states[conversation_id] = ConsultationState.INITIAL
-
-        # Handle identity questions
-        if self._is_identity_question(message):
-            return ChatResponse(
-                response="I am Nurse Oge, a medical AI assistant dedicated to helping patients in Nigeria. "
-                        "I'm here to provide medical guidance while ensuring I gather all necessary health information "
-                        "for accurate assessments.",
-                finished=True
-            )
+        try:
+            # Initialize state for new conversations
+            if conversation_id not in self.consultation_states:
+                self.consultation_states[conversation_id] = ConsultationState.INITIAL
 
-        # Handle location questions
-        if self._is_location_question(message):
-            return ChatResponse(
-                response="I am based in Nigeria and specifically trained to serve Nigerian communities, "
-                        "taking into account local healthcare contexts and needs.",
-                finished=True
-            )
+            # Handle identity questions
+            if self._is_identity_question(message):
+                return ChatResponse(
+                    response="I am Nurse Oge, a medical AI assistant dedicated to helping patients in Nigeria. "
+                            "I'm here to provide medical guidance while ensuring I gather all necessary health information "
+                            "for accurate assessments.",
+                    finished=True
+                )
 
-        # Start health assessment if it's a medical query
-        if self.consultation_states[conversation_id] == ConsultationState.INITIAL:
-            self.consultation_states[conversation_id] = ConsultationState.GATHERING_INFO
-            next_question = self._get_next_assessment_question(conversation_id)
-            return ChatResponse(
-                response=f"Before I can provide any medical advice, I need to gather some important health information. "
-                        f"{next_question}",
-                finished=False
-            )
+            # Handle location questions
+            if self._is_location_question(message):
+                return ChatResponse(
+                    response="I am based in Nigeria and specifically trained to serve Nigerian communities, "
+                            "taking into account local healthcare contexts and needs.",
+                    finished=True
+                )
 
-        # Continue gathering information
-        if self.consultation_states[conversation_id] == ConsultationState.GATHERING_INFO:
-            self.gathered_info[conversation_id].append(message)
-            next_question = self._get_next_assessment_question(conversation_id)
-            
-            if next_question:
+            # Start health assessment for medical queries
+            if self.consultation_states[conversation_id] == ConsultationState.INITIAL:
+                self.consultation_states[conversation_id] = ConsultationState.GATHERING_INFO
+                next_question = self._get_next_assessment_question(conversation_id)
                 return ChatResponse(
-                    response=f"Thank you for that information. {next_question}",
+                    response=f"Before I can provide any medical advice, I need to gather some important health information. "
+                            f"{next_question}",
                     finished=False
                 )
-            else:
-                self.consultation_states[conversation_id] = ConsultationState.DIAGNOSIS
-                context = "\n".join([
-                    f"Q: {q}\nA: {a}" for q, a in 
-                    zip(HEALTH_ASSESSMENT_QUESTIONS, self.gathered_info[conversation_id])
-                ])
-                
-                messages = [
-                    {"role": "system", "content": NURSE_OGE_IDENTITY},
-                    {"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(
-                    messages=messages,
-                    max_tokens=1024,
-                    temperature=0.7
-                )
-                
-                self.consultation_states[conversation_id] = ConsultationState.INITIAL
-                self.gathered_info[conversation_id] = []
+
+            # Continue gathering information
+            if self.consultation_states[conversation_id] == ConsultationState.GATHERING_INFO:
+                self.gathered_info[conversation_id].append(message)
+                next_question = self._get_next_assessment_question(conversation_id)
                 
-                return ChatResponse(
-                    response=response['choices'][0]['message']['content'],
-                    finished=True
-                )
+                if next_question:
+                    return ChatResponse(
+                        response=f"Thank you for that information. {next_question}",
+                        finished=False
+                    )
+                else:
+                    self.consultation_states[conversation_id] = ConsultationState.DIAGNOSIS
+                    context = "\n".join([
+                        f"Q: {q}\nA: {a}" for q, a in 
+                        zip(HEALTH_ASSESSMENT_QUESTIONS, self.gathered_info[conversation_id])
+                    ])
+                    
+                    messages = [
+                        {"role": "system", "content": NURSE_OGE_IDENTITY},
+                        {"role": "user", "content": f"Based on the following patient information, provide a thorough assessment and recommendations:\n\n{context}\n\nOriginal query: {message}"}
+                    ]
+                    
+                    # Implement retry logic for API calls
+                    max_retries = 3
+                    retry_delay = 2
+                    
+                    for attempt in range(max_retries):
+                        try:
+                            response = self.llm.create_chat_completion(
+                                messages=messages,
+                                max_tokens=512,  # Reduced for free tier
+                                temperature=0.7
+                            )
+                            break
+                        except Exception as e:
+                            if attempt < max_retries - 1:
+                                time.sleep(retry_delay)
+                                continue
+                            return ChatResponse(
+                                response="I'm sorry, I'm experiencing some technical difficulties. Please try again in a moment.",
+                                finished=True
+                            )
+                    
+                    self.consultation_states[conversation_id] = ConsultationState.INITIAL
+                    self.gathered_info[conversation_id] = []
+                    
+                    return ChatResponse(
+                        response=response['choices'][0]['message']['content'],
+                        finished=True
+                    )
+                    
+        except Exception as e:
+            return ChatResponse(
+                response=f"An error occurred while processing your request. Please try again.",
+                finished=True
+            )
 
 # Initialize FastAPI
 app = FastAPI()
@@ -180,6 +202,14 @@ app = FastAPI()
 # Create a global variable for our assistant
 nurse_oge = None
 
+# Add memory management middleware
+@app.middleware("http")
+async def add_memory_management(request: Request, call_next):
+    gc.collect()  # Force garbage collection before processing request
+    response = await call_next(request)
+    gc.collect()  # Clean up after request
+    return response
+
 @app.on_event("startup")
 async def startup_event():
     global nurse_oge
@@ -187,7 +217,10 @@ async def startup_event():
         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.get("/health")
+async def health_check():
+    return {"status": "healthy", "model_loaded": nurse_oge is not None}
 
 @app.post("/chat")
 async def chat_endpoint(request: ChatRequest):
@@ -197,15 +230,13 @@ async def chat_endpoint(request: ChatRequest):
             detail="The medical assistant is not available at the moment. Please try again later."
         )
     
-    conversation_id = "default"
-    
     if not request.messages:
         raise HTTPException(status_code=400, detail="No messages provided")
     
     latest_message = request.messages[-1].content
     
     response = await nurse_oge.process_message(
-        conversation_id=conversation_id,
+        conversation_id="default",
         message=latest_message,
         history=request.messages[:-1]
     )
@@ -220,6 +251,7 @@ def gradio_chat(message, history):
     response = nurse_oge.process_message("gradio_user", message, history)
     return response.response
 
+# Create and configure Gradio interface
 demo = gr.ChatInterface(
     fn=gradio_chat,
     title="Nurse Oge",