invincible-jha commited on
Commit
4464dfa
·
1 Parent(s): f0b146e

Update agent initialization for CrewAI compatibility

Browse files
mentalwellness_space/agents/conversation_agent.py ADDED
@@ -0,0 +1,96 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List
2
+ from .base_agent import BaseWellnessAgent
3
+ from transformers import pipeline
4
+
5
+ class ConversationAgent(BaseWellnessAgent):
6
+ """Agent specialized in therapeutic conversations"""
7
+
8
+ def __init__(self, model_config: Dict, verbose: bool = False):
9
+ super().__init__(
10
+ name="Therapeutic Conversation Agent",
11
+ role="Mental Health Conversationalist",
12
+ goal="Engage in supportive therapeutic conversations while maintaining empathy and professionalism",
13
+ backstory="""I am a specialized AI agent trained in therapeutic conversation techniques.
14
+ I use evidence-based approaches to provide emotional support and guidance while maintaining
15
+ appropriate boundaries and recognizing when to escalate to crisis intervention.""",
16
+ tools=["emotion_detection", "conversation"],
17
+ model_config=model_config,
18
+ verbose=verbose
19
+ )
20
+ self._initialize_conversation_model()
21
+
22
+ def _initialize_conversation_model(self):
23
+ """Initialize the conversation model"""
24
+ self.models["conversation"] = pipeline(
25
+ "text-generation",
26
+ model=self.model_config["conversation"]["model_id"]
27
+ )
28
+
29
+ def _generate_response(self, prompt: str) -> str:
30
+ """Generate response using the conversation model"""
31
+ response = self.models["conversation"](
32
+ prompt,
33
+ max_length=self.model_config["conversation"]["max_length"],
34
+ temperature=self.model_config["conversation"]["temperature"]
35
+ )
36
+ return response[0]["generated_text"] if response else ""
37
+
38
+ def process_message(self, message: str) -> Dict:
39
+ """Process user message and generate therapeutic response"""
40
+ # Analyze emotion
41
+ emotion = self.analyze_emotion(message)
42
+
43
+ # Update context with emotion
44
+ self.update_context({"last_emotion": emotion})
45
+
46
+ # Generate appropriate response based on emotion and context
47
+ prompt = self._create_prompt(message, emotion)
48
+ response = self._generate_response(prompt)
49
+
50
+ # Add to conversation history
51
+ interaction = {
52
+ "user_message": message,
53
+ "emotion": emotion,
54
+ "response": response
55
+ }
56
+ self.add_to_history(interaction)
57
+
58
+ return self.format_response(response)
59
+
60
+ def _create_prompt(self, message: str, emotion: Dict) -> str:
61
+ """Create context-aware prompt for response generation"""
62
+ history = self.get_history()
63
+ context = self.get_context()
64
+
65
+ # Build prompt with context
66
+ prompt = f"""As a therapeutic conversation agent, respond to the following message with empathy and support.
67
+ User's message: {message}
68
+ Detected emotion: {emotion}
69
+ Conversation history: {history[-3:] if history else 'No history'}
70
+ Additional context: {context}
71
+
72
+ Response:"""
73
+
74
+ return prompt
75
+
76
+ def get_conversation_summary(self) -> Dict:
77
+ """Generate summary of the conversation"""
78
+ history = self.get_history()
79
+ emotions = [interaction["emotion"] for interaction in history]
80
+
81
+ return {
82
+ "interaction_count": len(history),
83
+ "emotional_trajectory": emotions,
84
+ "key_topics": self._extract_key_topics(history),
85
+ "recommendations": self._generate_recommendations(history)
86
+ }
87
+
88
+ def _extract_key_topics(self, history: List[Dict]) -> List[str]:
89
+ """Extract main topics from conversation history"""
90
+ # Implement topic extraction logic
91
+ return []
92
+
93
+ def _generate_recommendations(self, history: List[Dict]) -> List[str]:
94
+ """Generate recommendations based on conversation"""
95
+ # Implement recommendation logic
96
+ return []
mentalwellness_space/agents/orchestrator.py ADDED
@@ -0,0 +1,195 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List
2
+ from crewai import Crew, Process, Task
3
+ from agents.conversation_agent import ConversationAgent
4
+ from agents.assessment_agent import AssessmentAgent
5
+ from agents.mindfulness_agent import MindfulnessAgent
6
+ from agents.crisis_agent import CrisisAgent
7
+ import logging
8
+ from utils.log_manager import LogManager
9
+
10
+ class WellnessOrchestrator:
11
+ """Orchestrates the coordination between different agents"""
12
+
13
+ def __init__(self, model_config: Dict):
14
+ self.model_config = model_config
15
+ self.log_manager = LogManager()
16
+ self.logger = self.log_manager.get_agent_logger("orchestrator")
17
+
18
+ # Initialize agents
19
+ self.initialize_agents()
20
+
21
+ # Initialize CrewAI
22
+ self.initialize_crew()
23
+
24
+ def initialize_agents(self):
25
+ """Initialize all agents with their specific roles and tools"""
26
+ self.logger.info("Initializing agents")
27
+
28
+ try:
29
+ # Conversation Agent
30
+ self.conversation_agent = ConversationAgent(
31
+ name="Therapeutic Conversation Agent",
32
+ role="Lead conversation therapist",
33
+ goal="Guide therapeutic conversations and provide emotional support",
34
+ backstory="Expert in therapeutic dialogue and emotional support",
35
+ tools=["chat", "emotion_detection"],
36
+ model_config=self.model_config
37
+ )
38
+
39
+ # Assessment Agent
40
+ self.assessment_agent = AssessmentAgent(
41
+ name="Mental Health Assessment Agent",
42
+ role="Mental health evaluator",
43
+ goal="Conduct mental health assessments and track progress",
44
+ backstory="Specialist in mental health evaluation and monitoring",
45
+ tools=["assessment_tools", "progress_tracking"],
46
+ model_config=self.model_config
47
+ )
48
+
49
+ # Mindfulness Agent
50
+ self.mindfulness_agent = MindfulnessAgent(
51
+ name="Mindfulness Guide Agent",
52
+ role="Mindfulness and meditation instructor",
53
+ goal="Guide mindfulness exercises and meditation sessions",
54
+ backstory="Expert in mindfulness techniques and meditation",
55
+ tools=["meditation_guide", "breathing_exercises"],
56
+ model_config=self.model_config
57
+ )
58
+
59
+ # Crisis Agent
60
+ self.crisis_agent = CrisisAgent(
61
+ name="Crisis Intervention Agent",
62
+ role="Emergency response specialist",
63
+ goal="Provide immediate support in crisis situations",
64
+ backstory="Trained in crisis intervention and emergency response",
65
+ tools=["crisis_protocol", "emergency_resources"],
66
+ model_config=self.model_config
67
+ )
68
+
69
+ self.logger.info("All agents initialized successfully")
70
+
71
+ except Exception as e:
72
+ self.logger.error(f"Error initializing agents: {str(e)}")
73
+ raise
74
+
75
+ def initialize_crew(self):
76
+ """Initialize CrewAI with agents and tasks"""
77
+ self.logger.info("Initializing CrewAI")
78
+
79
+ try:
80
+ # Create the crew
81
+ self.crew = Crew(
82
+ agents=[
83
+ self.conversation_agent,
84
+ self.assessment_agent,
85
+ self.mindfulness_agent,
86
+ self.crisis_agent
87
+ ],
88
+ tasks=[], # Tasks will be added dynamically
89
+ process=Process.sequential # Can be changed to parallel if needed
90
+ )
91
+
92
+ self.logger.info("CrewAI initialized successfully")
93
+
94
+ except Exception as e:
95
+ self.logger.error(f"Error initializing CrewAI: {str(e)}")
96
+ raise
97
+
98
+ def create_task(self, task_type: str, description: str, agent) -> Task:
99
+ """Create a task for an agent"""
100
+ return Task(
101
+ description=description,
102
+ agent=agent,
103
+ expected_output="Detailed response with next steps"
104
+ )
105
+
106
+ def process_message(self, message: str, context: Dict = None) -> Dict:
107
+ """Process user message through appropriate agents"""
108
+ self.logger.info("Processing message through agents")
109
+
110
+ try:
111
+ # Clear previous tasks
112
+ self.crew.tasks = []
113
+
114
+ # Initial assessment by conversation agent
115
+ initial_task = self.create_task(
116
+ "initial_assessment",
117
+ f"Analyze this message and determine required support: {message}",
118
+ self.conversation_agent
119
+ )
120
+ self.crew.tasks.append(initial_task)
121
+
122
+ # Analyze for crisis indicators
123
+ crisis_check = self.create_task(
124
+ "crisis_check",
125
+ f"Check for crisis indicators in: {message}",
126
+ self.crisis_agent
127
+ )
128
+ self.crew.tasks.append(crisis_check)
129
+
130
+ # Execute the crew tasks
131
+ result = self.crew.kickoff()
132
+
133
+ # Process results and determine next steps
134
+ if "crisis" in result.lower():
135
+ # Add crisis intervention task
136
+ crisis_task = self.create_task(
137
+ "crisis_intervention",
138
+ f"Provide crisis intervention for: {message}",
139
+ self.crisis_agent
140
+ )
141
+ self.crew.tasks = [crisis_task]
142
+ response = self.crew.kickoff()
143
+
144
+ elif "assessment" in result.lower():
145
+ # Add assessment task
146
+ assessment_task = self.create_task(
147
+ "mental_health_assessment",
148
+ f"Conduct mental health assessment based on: {message}",
149
+ self.assessment_agent
150
+ )
151
+ self.crew.tasks = [assessment_task]
152
+ response = self.crew.kickoff()
153
+
154
+ elif "mindfulness" in result.lower():
155
+ # Add mindfulness task
156
+ mindfulness_task = self.create_task(
157
+ "mindfulness_session",
158
+ f"Guide mindfulness exercise based on: {message}",
159
+ self.mindfulness_agent
160
+ )
161
+ self.crew.tasks = [mindfulness_task]
162
+ response = self.crew.kickoff()
163
+
164
+ else:
165
+ # Continue therapeutic conversation
166
+ conversation_task = self.create_task(
167
+ "therapeutic_conversation",
168
+ f"Continue therapeutic conversation: {message}",
169
+ self.conversation_agent
170
+ )
171
+ self.crew.tasks = [conversation_task]
172
+ response = self.crew.kickoff()
173
+
174
+ return {
175
+ "message": response,
176
+ "agent_type": self.crew.tasks[-1].agent.name,
177
+ "task_type": self.crew.tasks[-1].task_type
178
+ }
179
+
180
+ except Exception as e:
181
+ self.logger.error(f"Error processing message: {str(e)}")
182
+ return {
183
+ "message": "I apologize, but I encountered an error. Please try again.",
184
+ "agent_type": "error",
185
+ "task_type": "error_handling"
186
+ }
187
+
188
+ def get_agent_status(self) -> Dict:
189
+ """Get status of all agents"""
190
+ return {
191
+ "conversation": self.conversation_agent.get_status(),
192
+ "assessment": self.assessment_agent.get_status(),
193
+ "mindfulness": self.mindfulness_agent.get_status(),
194
+ "crisis": self.crisis_agent.get_status()
195
+ }