Implement Mistral LLM support for all agents

agents/assessment_agent.py

@@ -1,188 +1,20 @@
 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):
+    def __init__(self, model_config: Dict, **kwargs):
         super().__init__(
-            name="Mental Health Assessment Agent",
+            model_config=model_config,
+            agent_type="assessment",
             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()) 
+            verbose=True,
+            allow_delegation=False,
+            tools=[],  # Tools will be added as needed
+            **kwargs
+        ) 

agents/base_agent.py

@@ -1,91 +1,139 @@
-from typing import Dict, List, Optional
-from crewai import Agent
-from pydantic import BaseModel
-from transformers import pipeline
+from typing import Dict, List, Optional, Any
+from crewai import Agent, Task
+import logging
+from utils.log_manager import LogManager
+from pydantic import Field, BaseModel, ConfigDict
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-class AgentState(BaseModel):
-    """State management for agents"""
-    context: Dict = {}
-    history: List[Dict] = []
-    current_task: Optional[str] = None
-    
 class BaseWellnessAgent(Agent):
-    """Base class for all wellness platform agents"""
+    """Base agent class with Mistral LLM support"""
+    
+    # Allow arbitrary types in model
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    
+    # Define fields that will be used
+    log_manager: LogManager = Field(default_factory=LogManager)
+    logger: logging.Logger = Field(default=None)
+    config: Dict = Field(default_factory=dict)
+    model: Any = Field(default=None)
+    tokenizer: Any = Field(default=None)
+    agent_type: str = Field(default="base")
     
-    def __init__(
-        self,
-        name: str,
-        role: str,
-        goal: str,
-        backstory: str,
-        tools: List[str],
-        model_config: Dict,
-        verbose: bool = False
-    ):
+    def __init__(self, model_config: Dict, agent_type: str, **kwargs):
+        # Initialize the CrewAI agent first with required fields
         super().__init__(
-            name=name,
-            role=role,
-            goal=goal,
-            backstory=backstory,
-            verbose=verbose
+            role=kwargs.get("role", "Wellness Support Agent"),
+            goal=kwargs.get("goal", "Support mental wellness"),
+            backstory=kwargs.get("backstory", "I am an AI agent specialized in mental health support."),
+            verbose=kwargs.get("verbose", True),
+            allow_delegation=kwargs.get("allow_delegation", False),
+            tools=kwargs.get("tools", []),
+            **kwargs
         )
-        self.tools = tools
-        self.model_config = model_config
-        self.state = AgentState()
-        self._initialize_models()
-        
-    def _initialize_models(self):
-        """Initialize required AI models for the agent"""
-        self.models = {}
-        if "emotion_detection" in self.tools:
-            self.models["emotion"] = pipeline(
-                "text-classification",
-                model=self.model_config["emotion_detection"]["model_id"]
-            )
-            
-    def update_context(self, new_context: Dict):
-        """Update agent's context with new information"""
-        self.state.context.update(new_context)
         
-    def add_to_history(self, interaction: Dict):
-        """Add an interaction to agent's history"""
-        self.state.history.append(interaction)
+        # Initialize logging and configuration
+        self.config = model_config
+        self.agent_type = agent_type
+        self.logger = self.log_manager.get_agent_logger(agent_type)
         
-    def get_context(self) -> Dict:
-        """Get current context"""
-        return self.state.context
+        # Initialize Mistral model
+        self._initialize_model()
         
-    def get_history(self) -> List[Dict]:
-        """Get interaction history"""
-        return self.state.history
+        self.logger.info(f"{agent_type.capitalize()} Agent initialized")
         
-    def clear_state(self):
-        """Reset agent state"""
-        self.state = AgentState()
+    def _initialize_model(self):
+        """Initialize the Mistral model"""
+        try:
+            model_config = self.config[self.agent_type]
+            self.logger.info(f"Initializing Mistral model: {model_config['model_id']}")
+            
+            # Initialize tokenizer and model
+            self.tokenizer = AutoTokenizer.from_pretrained(model_config["model_id"])
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_config["model_id"],
+                torch_dtype=torch.float32,
+                device_map="auto",
+                load_in_4bit=True
+            )
+            
+            self.logger.info("Mistral model initialized successfully")
+            
+        except Exception as e:
+            self.logger.error(f"Error initializing Mistral model: {str(e)}")
+            raise
+            
+    def _generate_response(self, input_text: str) -> str:
+        """Generate response using Mistral model"""
+        try:
+            # Prepare input with instruction template
+            template = self.config[self.agent_type]["instruction_template"]
+            prompt = template.format(input=input_text)
+            
+            # Tokenize input
+            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+            
+            # Generate response
+            outputs = self.model.generate(
+                **inputs,
+                max_length=self.config[self.agent_type].get("max_length", 4096),
+                temperature=self.config[self.agent_type].get("temperature", 0.7),
+                top_p=self.config[self.agent_type].get("top_p", 0.95),
+                repetition_penalty=self.config[self.agent_type].get("repetition_penalty", 1.1),
+                do_sample=True,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
+            
+            # Decode and clean response
+            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            response = response.replace(prompt, "").strip()
+            
+            return response
+            
+        except Exception as e:
+            self.logger.error(f"Error generating response: {str(e)}")
+            return "I apologize, but I encountered an error generating a response."
         
-    def analyze_emotion(self, text: str) -> Dict:
-        """Analyze emotion in text if emotion detection is available"""
-        if "emotion" in self.models:
-            result = self.models["emotion"](text)
-            return result[0] if result else {}
-        return {}
+    def execute_task(self, task: Task) -> str:
+        """Execute a task assigned to the agent"""
+        self.logger.info(f"Executing task: {task.description}")
         
-    def format_response(self, response: str) -> Dict:
-        """Format agent response with metadata"""
-        return {
-            "content": response,
-            "agent_name": self.name,
-            "agent_role": self.role,
-            "timestamp": self._get_timestamp(),
-            "context": self.get_context(),
-            "emotion_analysis": self.analyze_emotion(response)
-        }
+        try:
+            # Process the task description as a message
+            result = self.process_message(task.description)
+            return result["message"]
+            
+        except Exception as e:
+            self.logger.error(f"Error executing task: {str(e)}")
+            return "I apologize, but I encountered an error processing your request."
         
-    def _get_timestamp(self) -> str:
-        """Get current timestamp"""
-        from datetime import datetime
-        return datetime.now().isoformat()
+    def process_message(self, message: str, context: Dict = None) -> Dict:
+        """Process a message and return a response"""
+        self.logger.info("Processing message")
+        context = context or {}
         
-    def process_message(self, message: str) -> Dict:
-        """Process incoming message - to be implemented by child classes"""
-        raise NotImplementedError("Child classes must implement process_message") 
+        try:
+            # Generate response using Mistral
+            response = self._generate_response(message)
+            
+            return {
+                "message": response,
+                "agent_type": self.agent_type,
+                "task_type": "dialogue"
+            }
+            
+        except Exception as e:
+            self.logger.error(f"Error processing message: {str(e)}")
+            return {
+                "message": "I apologize, but I encountered an error. Let me try a different approach.",
+                "agent_type": self.agent_type,
+                "task_type": "error_recovery"
+            }
+            
+    def get_status(self) -> Dict:
+        """Get the current status of the agent"""
+        return {
+            "type": self.agent_type,
+            "ready": bool(self.model and self.tokenizer),
+            "tools_available": len(self.tools)
+        }

agents/conversation_agent.py

@@ -1,23 +1,13 @@
-from typing import Dict, List, Optional, Any
-from crewai import Agent, Task
-import logging
-from utils.log_manager import LogManager
-from pydantic import Field, BaseModel, ConfigDict
+from typing import Dict
+from .base_agent import BaseWellnessAgent
 
-class ConversationAgent(Agent):
+class ConversationAgent(BaseWellnessAgent):
     """Agent specialized in therapeutic conversations"""
     
-    # Allow arbitrary types in model
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-    
-    # Define fields that will be used
-    log_manager: LogManager = Field(default_factory=LogManager)
-    logger: logging.Logger = Field(default=None)
-    config: Dict = Field(default_factory=dict)
-    
     def __init__(self, model_config: Dict, **kwargs):
-        # Initialize the CrewAI agent first with required fields
         super().__init__(
+            model_config=model_config,
+            agent_type="conversation",
             role="Therapeutic Conversation Specialist",
             goal="Guide therapeutic conversations and provide emotional support",
             backstory="""I am an AI agent specialized in therapeutic conversation techniques. 
@@ -27,83 +17,4 @@ class ConversationAgent(Agent):
             allow_delegation=False,
             tools=[],  # Tools will be added as needed
             **kwargs
-        )
-        
-        # Initialize logging and configuration after parent initialization
-        self.config = model_config
-        self.logger = self.log_manager.get_agent_logger("conversation")
-        
-        self.logger.info("Conversation Agent initialized")
-        
-    def execute_task(self, task: Task) -> str:
-        """Execute a task assigned to the agent"""
-        self.logger.info(f"Executing task: {task.description}")
-        
-        try:
-            # Process the task description as a message
-            result = self.process_message(task.description)
-            return result["message"]
-            
-        except Exception as e:
-            self.logger.error(f"Error executing task: {str(e)}")
-            return "I apologize, but I encountered an error processing your request."
-        
-    def process_message(self, message: str, context: Dict = None) -> Dict:
-        """Process a message and return a therapeutic response"""
-        self.logger.info("Processing message")
-        context = context or {}
-        
-        try:
-            # Generate therapeutic response
-            response = self._generate_therapeutic_response(message, context)
-            
-            return {
-                "message": response,
-                "agent_type": "conversation",
-                "task_type": "therapeutic_dialogue"
-            }
-            
-        except Exception as e:
-            self.logger.error(f"Error processing message: {str(e)}")
-            return {
-                "message": "I apologize, but I encountered an error. Let me try a different approach.",
-                "agent_type": "conversation",
-                "task_type": "error_recovery"
-            }
-            
-    def _generate_therapeutic_response(self, message: str, context: Dict) -> str:
-        """Generate a therapeutic response based on the input message and context"""
-        # Basic response templates
-        responses = {
-            "greeting": "Hello! I'm here to support you. How are you feeling today?",
-            "emotion": "I hear that you're feeling {}. Would you like to tell me more about that?",
-            "support": "That sounds challenging. I'm here to listen and support you.",
-            "clarification": "Could you tell me more about what you mean by that?",
-            "validation": "It's completely normal to feel that way. Your feelings are valid.",
-            "default": "I understand and I'm here to help. Could you tell me more about how you're feeling?"
-        }
-        
-        # Simple message analysis
-        message = message.lower()
-        
-        if "hello" in message or "hi" in message:
-            return responses["greeting"]
-        elif any(word in message for word in ["sad", "angry", "scared", "happy", "anxious"]):
-            emotion = next(word for word in ["sad", "angry", "scared", "happy", "anxious"] if word in message)
-            return responses["emotion"].format(emotion)
-        elif any(word in message for word in ["help", "support", "difficult", "hard"]):
-            return responses["support"]
-        elif len(message.split()) < 3:
-            return responses["clarification"]
-        elif "feel" in message or "feeling" in message:
-            return responses["validation"]
-        else:
-            return responses["default"]
-            
-    def get_status(self) -> Dict:
-        """Get the current status of the agent"""
-        return {
-            "type": "conversation",
-            "ready": True,
-            "tools_available": len(self.tools)
-        } 
+        ) 

agents/crisis_agent.py

@@ -1,251 +1,20 @@
-from typing import Dict, List
+from typing import Dict
 from .base_agent import BaseWellnessAgent
-import json
-from pathlib import Path
 
 class CrisisAgent(BaseWellnessAgent):
-    """Agent specialized in crisis intervention and emergency response"""
+    """Agent specialized in crisis intervention"""
     
-    def __init__(self, model_config: Dict, verbose: bool = False):
+    def __init__(self, model_config: Dict, **kwargs):
         super().__init__(
-            name="Crisis Intervention Agent",
-            role="Crisis Support Specialist",
+            model_config=model_config,
+            agent_type="crisis",
+            role="Crisis Intervention Specialist",
             goal="Provide immediate support and intervention in crisis situations",
             backstory="""I am an AI agent specialized in crisis intervention. 
-            I am trained to recognize and respond to emergency situations, 
-            provide immediate support, and connect users with appropriate resources.""",
-            tools=["emotion_detection", "conversation"],
-            model_config=model_config,
-            verbose=verbose
-        )
-        self.protocols = self._load_protocols()
-        self.current_intervention = None
-        self.emergency_contacts = self._load_emergency_contacts()
-        
-    def _load_protocols(self) -> Dict:
-        """Load crisis intervention protocols"""
-        protocols_path = Path(__file__).parent.parent / "knowledge_base" / "guidelines" / "crisis_protocols.json"
-        try:
-            with open(protocols_path) as f:
-                return json.load(f)
-        except FileNotFoundError:
-            return {
-                "suicide_risk": {
-                    "severity_levels": {
-                        "low": {
-                            "indicators": [
-                                "Passing thoughts of death",
-                                "No specific plan",
-                                "Some protective factors"
-                            ],
-                            "response": "Provide support and resources"
-                        },
-                        "medium": {
-                            "indicators": [
-                                "Frequent thoughts of death",
-                                "Vague plans",
-                                "Some risk factors present"
-                            ],
-                            "response": "Immediate counseling referral"
-                        },
-                        "high": {
-                            "indicators": [
-                                "Specific suicide plan",
-                                "Access to means",
-                                "Multiple risk factors"
-                            ],
-                            "response": "Emergency services contact"
-                        }
-                    },
-                    "immediate_actions": [
-                        "Express concern and care",
-                        "Assess immediate safety",
-                        "Connect with emergency services if needed"
-                    ]
-                },
-                "panic_attack": {
-                    "symptoms": [
-                        "Rapid heartbeat",
-                        "Difficulty breathing",
-                        "Feeling of doom"
-                    ],
-                    "interventions": [
-                        "Grounding techniques",
-                        "Breathing exercises",
-                        "Reassurance and support"
-                    ]
-                }
-            }
-            
-    def _load_emergency_contacts(self) -> Dict:
-        """Load emergency contact information"""
-        return {
-            "emergency_services": "911",
-            "crisis_hotline": "988",
-            "text_line": "741741",
-            "poison_control": "1-800-222-1222"
-        }
-        
-    def process_message(self, message: str) -> Dict:
-        """Process crisis-related message and provide appropriate response"""
-        # Analyze message for crisis indicators
-        risk_assessment = self._assess_risk(message)
-        
-        # Update intervention state
-        self._update_intervention_state(risk_assessment)
-        
-        # Generate appropriate response based on risk level
-        response = self._generate_crisis_response(risk_assessment)
-        
-        # Record interaction
-        self.add_to_history({
-            "message": message,
-            "risk_assessment": risk_assessment,
-            "response": response
-        })
-        
-        return self.format_response(response)
-        
-    def _assess_risk(self, message: str) -> Dict:
-        """Assess risk level from message content"""
-        # Analyze emotion
-        emotion = self.analyze_emotion(message)
-        
-        # Check for critical keywords
-        risk_level = self._determine_risk_level(message.lower())
-        
-        return {
-            "risk_level": risk_level,
-            "emotion": emotion,
-            "timestamp": self._get_timestamp()
-        }
-        
-    def _determine_risk_level(self, message: str) -> str:
-        """Determine risk level based on message content"""
-        high_risk_keywords = [
-            "kill myself", "end my life", "suicide plan",
-            "want to die", "no reason to live"
-        ]
-        medium_risk_keywords = [
-            "want to disappear", "everyone better off",
-            "cant take it", "hopeless"
-        ]
-        
-        if any(keyword in message for keyword in high_risk_keywords):
-            return "high"
-        elif any(keyword in message for keyword in medium_risk_keywords):
-            return "medium"
-        return "low"
-        
-    def _update_intervention_state(self, risk_assessment: Dict):
-        """Update or create intervention state"""
-        if not self.current_intervention:
-            self.current_intervention = {
-                "start_time": self._get_timestamp(),
-                "risk_assessments": [],
-                "escalation_level": "initial"
-            }
-            
-        self.current_intervention["risk_assessments"].append(risk_assessment)
-        self._update_escalation_level(risk_assessment["risk_level"])
-        
-    def _update_escalation_level(self, risk_level: str):
-        """Update intervention escalation level"""
-        if risk_level == "high":
-            self.current_intervention["escalation_level"] = "emergency"
-        elif risk_level == "medium" and self.current_intervention["escalation_level"] == "initial":
-            self.current_intervention["escalation_level"] = "elevated"
-            
-    def _generate_crisis_response(self, risk_assessment: Dict) -> str:
-        """Generate appropriate crisis response"""
-        risk_level = risk_assessment["risk_level"]
-        escalation_level = self.current_intervention["escalation_level"]
-        
-        if risk_level == "high" or escalation_level == "emergency":
-            return self._generate_emergency_response()
-        elif risk_level == "medium" or escalation_level == "elevated":
-            return self._generate_elevated_response()
-        else:
-            return self._generate_supportive_response()
-            
-    def _generate_emergency_response(self) -> str:
-        """Generate response for emergency situations"""
-        return f"""I'm very concerned about your safety right now. This is an emergency situation that requires immediate attention.
-
-EMERGENCY CONTACTS:
-- Emergency Services: {self.emergency_contacts['emergency_services']}
-- Crisis Hotline: {self.emergency_contacts['crisis_hotline']}
-- Crisis Text Line: Text HOME to {self.emergency_contacts['text_line']}
-
-Please:
-1. Call emergency services immediately
-2. Stay on the line with me
-3. If possible, contact a trusted person to be with you
-
-Your life has value and people care about you. Help is available 24/7."""
-        
-    def _generate_elevated_response(self) -> str:
-        """Generate response for elevated risk situations"""
-        return f"""I hear how much pain you're in, and I want to help ensure your safety.
-
-Please consider:
-1. Calling the Crisis Hotline: {self.emergency_contacts['crisis_hotline']}
-2. Texting HOME to {self.emergency_contacts['text_line']}
-3. Reaching out to a trusted friend, family member, or counselor
-
-Would you be willing to talk about what's troubling you? I'm here to listen without judgment."""
-        
-    def _generate_supportive_response(self) -> str:
-        """Generate supportive response for lower risk situations"""
-        return """I can hear that you're going through a difficult time. Your feelings are valid, and you're not alone.
-
-Let's talk about what's happening. I'm here to:
-1. Listen and understand
-2. Help you explore your feelings
-3. Discuss coping strategies
-4. Connect you with resources
-
-What would be most helpful for you right now?"""
-        
-    def end_intervention(self) -> Dict:
-        """End current crisis intervention"""
-        if not self.current_intervention:
-            return self.format_response(
-                "No active crisis intervention to end."
-            )
-            
-        # Generate intervention summary
-        summary = self._generate_intervention_summary()
-        
-        # Add final summary to history
-        self.add_to_history({
-            "intervention_summary": summary,
-            "end_time": self._get_timestamp()
-        })
-        
-        # Reset current intervention
-        self.current_intervention = None
-        
-        return self.format_response(summary)
-        
-    def _generate_intervention_summary(self) -> str:
-        """Generate summary of crisis intervention"""
-        assessments = self.current_intervention["risk_assessments"]
-        max_risk = max(a["risk_level"] for a in assessments)
-        
-        return f"""Crisis Intervention Summary
-        
-        Duration: {self._calculate_duration(
-            self.current_intervention["start_time"],
-            self._get_timestamp()
-        ) // 60} minutes
-        Maximum Risk Level: {max_risk}
-        Escalation Level: {self.current_intervention["escalation_level"]}
-        
-        Follow-up Recommendations:
-        1. Continue monitoring emotional state
-        2. Maintain contact with support system
-        3. Follow up with mental health professional
-        4. Keep crisis resources readily available
-        
-        Remember: Support is always available at {self.emergency_contacts['crisis_hotline']}""" 
+            I prioritize user safety and provide immediate support while ensuring 
+            appropriate escalation to emergency services when necessary.""",
+            verbose=True,
+            allow_delegation=False,
+            tools=[],  # Tools will be added as needed
+            **kwargs
+        ) 

agents/mindfulness_agent.py

@@ -1,227 +1,20 @@
-from typing import Dict, List
+from typing import Dict
 from .base_agent import BaseWellnessAgent
-import json
-from pathlib import Path
 
 class MindfulnessAgent(BaseWellnessAgent):
     """Agent specialized in mindfulness and meditation guidance"""
     
-    def __init__(self, model_config: Dict, verbose: bool = False):
+    def __init__(self, model_config: Dict, **kwargs):
         super().__init__(
-            name="Mindfulness & Meditation Agent",
-            role="Mindfulness Guide",
-            goal="Guide users through mindfulness exercises and meditation sessions",
-            backstory="""I am an AI agent specialized in mindfulness and meditation guidance. 
-            I provide calming exercises, breathing techniques, and meditation sessions 
-            tailored to users' needs and experience levels.""",
-            tools=["emotion_detection", "conversation"],
             model_config=model_config,
-            verbose=verbose
-        )
-        self.exercises = self._load_exercises()
-        self.current_session = None
-        
-    def _load_exercises(self) -> Dict:
-        """Load mindfulness exercises and meditation scripts"""
-        exercises_path = Path(__file__).parent.parent / "knowledge_base" / "resources" / "mindfulness_exercises.json"
-        try:
-            with open(exercises_path) as f:
-                return json.load(f)
-        except FileNotFoundError:
-            return {
-                "breathing": {
-                    "box_breathing": {
-                        "name": "Box Breathing",
-                        "description": "A simple technique to reduce stress and improve focus",
-                        "duration": 300,  # 5 minutes
-                        "steps": [
-                            "Inhale slowly for 4 counts",
-                            "Hold your breath for 4 counts",
-                            "Exhale slowly for 4 counts",
-                            "Hold for 4 counts before the next breath"
-                        ],
-                        "instructions": "Find a comfortable position. We'll practice box breathing for 5 minutes."
-                    },
-                    "deep_breathing": {
-                        "name": "Deep Breathing",
-                        "description": "Calming deep breathing exercise",
-                        "duration": 300,
-                        "steps": [
-                            "Take a deep breath in through your nose",
-                            "Feel your belly expand",
-                            "Exhale slowly through your mouth",
-                            "Feel your body relax"
-                        ],
-                        "instructions": "Sit comfortably with your back straight. Let's begin deep breathing."
-                    }
-                },
-                "meditation": {
-                    "body_scan": {
-                        "name": "Body Scan Meditation",
-                        "description": "Progressive relaxation through body awareness",
-                        "duration": 600,  # 10 minutes
-                        "steps": [
-                            "Focus on your toes and feet",
-                            "Move attention to your legs",
-                            "Progress through torso and arms",
-                            "End with neck and head"
-                        ],
-                        "instructions": "Lie down comfortably. We'll guide you through a full body scan."
-                    },
-                    "loving_kindness": {
-                        "name": "Loving-Kindness Meditation",
-                        "description": "Develop compassion for self and others",
-                        "duration": 600,
-                        "steps": [
-                            "Direct love to yourself",
-                            "Extend to loved ones",
-                            "Include neutral people",
-                            "Embrace all beings"
-                        ],
-                        "instructions": "Sit in a relaxed position. We'll practice sending loving-kindness."
-                    }
-                }
-            }
-            
-    def start_session(self, session_type: str, exercise_name: str = None) -> Dict:
-        """Start a new mindfulness session"""
-        if session_type not in self.exercises:
-            raise ValueError(f"Unknown session type: {session_type}")
-            
-        exercises = self.exercises[session_type]
-        if exercise_name and exercise_name not in exercises:
-            raise ValueError(f"Unknown exercise: {exercise_name}")
-            
-        selected_exercise = exercises[exercise_name] if exercise_name else next(iter(exercises.values()))
-        
-        self.current_session = {
-            "type": session_type,
-            "exercise": selected_exercise,
-            "start_time": self._get_timestamp(),
-            "current_step": 0,
-            "completed_steps": []
-        }
-        
-        return self.format_response(self._create_session_intro())
-        
-    def process_message(self, message: str) -> Dict:
-        """Process user message during mindfulness session"""
-        if not self.current_session:
-            return self.format_response(
-                "No active session. Please start a mindfulness session first."
-            )
-            
-        # Check for session control commands
-        if message.lower() in ["pause", "stop", "end"]:
-            return self.end_session()
-            
-        # Progress through exercise steps
-        return self._progress_session()
-        
-    def _progress_session(self) -> Dict:
-        """Progress through the current session's steps"""
-        exercise = self.current_session["exercise"]
-        current_step = self.current_session["current_step"]
-        
-        if current_step >= len(exercise["steps"]):
-            return self.end_session()
-            
-        step = exercise["steps"][current_step]
-        self.current_session["completed_steps"].append({
-            "step": current_step,
-            "timestamp": self._get_timestamp()
-        })
-        
-        self.current_session["current_step"] += 1
-        
-        return self.format_response(self._create_step_guidance(step))
-        
-    def _create_session_intro(self) -> str:
-        """Create introduction for the session"""
-        exercise = self.current_session["exercise"]
-        return f"""Welcome to {exercise['name']}.
-        
-        {exercise['description']}
-        
-        Duration: {exercise['duration'] // 60} minutes
-        
-        {exercise['instructions']}
-        
-        When you're ready, respond with 'begin' to start."""
-        
-    def _create_step_guidance(self, step: str) -> str:
-        """Create guidance for the current step"""
-        return f"""{step}
-        
-        Take your time with this step.
-        When you're ready to continue, send any message."""
-        
-    def end_session(self) -> Dict:
-        """End the current mindfulness session"""
-        if not self.current_session:
-            return self.format_response(
-                "No active session to end."
-            )
-            
-        # Calculate session statistics
-        stats = self._calculate_session_stats()
-        
-        # Generate session summary
-        summary = self._generate_session_summary(stats)
-        
-        # Add to history
-        self.add_to_history({
-            "session_type": self.current_session["type"],
-            "exercise": self.current_session["exercise"]["name"],
-            "stats": stats,
-            "summary": summary
-        })
-        
-        # Reset current session
-        self.current_session = None
-        
-        return self.format_response(summary)
-        
-    def _calculate_session_stats(self) -> Dict:
-        """Calculate statistics for the completed session"""
-        start_time = self.current_session["start_time"]
-        end_time = self._get_timestamp()
-        completed_steps = len(self.current_session["completed_steps"])
-        total_steps = len(self.current_session["exercise"]["steps"])
-        
-        return {
-            "duration": self._calculate_duration(start_time, end_time),
-            "completion_rate": completed_steps / total_steps,
-            "steps_completed": completed_steps
-        }
-        
-    def _calculate_duration(self, start_time: str, end_time: str) -> int:
-        """Calculate session duration in seconds"""
-        from datetime import datetime
-        start = datetime.fromisoformat(start_time)
-        end = datetime.fromisoformat(end_time)
-        return int((end - start).total_seconds())
-        
-    def _generate_session_summary(self, stats: Dict) -> str:
-        """Generate summary of the completed session"""
-        exercise = self.current_session["exercise"]
-        
-        return f"""Session Complete: {exercise['name']}
-        
-        Duration: {stats['duration'] // 60} minutes
-        Steps Completed: {stats['steps_completed']} of {len(exercise['steps'])}
-        Completion Rate: {stats['completion_rate'] * 100:.1f}%
-        
-        Thank you for practicing mindfulness. Remember to:
-        1. Take these peaceful feelings with you
-        2. Practice regularly for best results
-        3. Be gentle with yourself
-        
-        Would you like to try another exercise?"""
-        
-    def get_available_exercises(self) -> Dict:
-        """Get list of available exercises"""
-        return {
-            category: list(exercises.keys())
-            for category, exercises in self.exercises.items()
-        } 
+            agent_type="mindfulness",
+            role="Mindfulness Guide",
+            goal="Guide users through mindfulness exercises and meditation",
+            backstory="""I am an AI agent specialized in mindfulness and meditation. 
+            I provide calming guidance and support for various mindfulness practices 
+            while being attentive to individual needs and comfort levels.""",
+            verbose=True,
+            allow_delegation=False,
+            tools=[],  # Tools will be added as needed
+            **kwargs
+        ) 

agents/orchestrator.py

@@ -26,21 +26,32 @@ class WellnessOrchestrator:
         self.logger.info("Initializing agents")
         
         try:
+            # Common agent configuration
+            agent_config = {
+                "llm_config": self.model_config.get("conversation", {}),
+                "temperature": 0.7,
+                "max_iterations": 3
+            }
+            
             # Initialize each agent
             self.conversation_agent = ConversationAgent(
-                model_config=self.model_config
+                model_config=self.model_config,
+                **agent_config
             )
             
             self.assessment_agent = ConversationAgent(  # Temporarily using ConversationAgent
-                model_config=self.model_config
+                model_config=self.model_config,
+                **agent_config
             )
             
             self.mindfulness_agent = ConversationAgent(  # Temporarily using ConversationAgent
-                model_config=self.model_config
+                model_config=self.model_config,
+                **agent_config
             )
             
             self.crisis_agent = ConversationAgent(  # Temporarily using ConversationAgent
-                model_config=self.model_config
+                model_config=self.model_config,
+                **agent_config
             )
             
             self.logger.info("All agents initialized successfully")

config/config.py

@@ -6,18 +6,50 @@ def load_config() -> Dict:
     return {
         "MODEL_CONFIGS": {
             "conversation": {
-                "model_id": "facebook/opt-125m",
-                "max_length": 2048,
+                "model_id": "mistralai/Mistral-7B-Instruct-v0.2",
+                "max_length": 4096,
                 "temperature": 0.7,
+                "top_p": 0.95,
+                "repetition_penalty": 1.1,
+                "context_length": 4096,
+                "instruction_template": """You are a therapeutic conversation agent specialized in mental health support.
+                Your goal is to provide empathetic, supportive responses while maintaining appropriate boundaries.
+                
+                User: {input}
+                Assistant: """
             },
             "emotion_detection": {
-                "model_id": "j-hartmann/emotion-english-distilroberta-base",
+                "model_id": "mistralai/Mistral-7B-Instruct-v0.2",
+                "instruction_template": """Analyze the emotional content of the following text and identify the primary emotions present.
+                Focus on: joy, sadness, anger, fear, surprise, disgust, anxiety, and hope.
+                
+                Text: {input}
+                Emotions: """
             },
-            "speech_to_text": {
-                "model_id": "openai/whisper-tiny",
+            "assessment": {
+                "model_id": "mistralai/Mistral-7B-Instruct-v0.2",
+                "instruction_template": """Conduct a mental health assessment based on the user's input.
+                Consider: mood, anxiety levels, sleep patterns, social support, and risk factors.
+                Provide a structured assessment while maintaining a supportive tone.
+                
+                User Input: {input}
+                Assessment: """
             },
-            "vision": {
-                "model_id": "microsoft/resnet-18",
+            "crisis": {
+                "model_id": "mistralai/Mistral-7B-Instruct-v0.2",
+                "instruction_template": """You are a crisis intervention specialist. Assess the severity of the situation
+                and provide appropriate support and resources. Always prioritize user safety.
+                
+                User: {input}
+                Response: """
+            },
+            "mindfulness": {
+                "model_id": "mistralai/Mistral-7B-Instruct-v0.2",
+                "instruction_template": """Guide the user through mindfulness exercises and meditation.
+                Provide clear, calming instructions while being attentive to their needs.
+                
+                User: {input}
+                Guide: """
             }
         },
         "INTERFACE_CONFIG": {

requirements.txt

@@ -20,4 +20,7 @@ psutil>=5.9.0
 python-json-logger>=2.0.7
 accelerate>=0.26.0
 bitsandbytes>=0.41.1
-safetensors>=0.4.0 
+safetensors>=0.4.0
+sentencepiece>=0.1.99
+einops>=0.7.0
+optimum>=1.16.1