Enhance medical consultation app with LangChain memory management and improved patient context tracking

app.py

@@ -2,9 +2,13 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import spaces
+from langchain.memory import ConversationBufferWindowMemory
+from langchain.schema import HumanMessage, AIMessage
+import json
+from datetime import datetime
 
 # Model configuration - Using correct Me-LLaMA model identifier
-ME_LLAMA_MODEL = "clinicalnlplab/me-llama-13b"  # Corrected model name
+ME_LLAMA_MODEL = "clinicalnlplab/me-llama-13b"  
 
 # System prompts for different phases
 CONSULTATION_PROMPT = """You are a professional virtual doctor. Your goal is to collect detailed information about the user's health condition, symptoms, medical history, medications, lifestyle, and other relevant data.
@@ -26,7 +30,8 @@ MEDICINE_PROMPT = """You are a specialized medical assistant. Based on the patie
 
 Be concise, practical, and focus only on general symptom relief. Do not diagnose. Include a disclaimer that you are not a licensed medical professional.
 
-Patient information: {patient_info}"""
+Patient information: {patient_info}
+Previous conversation context: {memory_context}"""
 
 # Global variables
 me_llama_model = None
@@ -34,13 +39,121 @@ me_llama_tokenizer = None
 conversation_turns = 0
 patient_data = []
 
+# LangChain Memory Configuration
+class MedicalMemoryManager:
+    def __init__(self, k=10):  # Keep last 10 conversation turns
+        self.conversation_memory = ConversationBufferWindowMemory(k=k, return_messages=True)
+        self.patient_context = {
+            "symptoms": [],
+            "medical_history": [],
+            "medications": [],
+            "allergies": [],
+            "lifestyle_factors": [],
+            "timeline": [],
+            "severity_scores": {},
+            "session_start": datetime.now().isoformat()
+        }
+        
+    def add_interaction(self, human_input, ai_response):
+        """Add human-AI interaction to memory"""
+        self.conversation_memory.chat_memory.add_user_message(human_input)
+        self.conversation_memory.chat_memory.add_ai_message(ai_response)
+        
+        # Extract and categorize medical information
+        self._extract_medical_info(human_input)
+    
+    def _extract_medical_info(self, user_input):
+        """Extract medical information from user input and categorize it"""
+        user_lower = user_input.lower()
+        
+        # Extract symptoms (simple keyword matching)
+        symptom_keywords = ["pain", "ache", "hurt", "sore", "cough", "fever", "nausea", 
+                           "headache", "dizzy", "tired", "fatigue", "vomit", "swollen", 
+                           "rash", "itch", "burn", "cramp", "bleed", "shortness of breath"]
+        
+        for keyword in symptom_keywords:
+            if keyword in user_lower and keyword not in [s.lower() for s in self.patient_context["symptoms"]]:
+                self.patient_context["symptoms"].append(user_input)
+                break
+        
+        # Extract timeline information
+        time_keywords = ["days", "weeks", "months", "hours", "yesterday", "today", "started", "began"]
+        if any(keyword in user_lower for keyword in time_keywords):
+            self.patient_context["timeline"].append(user_input)
+        
+        # Extract severity (look for numbers 1-10)
+        import re
+        severity_match = re.search(r'\b([1-9]|10)\b.*(?:pain|severity|scale)', user_lower)
+        if severity_match:
+            self.patient_context["severity_scores"][datetime.now().isoformat()] = severity_match.group(1)
+        
+        # Extract medications
+        med_keywords = ["taking", "medication", "medicine", "pills", "prescribed", "drug"]
+        if any(keyword in user_lower for keyword in med_keywords):
+            self.patient_context["medications"].append(user_input)
+        
+        # Extract allergies
+        allergy_keywords = ["allergic", "allergy", "allergies", "reaction"]
+        if any(keyword in user_lower for keyword in allergy_keywords):
+            self.patient_context["allergies"].append(user_input)
+    
+    def get_memory_context(self):
+        """Get formatted memory context for the model"""
+        messages = self.conversation_memory.chat_memory.messages
+        context = []
+        
+        for msg in messages[-6:]:  # Last 6 messages (3 exchanges)
+            if isinstance(msg, HumanMessage):
+                context.append(f"Patient: {msg.content}")
+            elif isinstance(msg, AIMessage):
+                context.append(f"Doctor: {msg.content}")
+        
+        return "\n".join(context)
+    
+    def get_patient_summary(self):
+        """Get structured patient information summary"""
+        summary = {
+            "conversation_turns": len(self.conversation_memory.chat_memory.messages) // 2,
+            "session_duration": datetime.now().isoformat(),
+            "key_symptoms": self.patient_context["symptoms"][-3:],  # Last 3 symptoms mentioned
+            "timeline_info": self.patient_context["timeline"][-2:],  # Last 2 timeline mentions
+            "medications": self.patient_context["medications"],
+            "allergies": self.patient_context["allergies"],
+            "severity_scores": self.patient_context["severity_scores"]
+        }
+        return json.dumps(summary, indent=2)
+    
+    def reset_session(self):
+        """Reset memory for new consultation"""
+        self.conversation_memory.clear()
+        self.patient_context = {
+            "symptoms": [],
+            "medical_history": [],
+            "medications": [],
+            "allergies": [],
+            "lifestyle_factors": [],
+            "timeline": [],
+            "severity_scores": {},
+            "session_start": datetime.now().isoformat()
+        }
+
+# Initialize memory manager
+memory_manager = MedicalMemoryManager()
+
 def build_me_llama_prompt(system_prompt, history, user_input):
-    """Format the conversation for Me-LLaMA chat model."""
+    """Format the conversation for Me-LLaMA chat model with memory context."""
+    # Get memory context from LangChain
+    memory_context = memory_manager.get_memory_context()
+    
+    # Enhance system prompt with memory context
+    enhanced_system_prompt = f"{system_prompt}\n\nPrevious conversation context:\n{memory_context}"
+    
     # Use standard Llama-2 chat format since Me-LLaMA is based on Llama-2
-    prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"
+    prompt = f"<s>[INST] <<SYS>>\n{enhanced_system_prompt}\n<</SYS>>\n\n"
     
-    # Add conversation history
-    for user_msg, assistant_msg in history:
+    # Add only recent history to avoid token limit issues
+    recent_history = history[-3:] if len(history) > 3 else history
+    for user_msg, assistant_msg in recent_history:
         prompt += f"{user_msg} [/INST] {assistant_msg} </s><s>[INST] "
     
     # Add the current user input
@@ -80,12 +193,12 @@ def load_model_if_needed():
             print("Fallback model loaded successfully!")
 
 @spaces.GPU
-def generate_medicine_suggestions(patient_info):
-    """Use Me-LLaMA to generate medicine and remedy suggestions."""
+def generate_medicine_suggestions(patient_info, memory_context):
+    """Use Me-LLaMA to generate medicine and remedy suggestions with memory context."""
     load_model_if_needed()
     
-    # Create a simple prompt for medicine suggestions
-    prompt = f"<s>[INST] {MEDICINE_PROMPT.format(patient_info=patient_info)} [/INST] "
+    # Create a prompt with both patient info and memory context
+    prompt = f"<s>[INST] {MEDICINE_PROMPT.format(patient_info=patient_info, memory_context=memory_context)} [/INST] "
     
     inputs = me_llama_tokenizer(prompt, return_tensors="pt")
     
@@ -109,7 +222,7 @@ def generate_medicine_suggestions(patient_info):
 
 @spaces.GPU
 def generate_response(message, history):
-    """Generate response using only Me-LLaMA for both consultation and medicine suggestions."""
+    """Generate response using Me-LLaMA with LangChain memory management."""
     global conversation_turns, patient_data
     
     try:
@@ -119,12 +232,12 @@ def generate_response(message, history):
         # Track conversation turns
         conversation_turns += 1
         
-        # Store patient data
+        # Store patient data (legacy support)
         patient_data.append(message)
         
-        # Phase 1-3: Information gathering
+        # Phase 1-3: Information gathering with memory
         if conversation_turns < 4:
-            # Build consultation prompt
+            # Build consultation prompt with memory context
             prompt = build_me_llama_prompt(CONSULTATION_PROMPT, history, message)
             
             inputs = me_llama_tokenizer(prompt, return_tensors="pt")
@@ -149,13 +262,20 @@ def generate_response(message, history):
             full_response = me_llama_tokenizer.decode(outputs[0], skip_special_tokens=False)
             response = full_response.split('[/INST]')[-1].split('</s>')[0].strip()
             
+            # Add interaction to memory
+            memory_manager.add_interaction(message, response)
+            
             return response
         
-        # Phase 4+: Summary and medicine suggestions
+        # Phase 4+: Summary and medicine suggestions with full memory context
         else:
-            # First, get summary from consultation
+            # Get comprehensive patient summary from memory
+            patient_summary = memory_manager.get_patient_summary()
+            memory_context = memory_manager.get_memory_context()
+            
+            # First, get summary from consultation with memory context
             summary_prompt = build_me_llama_prompt(
-                CONSULTATION_PROMPT + "\n\nNow summarize what you've learned and suggest when professional care may be needed.",
+                CONSULTATION_PROMPT + "\n\nNow provide a comprehensive summary based on all the information gathered. Include when professional care may be needed.",
                 history, 
                 message
             )
@@ -180,34 +300,81 @@ def generate_response(message, history):
             summary_response = me_llama_tokenizer.decode(outputs[0], skip_special_tokens=False)
             summary = summary_response.split('[/INST]')[-1].split('</s>')[0].strip()
             
-            # Then get medicine suggestions using the same model
-            full_patient_info = "\n".join(patient_data) + f"\n\nMedical Summary: {summary}"
-            medicine_suggestions = generate_medicine_suggestions(full_patient_info)
+            # Get medicine suggestions using memory context
+            full_patient_info = f"Patient Summary: {patient_summary}\n\nDetailed Summary: {summary}"
+            medicine_suggestions = generate_medicine_suggestions(full_patient_info, memory_context)
             
             # Combine both responses
             final_response = (
-                f"**MEDICAL SUMMARY:**\n{summary}\n\n"
+                f"**COMPREHENSIVE MEDICAL SUMMARY:**\n{summary}\n\n"
                 f"**MEDICATION AND HOME CARE SUGGESTIONS:**\n{medicine_suggestions}\n\n"
+                f"**PATIENT CONTEXT SUMMARY:**\n{patient_summary}\n\n"
                 f"**DISCLAIMER:** This is AI-generated advice for informational purposes only. Please consult a licensed healthcare provider for proper medical diagnosis and treatment."
             )
             
+            # Add final interaction to memory
+            memory_manager.add_interaction(message, final_response)
+            
             return final_response
             
     except Exception as e:
-        return f"I apologize, but I'm experiencing technical difficulties. Please try again. Error: {str(e)}"
+        error_msg = f"I apologize, but I'm experiencing technical difficulties. Please try again. Error: {str(e)}"
+        # Still try to add to memory even on error
+        try:
+            memory_manager.add_interaction(message, error_msg)
+        except:
+            pass
+        return error_msg
+
+def reset_consultation():
+    """Reset the consultation and memory for a new patient."""
+    global conversation_turns, patient_data, memory_manager
+    
+    conversation_turns = 0
+    patient_data = []
+    memory_manager.reset_session()
+    
+    return "New consultation started. Please tell me about your symptoms or health concerns."
 
-# Create the Gradio interface
-demo = gr.ChatInterface(
-    fn=generate_response,
-    title="🏥 Complete Medical Assistant - Me-LLaMA 13B",
-    description="Comprehensive medical consultation powered by Me-LLaMA 13B. One model handles both consultation and medicine suggestions. Tell me about your symptoms!",
-    examples=[
-        "I have a persistent cough and sore throat for 3 days",
-        "I've been having severe headaches and feel dizzy",
-        "My stomach hurts and I feel nauseous after eating"
-    ],
-    theme="soft"
-)
+# Create the Gradio interface with memory reset option
+with gr.Blocks(theme="soft") as demo:
+    gr.Markdown("# 🏥 Complete Medical Assistant - Me-LLaMA 13B with Memory")
+    gr.Markdown("Comprehensive medical consultation powered by Me-LLaMA 13B with LangChain memory management. One model handles both consultation and medicine suggestions with full context awareness.")
+    
+    with gr.Row():
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot(height=500)
+            msg = gr.Textbox(
+                placeholder="Tell me about your symptoms or health concerns...",
+                label="Your Message"
+            )
+            
+        with gr.Column(scale=1):
+            reset_btn = gr.Button("🔄 Start New Consultation", variant="secondary")
+            gr.Markdown("**Memory Features:**\n- Tracks symptoms & timeline\n- Remembers medications & allergies\n- Maintains conversation context\n- Provides comprehensive summaries")
+    
+    # Examples
+    gr.Examples(
+        examples=[
+            "I have a persistent cough and sore throat for 3 days",
+            "I've been having severe headaches and feel dizzy",
+            "My stomach hurts and I feel nauseous after eating"
+        ],
+        inputs=msg
+    )
+    
+    # Event handlers
+    def respond(message, chat_history):
+        bot_message = generate_response(message, chat_history)
+        chat_history.append((message, bot_message))
+        return "", chat_history
+    
+    def reset_chat():
+        reset_msg = reset_consultation()
+        return [(None, reset_msg)], ""
+    
+    msg.submit(respond, [msg, chatbot], [msg, chatbot])
+    reset_btn.click(reset_chat, [], [chatbot, msg])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()