Spaces:
Runtime error
Runtime error
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 |
+
}
|