utils/medical_agent_lite.py

import json
from typing import Dict, Any, List
from collections import defaultdict

class LiteMedicalAgent:
    """Lightweight medical agent optimized for ZeroGPU - no heavy models"""
    def __init__(self):
        self.sessions = defaultdict(dict)
        self.max_questions = 3
        print("✅ Lightweight Medical Agent initialized (rule-based, no LLM)")

    def process_input(self, english_text: str, session_id: str) -> Dict[str, Any]:
        """Main entry point with session management"""
        # Get or initialize session state
        if session_id not in self.sessions:
            self.sessions[session_id] = {
                'question_count': 0,
                'conversation_history': [],
                'symptoms': []
            }
        
        session_state = self.sessions[session_id]
        current_count = session_state['question_count'] + 1
        session_state['question_count'] = current_count
        session_state['symptoms'].append(english_text)
        
        # Generate response based on question count
        state = 'questioning' if current_count < self.max_questions else 'summary'
        
        if state == 'summary':
            # Create summary
            all_symptoms = ", ".join(session_state['symptoms'])
            response = f"Thank you. Patient reported: {all_symptoms}. Please consult with a healthcare provider for proper diagnosis."
        else:
            # Get next question based on symptoms and count
            response = self._get_contextual_question(english_text, current_count, session_state['symptoms'])
        
        return {
            'response': response,
            'question_count': current_count,
            'state': state,
            'workflow_used': True
        }
    
    def _get_contextual_question(self, current_input: str, question_num: int, previous_symptoms: List[str]) -> str:
        """Generate contextual medical follow-up questions"""
        current_lower = current_input.lower()
        
        # First question - get duration
        if question_num == 1:
            if any(word in current_lower for word in ['pain', 'hurt', 'ache', 'sore']):
                return "How long have you had this pain?"
            elif any(word in current_lower for word in ['cough', 'fever', 'cold']):
                return "When did symptoms start?"
            else:
                return "How long have symptoms lasted?"
        
        # Second question - get severity/location
        elif question_num == 2:
            if any(word in current_lower for word in ['pain', 'hurt', 'ache']):
                return "Where exactly is the pain?"
            elif any(word in current_lower for word in ['fever', 'temperature']):
                return "Do you have high fever?"
            elif any(word in current_lower for word in ['days', 'weeks', 'hours']):
                return "Rate severity from 1 to 10?"
            else:
                return "Any other associated symptoms?"
        
        # Third question - get additional details
        else:
            if any(word in current_lower for word in ['severe', 'bad', 'terrible']):
                return "Any difficulty breathing?"
            elif any(word in current_lower for word in ['head', 'chest', 'stomach', 'back']):
                return "What makes it worse or better?"
            else:
                return "Any recent changes or triggers?"
    
    def process_doctor_input(self, doctor_text: str) -> str:
        """Process doctor's input and simplify for patient"""
        doctor_lower = doctor_text.lower()
        
        # Map doctor's complex questions to simple ones
        if any(word in doctor_lower for word in ['duration', 'how long', 'when']):
            return "How long have symptoms lasted?"
        elif any(word in doctor_lower for word in ['location', 'where']):
            return "Where is the problem?"
        elif any(word in doctor_lower for word in ['severity', 'rate', 'scale']):
            return "Rate severity 1-10?"
        elif any(word in doctor_lower for word in ['associate', 'other', 'additional']):
            return "Any other symptoms?"
        elif any(word in doctor_lower for word in ['worsen', 'better', 'trigger']):
            return "What makes it worse?"
        else:
            return "Please describe more details?"