update meditron

app.py

@@ -1,530 +1,255 @@
 import gradio as gr
+import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from langgraph.graph import StateGraph, END
-from typing import TypedDict, List, Dict
-from datetime import datetime
-import re
-
-# Enhanced State Management with RAG
-class MedicalState(TypedDict):
-    patient_id: str
-    conversation_history: List[Dict]
-    symptoms: Dict[str, any]
-    vital_questions_asked: List[str]
-    medical_history: Dict
-    current_medications: List[str]
-    allergies: List[str]
-    severity_scores: Dict[str, int]
-    red_flags: List[str]
-    assessment_complete: bool
-    suggested_actions: List[str]
-    consultation_stage: str
-    retrieved_knowledge: List[Dict]
-    confidence_scores: Dict[str, float]
-
-# Medical Knowledge Base for RAG
-MEDICAL_KNOWLEDGE_BASE = {
-    "conditions": {
-        "common_cold": {
-            "symptoms": ["runny nose", "cough", "sneezing", "sore throat", "mild fever"],
-            "treatment": "Rest, fluids, OTC pain relievers",
-            "otc_medications": [
-                {"name": "Acetaminophen", "dose": "500-1000mg every 4-6 hours", "max_daily": "3000mg"},
-                {"name": "Ibuprofen", "dose": "200-400mg every 4-6 hours", "max_daily": "1200mg"}
-            ],
-            "home_remedies": ["Warm salt water gargle", "Honey and lemon tea", "Steam inhalation"],
-            "when_to_seek_care": "If symptoms worsen after 7-10 days or fever above 101.3°F"
-        },
-        "headache": {
-            "symptoms": ["head pain", "pressure", "throbbing"],
-            "treatment": "Pain relief, rest, hydration",
-            "otc_medications": [
-                {"name": "Acetaminophen", "dose": "500-1000mg every 4-6 hours", "max_daily": "3000mg"},
-                {"name": "Ibuprofen", "dose": "400-600mg every 6-8 hours", "max_daily": "1200mg"}
-            ],
-            "home_remedies": ["Cold or warm compress", "Rest in dark room", "Stay hydrated"],
-            "when_to_seek_care": "Sudden severe headache, fever, neck stiffness, vision changes"
-        },
-        "stomach_pain": {
-            "symptoms": ["abdominal pain", "nausea", "bloating", "cramps"],
-            "treatment": "Bland diet, rest, hydration",
-            "otc_medications": [
-                {"name": "Pepto-Bismol", "dose": "525mg every 30 minutes as needed", "max_daily": "8 doses"},
-                {"name": "TUMS", "dose": "2-4 tablets as needed", "max_daily": "15 tablets"}
-            ],
-            "home_remedies": ["BRAT diet", "Ginger tea", "Warm compress on stomach"],
-            "when_to_seek_care": "Severe pain, fever, vomiting, blood in stool"
-        }
-    }
-}
-
-MEDICAL_CATEGORIES = {
-    "respiratory": ["cough", "shortness of breath", "chest pain", "wheezing", "runny nose", "sore throat"],
-    "gastrointestinal": ["nausea", "vomiting", "diarrhea", "stomach pain", "heartburn", "bloating"],
-    "neurological": ["headache", "dizziness", "numbness", "tingling"],
-    "musculoskeletal": ["joint pain", "muscle pain", "back pain", "stiffness"],
-    "cardiovascular": ["chest pain", "palpitations", "swelling", "fatigue"],
-    "dermatological": ["rash", "itching", "skin changes", "wounds"],
-    "mental_health": ["anxiety", "depression", "sleep issues", "stress"]
-}
-
-RED_FLAGS = [
-    "chest pain", "difficulty breathing", "severe headache", "high fever",
-    "blood in stool", "blood in urine", "severe abdominal pain",
-    "sudden vision changes", "loss of consciousness", "severe allergic reaction"
-]
-
-class SimpleRAGSystem:
-    def __init__(self):
-        self.knowledge_base = MEDICAL_KNOWLEDGE_BASE
-        self.setup_simple_retrieval()
-    
-    def setup_simple_retrieval(self):
-        """Setup simple keyword-based retrieval system"""
-        self.symptom_to_condition = {}
-        
-        for condition, data in self.knowledge_base["conditions"].items():
-            for symptom in data["symptoms"]:
-                if symptom not in self.symptom_to_condition:
-                    self.symptom_to_condition[symptom] = []
-                self.symptom_to_condition[symptom].append(condition)
-    
-    def retrieve_relevant_conditions(self, symptoms: List[str]) -> List[Dict]:
-        """Retrieve relevant medical conditions based on symptoms"""
-        relevant_conditions = {}
-        
-        for symptom in symptoms:
-            symptom_lower = symptom.lower()
-            
-            # Direct match
-            if symptom_lower in self.symptom_to_condition:
-                for condition in self.symptom_to_condition[symptom_lower]:
-                    if condition not in relevant_conditions:
-                        relevant_conditions[condition] = self.knowledge_base["conditions"][condition]
-            
-            # Partial match
-            for kb_symptom, conditions in self.symptom_to_condition.items():
-                if symptom_lower in kb_symptom or kb_symptom in symptom_lower:
-                    for condition in conditions:
-                        if condition not in relevant_conditions:
-                            relevant_conditions[condition] = self.knowledge_base["conditions"][condition]
-        
-        return [{"condition": k, "data": v} for k, v in relevant_conditions.items()]
-
-class EnhancedMedicalAssistant:
-    def __init__(self):
-        self.load_models()
-        self.rag_system = SimpleRAGSystem()
-        self.setup_langgraph()
-        self.conversation_count = {}
-    
-    def load_models(self):
-        """Load the AI models with fallback options"""
-        print("Loading models...")
-        try:
-            # Llama-2 for conversation
-            self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
-            if self.tokenizer.pad_token is None:
-                self.tokenizer.pad_token = self.tokenizer.eos_token
-            
-            self.model = AutoModelForCausalLM.from_pretrained(
-                "meta-llama/Llama-2-7b-chat-hf",
-                torch_dtype=torch.float16,
-                device_map="auto"
-            )
-            
-            # Meditron for medical suggestions
-            self.meditron_tokenizer = AutoTokenizer.from_pretrained("epfl-llm/meditron-7b")
-            if self.meditron_tokenizer.pad_token is None:
-                self.meditron_tokenizer.pad_token = self.meditron_tokenizer.eos_token
-                
-            self.meditron_model = AutoModelForCausalLM.from_pretrained(
-                "epfl-llm/meditron-7b",
-                torch_dtype=torch.float16,
-                device_map="auto"
-            )
-            print("Models loaded successfully!")
-            
-        except Exception as e:
-            print(f"Error loading models: {e}")
-            # Fallback - use smaller models
-            self.tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
-            self.model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
-            self.meditron_tokenizer = self.tokenizer
-            self.meditron_model = self.model
-    
-    def setup_langgraph(self):
-        """Setup LangGraph workflow"""
-        workflow = StateGraph(MedicalState)
-        
-        workflow.add_node("intake", self.patient_intake)
-        workflow.add_node("generate_recommendations", self.generate_recommendations)
-        workflow.add_node("emergency_triage", self.emergency_triage)
-        
-        workflow.set_entry_point("intake")
-        workflow.add_conditional_edges(
-            "intake",
-            self.route_after_intake,
-            {
-                "emergency": "emergency_triage",
-                "recommendations": "generate_recommendations"
-            }
+from typing import TypedDict, List, Tuple
+import json
+
+# Model configuration
+LLAMA_MODEL = "meta-llama/Llama-2-7b-chat-hf"
+MEDITRON_MODEL = "epfl-llm/meditron-7b"
+
+SYSTEM_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.
+Ask 1-2 follow-up questions at a time to gather more details about:
+- Detailed description of symptoms
+- Duration (when did it start?)
+- Severity (scale of 1-10)
+- Aggravating or alleviating factors
+- Related symptoms
+- Medical history
+- Current medications and allergies
+After collecting sufficient information (4-5 exchanges), summarize findings and suggest when they should seek professional care. Do NOT make specific diagnoses or recommend specific treatments.
+Respond empathetically and clearly. Always be professional and thorough."""
+
+MEDITRON_PROMPT = """<|im_start|>system
+You are a specialized medical assistant focusing ONLY on suggesting over-the-counter medicines and home remedies based on patient information.
+Based on the following patient information, provide ONLY:
+1. One specific over-the-counter medicine with proper adult dosing instructions
+2. One practical home remedy that might help
+3. Clear guidance on when to seek professional medical care
+Be concise, practical, and focus only on general symptom relief. Do not diagnose. Include a disclaimer that you are not a licensed medical professional.
+<|im_end|>
+<|im_start|>user
+Patient information: {patient_info}
+<|im_end|>
+<|im_start|>assistant
+"""
+
+# Load models
+print("Loading Llama-2 model...")
+tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL)
+model = AutoModelForCausalLM.from_pretrained(
+    LLAMA_MODEL,
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
+print("Llama-2 model loaded successfully!")
+
+print("Loading Meditron model...")
+meditron_tokenizer = AutoTokenizer.from_pretrained(MEDITRON_MODEL)
+meditron_model = AutoModelForCausalLM.from_pretrained(
+    MEDITRON_MODEL,
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
+print("Meditron model loaded successfully!")
+
+# Define the state for LangGraph
+class ConversationState(TypedDict):
+    messages: List[str]
+    history: List[Tuple[str, str]]
+    current_message: str
+    conversation_turns: int
+    patient_data: List[str]
+    llama_response: str
+    final_response: str
+    should_get_suggestions: bool
+
+def build_llama2_prompt(system_prompt, history, user_input):
+    """Format the conversation history and user input for Llama-2 chat models."""
+    prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"
+    
+    # Add conversation history
+    for user_msg, assistant_msg in history:
+        prompt += f"{user_msg} [/INST] {assistant_msg} </s><s>[INST] "
+    
+    # Add the current user input
+    prompt += f"{user_input} [/INST] "
+    
+    return prompt
+
+def get_meditron_suggestions(patient_info):
+    """Use Meditron model to generate medicine and remedy suggestions."""
+    prompt = MEDITRON_PROMPT.format(patient_info=patient_info)
+    inputs = meditron_tokenizer(prompt, return_tensors="pt").to(meditron_model.device)
+    
+    with torch.no_grad():
+        outputs = meditron_model.generate(
+            inputs.input_ids,
+            attention_mask=inputs.attention_mask,
+            max_new_tokens=256,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True
         )
-        workflow.add_edge("generate_recommendations", END)
-        workflow.add_edge("emergency_triage", END)
-        
-        self.workflow = workflow.compile()
-    
-    def patient_intake(self, state: MedicalState) -> MedicalState:
-        """Enhanced patient intake with RAG"""
-        last_message = state["conversation_history"][-1]["content"] if state["conversation_history"] else ""
-        
-        # Extract symptoms
-        detected_symptoms = self.extract_symptoms(last_message)
-        state["symptoms"].update(detected_symptoms)
-        
-        # Use RAG to get relevant medical knowledge
-        if detected_symptoms:
-            symptom_names = list(detected_symptoms.keys())
-            relevant_conditions = self.rag_system.retrieve_relevant_conditions(symptom_names)
-            state["retrieved_knowledge"] = relevant_conditions
-        
-        # Check for red flags
-        red_flags = self.check_red_flags(last_message)
-        state["red_flags"].extend(red_flags)
-        
-        # Determine consultation stage
-        if red_flags:
-            state["consultation_stage"] = "emergency"
-        else:
-            state["consultation_stage"] = "recommendations"
-        
-        return state
-    
-    def generate_recommendations(self, state: MedicalState) -> MedicalState:
-        """Generate RAG-enhanced recommendations"""
-        # Create structured recommendations from RAG knowledge
-        recommendations = self.create_structured_recommendations(state)
-        state["suggested_actions"] = recommendations
-        return state
-    
-    def create_structured_recommendations(self, state: MedicalState) -> List[str]:
-        """Create structured recommendations using RAG knowledge"""
-        recommendations = []
-        
-        if not state["retrieved_knowledge"]:
-            recommendations.append("I need more specific information about your symptoms to provide targeted recommendations.")
-            return recommendations
-        
-        # Process each relevant condition
-        for knowledge_item in state["retrieved_knowledge"][:2]:  # Limit to top 2 conditions
-            condition = knowledge_item["condition"]
-            data = knowledge_item["data"]
-            
-            # Format condition information
-            condition_info = f"\n**Possible Condition: {condition.replace('_', ' ').title()}**\n"
-            
-            # Add medications
-            if "otc_medications" in data:
-                condition_info += "\n**💊 Over-the-Counter Medications:**\n"
-                for med in data["otc_medications"]:
-                    condition_info += f"• **{med['name']}**: {med['dose']} (Max daily: {med['max_daily']})\n"
-            
-            # Add home remedies
-            if "home_remedies" in data:
-                condition_info += "\n**🏠 Home Remedies:**\n"
-                for remedy in data["home_remedies"]:
-                    condition_info += f"• {remedy}\n"
-            
-            # Add when to seek care
-            if "when_to_seek_care" in data:
-                condition_info += f"\n**⚠️ Seek Medical Care If:** {data['when_to_seek_care']}\n"
-            
-            recommendations.append(condition_info)
-        
-        # Add general advice
-        recommendations.append("""
-**📋 General Recommendations:**
-• Monitor your symptoms for any changes
-• Stay hydrated and get adequate rest
-• Follow medication instructions carefully
-• Don't exceed recommended dosages
-
-**🚨 Emergency Warning Signs:**
-• Severe worsening of symptoms
-• High fever (>101.3°F/38.5°C)
-• Difficulty breathing
-• Severe pain
-• Signs of dehydration
-        """)
-        
-        return recommendations
-    
-    def emergency_triage(self, state: MedicalState) -> MedicalState:
-        """Handle emergency situations"""
-        emergency_response = f"""
-🚨 **URGENT MEDICAL ATTENTION NEEDED** 🚨
-
-Based on your symptoms, I strongly recommend seeking immediate medical care because you mentioned: {', '.join(state['red_flags'])}
-
-**Immediate Actions:**
-• Go to the nearest emergency room, OR
-• Call emergency services (911), OR  
-• Contact your doctor immediately
-
-**Why This is Urgent:**
-These symptoms can indicate serious conditions that require professional medical evaluation and treatment.
-
-⚠️ **Disclaimer:** This is not a medical diagnosis, but these symptoms warrant immediate professional assessment.
-        """
-        
-        state["suggested_actions"] = [emergency_response]
-        return state
     
-    def route_after_intake(self, state: MedicalState):
-        """Route decision after intake"""
-        if state["red_flags"]:
-            return "emergency"
-        else:
-            return "recommendations"
-    
-    def extract_symptoms(self, text: str) -> Dict:
-        """Extract and categorize symptoms from patient text"""
-        symptoms = {}
-        text_lower = text.lower()
-        
-        for category, symptom_list in MEDICAL_CATEGORIES.items():
-            for symptom in symptom_list:
-                if symptom in text_lower:
-                    symptoms[symptom] = {
-                        "category": category,
-                        "mentioned_at": datetime.now().isoformat(),
-                        "context": text
-                    }
-        
-        return symptoms
+    suggestion = meditron_tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+    return suggestion
+
+# LangGraph Node Functions
+def initialize_conversation(state: ConversationState) -> ConversationState:
+    """Initialize or update conversation state."""
+    # Update conversation turns
+    state["conversation_turns"] = state.get("conversation_turns", 0) + 1
+    
+    # Add current message to patient data
+    if "patient_data" not in state:
+        state["patient_data"] = []
+    state["patient_data"].append(state["current_message"])
+    
+    # Determine if we should get suggestions (after 4 turns)
+    state["should_get_suggestions"] = state["conversation_turns"] >= 4
+    
+    return state
+
+def generate_llama_response(state: ConversationState) -> ConversationState:
+    """Generate response using Llama-2 model."""
+    # Build the prompt with proper Llama-2 formatting
+    prompt = build_llama2_prompt(SYSTEM_PROMPT, state["history"], state["current_message"])
+    
+    # Add summarization instruction after 4 turns
+    if state["conversation_turns"] >= 4:
+        prompt = prompt.replace("[/INST] ", "[/INST] Now summarize what you've learned and suggest when professional care may be needed. ")
+    
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    # Generate the Llama-2 response
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs.input_ids,
+            attention_mask=inputs.attention_mask,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
     
-    def check_red_flags(self, text: str) -> List[str]:
-        """Check for emergency red flags"""
-        found_flags = []
-        text_lower = text.lower()
-        
-        for flag in RED_FLAGS:
-            if flag in text_lower:
-                found_flags.append(flag)
-        
-        return found_flags
+    # Decode and extract Llama-2's response
+    full_response = tokenizer.decode(outputs[0], skip_special_tokens=False)
+    llama_response = full_response.split('[/INST]')[-1].split('</s>')[0].strip()
+    
+    state["llama_response"] = llama_response
+    return state
+
+def generate_medicine_suggestions(state: ConversationState) -> ConversationState:
+    """Generate medicine suggestions using Meditron model."""
+    # Collect full patient conversation
+    full_patient_info = "\n".join(state["patient_data"]) + "\n\nSummary: " + state["llama_response"]
+    
+    # Get medicine suggestions
+    medicine_suggestions = get_meditron_suggestions(full_patient_info)
+    
+    # Format final response
+    final_response = (
+        f"{state['llama_response']}\n\n"
+        f"--- MEDICATION AND HOME CARE SUGGESTIONS ---\n\n"
+        f"{medicine_suggestions}"
+    )
+    
+    state["final_response"] = final_response
+    return state
+
+def finalize_response(state: ConversationState) -> ConversationState:
+    """Finalize the response without medicine suggestions."""
+    state["final_response"] = state["llama_response"]
+    return state
+
+def should_get_suggestions(state: ConversationState) -> str:
+    """Conditional edge to determine next step."""
+    if state["should_get_suggestions"]:
+        return "get_suggestions"
+    else:
+        return "finalize"
+
+# Create the LangGraph workflow
+def create_medical_workflow():
+    """Create the LangGraph workflow for medical assistant."""
+    workflow = StateGraph(ConversationState)
+    
+    # Add nodes
+    workflow.add_node("initialize", initialize_conversation)
+    workflow.add_node("generate_llama", generate_llama_response)
+    workflow.add_node("get_suggestions", generate_medicine_suggestions)
+    workflow.add_node("finalize", finalize_response)
+    
+    # Define the flow
+    workflow.set_entry_point("initialize")
+    workflow.add_edge("initialize", "generate_llama")
+    workflow.add_conditional_edges(
+        "generate_llama",
+        should_get_suggestions,
+        {
+            "get_suggestions": "get_suggestions",
+            "finalize": "finalize"
+        }
+    )
+    workflow.add_edge("get_suggestions", END)
+    workflow.add_edge("finalize", END)
+    
+    return workflow.compile()
+
+# Initialize the workflow
+medical_workflow = create_medical_workflow()
+
+# Conversation state tracking (for Gradio session management)
+conversation_states = {}
+
+@spaces.GPU
+def generate_response(message, history):
+    """Generate a response using the LangGraph workflow."""
+    session_id = "default-session"
+    
+    # Initialize or get existing conversation state
+    if session_id not in conversation_states:
+        conversation_states[session_id] = {
+            "messages": [],
+            "history": [],
+            "conversation_turns": 0,
+            "patient_data": []
+        }
     
-    def generate_response(self, message: str, history: List) -> str:
-        """Main response generation function"""
-        session_id = "default_session"
-        
-        # Track conversation count
-        if session_id not in self.conversation_count:
-            self.conversation_count[session_id] = 0
-        self.conversation_count[session_id] += 1
-        
-        # Initialize state
-        state = MedicalState(
-            patient_id=session_id,
-            conversation_history=history + [{"role": "user", "content": message}],
-            symptoms={},
-            vital_questions_asked=[],
-            medical_history={},
-            current_medications=[],
-            allergies=[],
-            severity_scores={},
-            red_flags=[],
-            assessment_complete=False,
-            suggested_actions=[],
-            consultation_stage="intake",
-            retrieved_knowledge=[],
-            confidence_scores={}
-        )
-        
-        # For first few messages, do conversational intake
-        if self.conversation_count[session_id] <= 3:
-            return self.generate_conversational_response(message, history)
-        
-        # After gathering info, run workflow for recommendations
-        try:
-            result = self.workflow.invoke(state)
-            return self.format_final_response(result)
-        except Exception as e:
-            print(f"Workflow error: {e}")
-            return self.generate_conversational_response(message, history)
+    # Update state with current message and history
+    state = conversation_states[session_id].copy()
+    state["current_message"] = message
+    state["history"] = history
     
-    def generate_conversational_response(self, message: str, history: List) -> str:
-        """Generate conversational response for intake phase"""
-        # Extract symptoms for context
-        symptoms = self.extract_symptoms(message)
-        red_flags = self.check_red_flags(message)
-        
-        # Handle emergencies immediately
-        if red_flags:
-            return f"""
-🚨 **URGENT MEDICAL ATTENTION NEEDED** 🚨
-
-I notice you mentioned: {', '.join(red_flags)}
-
-Please seek immediate medical care:
-• Go to the nearest emergency room
-• Call emergency services (911)  
-• Contact your doctor immediately
-
-These symptoms require professional medical evaluation right away.
-            """
-        
-        # Generate contextual questions based on symptoms
-        if symptoms:
-            symptom_names = list(symptoms.keys())
-            prompt = f"""
-You are a caring medical assistant. The patient mentioned these symptoms: {', '.join(symptom_names)}.
-
-Respond empathetically and ask 1-2 relevant follow-up questions about:
-- How long they've had these symptoms
-- Severity (mild, moderate, severe)
-- What makes it better or worse
-- Any other symptoms they're experiencing
-
-Be professional, caring, and concise. Don't provide treatment advice yet.
-            """
-        else:
-            prompt = f"""
-You are a caring medical assistant. The patient said: "{message}"
-
-Respond empathetically and ask relevant questions to understand their health concern better.
-Be professional and caring.
-            """
-        
-        return self.generate_llama_response(prompt)
+    # Run the workflow
+    result = medical_workflow.invoke(state)
     
-    def generate_llama_response(self, prompt: str) -> str:
-        """Generate response using Llama-2 with better formatting"""
-        try:
-            formatted_prompt = f"<s>[INST] {prompt} [/INST]"
-            inputs = self.tokenizer(formatted_prompt, return_tensors="pt", truncation=True, max_length=512)
-            
-            if torch.cuda.is_available():
-                inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
-            
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    **inputs,
-                    max_new_tokens=200,
-                    temperature=0.7,
-                    top_p=0.9,
-                    do_sample=True,
-                    pad_token_id=self.tokenizer.eos_token_id
-                )
-            
-            # Decode response
-            response = self.tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
-            
-            # Clean up the response
-            response = response.split('</s>')[0].strip()
-            response = response.replace('<s>', '').replace('[INST]', '').replace('[/INST]', '').strip()
-            
-            # Remove any XML-like tags
-            response = re.sub(r'<[^>]+>', '', response)
-            
-            return response if response else "I understand your concern. Can you tell me more about what you're experiencing?"
-            
-        except Exception as e:
-            print(f"Error generating response: {e}")
-            return "I understand your concern. Can you tell me more about your symptoms?"
+    # Update the stored conversation state
+    conversation_states[session_id] = {
+        "messages": result["messages"] if "messages" in result else [],
+        "history": history,
+        "conversation_turns": result["conversation_turns"],
+        "patient_data": result["patient_data"]
+    }
     
-    def format_final_response(self, state: MedicalState) -> str:
-        """Format the final response with recommendations"""
-        if state["consultation_stage"] == "emergency":
-            return state["suggested_actions"][0] if state["suggested_actions"] else "Please seek immediate medical attention."
-        
-        # Format recommendations nicely
-        if state["suggested_actions"]:
-            response = "## 🏥 Medical Assessment & Recommendations\n\n"
-            response += "Based on our conversation, here's what I recommend:\n"
-            
-            for action in state["suggested_actions"]:
-                response += f"{action}\n"
-            
-            response += "\n---\n"
-            response += "**Important Disclaimer:** I'm an AI assistant providing general health information. "
-            response += "This is not a substitute for professional medical advice, diagnosis, or treatment. "
-            response += "Always consult with qualified healthcare providers for medical concerns."
-            
-            return response
-        else:
-            return "Please provide more details about your symptoms so I can offer better guidance."
+    return result["final_response"]
 
-# Initialize the medical assistant
-medical_assistant = EnhancedMedicalAssistant()
-
-# Gradio chat interface function
-def chat_interface(message, history):
-    """Gradio chat interface"""
-    try:
-        return medical_assistant.generate_response(message, history)
-    except Exception as e:
-        print(f"Chat interface error: {e}")
-        return f"I apologize, but I encountered an error. Please try rephrasing your question. Error: {str(e)}"
-
-# Create Gradio interface with enhanced styling
+# Create the Gradio interface
 demo = gr.ChatInterface(
-    fn=chat_interface,
-    title="🏥 Medical AI Assistant with medRAG",
-    description="""
-    I'm an AI medical assistant powered by medical knowledge retrieval (medRAG). 
-    I can help assess your symptoms and provide evidence-based recommendations.
-    
-    **How it works:**
-    1. Tell me about your symptoms
-    2. I'll ask follow-up questions  
-    3. I'll provide personalized recommendations based on medical knowledge
-    
-    ⚠️ **Important**: I'm not a replacement for professional medical care. Always consult healthcare providers for serious concerns.
-    """,
+    fn=generate_response,
+    title="Medical Assistant with LangGraph & Medicine Suggestions",
+    description="Tell me about your symptoms, and after gathering enough information, I'll suggest potential remedies using an AI workflow.",
     examples=[
-        "I have a bad cough and sore throat",
-        "I've been having headaches for the past few days", 
-        "My stomach has been hurting after meals",
-        "I have chest pain and trouble breathing"
+        "I have a cough and my throat hurts",
+        "I've been having headaches for a week",
+        "My stomach has been hurting since yesterday"
     ],
-    theme="soft",
-    css="""
-    /* Main container styling */
-    .gradio-container {
-        background: linear-gradient(to right, #f0f8ff, #f5f5f5);
-        font-family: 'Arial', sans-serif;
-    }
-    
-    /* Chat message styling */
-    .message.user { 
-        background-color: #e3f2fd; 
-        border-radius: 12px;
-        padding: 12px;
-        margin: 8px;
-        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
-        border-left: 4px solid #2196F3;
-    }
-    
-    .message.bot { 
-        background-color: #f1f8e9; 
-        border-radius: 12px;
-        padding: 12px;
-        margin: 8px;
-        box-shadow: 0 2px 5px rgba(0,0,0,0.1);
-        border-left: 4px solid #4CAF50;
-    }
-    
-    /* Enhanced medical styling */
-    .bot h2 {
-        color: #1976D2 !important;
-        border-bottom: 2px solid #E0E0E0 !important;
-        padding-bottom: 8px !important;
-    }
-    """
+    theme="soft"
 )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch()

requirements.txt

@@ -1,45 +1,18 @@
-# Core dependencies
-gradio>=4.0.0
-torch>=2.0.0
-transformers>=4.30.0
-langgraph>=0.3.27
-langchain-core>=0.2.38
+# Core packages
+gradio==4.24.0
+spaces==0.21.1
 
-# Vector search and embeddings
-sentence-transformers>=2.2.2
-faiss-cpu>=1.7.4
+# Transformers & tokenization
+transformers==4.40.1
+torch>=2.1.0
 
-# Data processing
-numpy>=1.24.0
-typing-extensions>=4.5.0
-
-# LangGraph ecosystem components
-langsmith>=0.1.63
-langgraph-sdk>=0.1.66
-langgraph-checkpoint>=2.0.23
-
-# Web serving utilities
-httpx>=0.25.0
-uvicorn>=0.26.0
-sse-starlette>=2.1.0,<2.2.0
-uvloop>=0.18.0
-httptools>=0.5.0
+# LangGraph
+langgraph==0.0.41
 
-# Serialization and utilities
-orjson>=3.9.7,<3.10.17
-jsonschema-rs>=0.20.0
-structlog>=24.1.0
-cloudpickle>=3.0.0
-tenacity>=8.0.0
+# Optional but often required for transformers
+accelerate==0.30.1
+sentencepiece==0.1.99
+protobuf==4.25.3
 
-# Model acceleration (optional but recommended)
-accelerate>=0.20.0
-safetensors>=0.3.1
-bitsandbytes>=0.40.0
-
-# For Hugging Face model access
-huggingface_hub>=0.16.0
-
-# Optional - specific model support
-langchain_anthropic>=0.0.5
-langchain_openai>=0.0.2
+# Utility
+typing-extensions>=4.5.0