agents/assessment_agent.py

from typing import Dict, List
from .base_agent import BaseWellnessAgent
import json
from pathlib import Path

class AssessmentAgent(BaseWellnessAgent):
    """Agent specialized in mental health assessments"""
    
    def __init__(self, model_config: Dict, verbose: bool = False):
        super().__init__(
            name="Mental Health Assessment Agent",
            role="Mental Health Evaluator",
            goal="Conduct thorough mental health assessments and provide actionable insights",
            backstory="""I am an AI agent specialized in mental health assessment. 
            I use validated assessment tools and techniques to evaluate mental well-being 
            and provide personalized recommendations.""",
            tools=["emotion_detection", "conversation"],
            model_config=model_config,
            verbose=verbose
        )
        self.assessment_tools = self._load_assessment_tools()
        self.current_assessment = None
        
    def _load_assessment_tools(self) -> Dict:
        """Load assessment questionnaires and scoring guides"""
        tools_path = Path(__file__).parent.parent / "knowledge_base" / "assessment_tools.json"
        try:
            with open(tools_path) as f:
                return json.load(f)
        except FileNotFoundError:
            return {
                "phq9": {
                    "name": "PHQ-9 Depression Scale",
                    "questions": [
                        "Little interest or pleasure in doing things?",
                        "Feeling down, depressed, or hopeless?",
                        # Add more questions...
                    ],
                    "scoring": {
                        "ranges": {
                            "0-4": "Minimal depression",
                            "5-9": "Mild depression",
                            "10-14": "Moderate depression",
                            "15-19": "Moderately severe depression",
                            "20-27": "Severe depression"
                        }
                    }
                },
                "gad7": {
                    "name": "GAD-7 Anxiety Scale",
                    "questions": [
                        "Feeling nervous, anxious, or on edge?",
                        "Not being able to stop or control worrying?",
                        # Add more questions...
                    ],
                    "scoring": {
                        "ranges": {
                            "0-4": "Minimal anxiety",
                            "5-9": "Mild anxiety",
                            "10-14": "Moderate anxiety",
                            "15-21": "Severe anxiety"
                        }
                    }
                }
            }
            
    def start_assessment(self, assessment_type: str) -> Dict:
        """Start a new assessment"""
        if assessment_type not in self.assessment_tools:
            raise ValueError(f"Unknown assessment type: {assessment_type}")
            
        self.current_assessment = {
            "type": assessment_type,
            "tool": self.assessment_tools[assessment_type],
            "responses": [],
            "current_question": 0
        }
        
        return self.get_next_question()
        
    def process_message(self, message: str) -> Dict:
        """Process user response and continue assessment"""
        if not self.current_assessment:
            return self.format_response(
                "No active assessment. Please start an assessment first."
            )
            
        # Record response
        self.current_assessment["responses"].append({
            "question": self.current_assessment["current_question"],
            "response": message,
            "emotion": self.analyze_emotion(message)
        })
        
        # Move to next question or finish assessment
        return (
            self.finish_assessment()
            if self._is_assessment_complete()
            else self.get_next_question()
        )
        
    def get_next_question(self) -> Dict:
        """Get the next question in the assessment"""
        if not self.current_assessment:
            return self.format_response(
                "No active assessment. Please start an assessment first."
            )
            
        tool = self.current_assessment["tool"]
        question_idx = self.current_assessment["current_question"]
        
        if question_idx >= len(tool["questions"]):
            return self.finish_assessment()
            
        question = tool["questions"][question_idx]
        self.current_assessment["current_question"] += 1
        
        return self.format_response(question)
        
    def _is_assessment_complete(self) -> bool:
        """Check if all questions have been answered"""
        if not self.current_assessment:
            return False
            
        return len(self.current_assessment["responses"]) >= len(
            self.current_assessment["tool"]["questions"]
        )
        
    def finish_assessment(self) -> Dict:
        """Complete the assessment and generate report"""
        if not self.current_assessment:
            return self.format_response(
                "No active assessment to finish."
            )
            
        # Calculate scores
        scores = self._calculate_scores()
        
        # Generate report
        report = self._generate_report(scores)
        
        # Add to history
        self.add_to_history({
            "assessment_type": self.current_assessment["type"],
            "scores": scores,
            "report": report
        })
        
        # Reset current assessment
        self.current_assessment = None
        
        return self.format_response(report)
        
    def _calculate_scores(self) -> Dict:
        """Calculate assessment scores"""
        # Implement scoring logic based on assessment type
        return {
            "total_score": 0,  # Calculate actual score
            "subscores": {},
            "severity": "Unknown"  # Determine from scoring ranges
        }
        
    def _generate_report(self, scores: Dict) -> str:
        """Generate detailed assessment report"""
        report = f"""Assessment Report
        
        Type: {self.current_assessment["type"].upper()}
        Total Score: {scores["total_score"]}
        Severity: {scores["severity"]}
        
        Recommendations:
        1. Continue monitoring your mental health
        2. Consider discussing results with a mental health professional
        3. Practice self-care and stress management techniques
        
        Note: This assessment is for screening purposes only and does not constitute 
        a clinical diagnosis. Please consult with a qualified mental health professional 
        for proper evaluation and treatment."""
        
        return report
        
    def get_assessment_history(self) -> List[Dict]:
        """Get history of completed assessments"""
        return self.get_history()
        
    def get_available_assessments(self) -> List[str]:
        """Get list of available assessment tools"""
        return list(self.assessment_tools.keys())