Spaces:
Sleeping
Sleeping
| import requests | |
| import json | |
| import time | |
| from typing import Dict, List, Optional, Tuple | |
| class GraniteAgent: | |
| def __init__(self, config): | |
| self.config = config | |
| self.base_url = config.OLLAMA_BASE_URL | |
| self.model = config.GRANITE_MODEL | |
| self.max_retries = 3 | |
| self.retry_delay = 2 # seconds | |
| def assign_wellness_tasks(self, user_profile: Dict, mental_health_assessment: str, risk_level: int) -> List[Dict]: | |
| """Assign wellness tasks based on mental health assessment using Granite AI""" | |
| # First attempt: Comprehensive AI generation | |
| tasks = self._generate_ai_tasks(user_profile, mental_health_assessment, risk_level) | |
| if tasks: | |
| return tasks | |
| # Second attempt: Simplified AI generation with retry | |
| tasks = self._generate_simplified_tasks(user_profile, risk_level) | |
| if tasks: | |
| return tasks | |
| # Final fallback: Basic AI generation with minimal context | |
| tasks = self._generate_basic_tasks(risk_level) | |
| if tasks: | |
| return tasks | |
| # Emergency fallback: Only if AI is completely unavailable | |
| print("WARNING: Granite AI completely unavailable, using emergency preset tasks") | |
| return self._get_emergency_preset_tasks(risk_level) | |
| def _generate_ai_tasks(self, user_profile: Dict, mental_health_assessment: str, risk_level: int) -> Optional[List[Dict]]: | |
| """Primary method: Generate comprehensive personalized tasks using Granite AI""" | |
| prompt = self._build_comprehensive_prompt(user_profile, mental_health_assessment, risk_level) | |
| for attempt in range(self.max_retries): | |
| try: | |
| response = self._call_granite_api(prompt, temperature=0.5) | |
| if response: | |
| tasks = self._extract_and_validate_tasks(response, risk_level) | |
| if tasks and len(tasks) >= 3: # Ensure we have enough tasks | |
| print(f"✓ Generated {len(tasks)} comprehensive AI tasks") | |
| return tasks | |
| except Exception as e: | |
| print(f"Attempt {attempt + 1} failed for comprehensive generation: {e}") | |
| if attempt < self.max_retries - 1: | |
| time.sleep(self.retry_delay) | |
| return None | |
| def _generate_simplified_tasks(self, user_profile: Dict, risk_level: int) -> Optional[List[Dict]]: | |
| """Fallback method: Generate simplified tasks with reduced context""" | |
| prompt = self._build_simplified_prompt(user_profile, risk_level) | |
| for attempt in range(self.max_retries): | |
| try: | |
| response = self._call_granite_api(prompt, temperature=0.4) | |
| if response: | |
| tasks = self._extract_and_validate_tasks(response, risk_level) | |
| if tasks and len(tasks) >= 2: | |
| print(f"✓ Generated {len(tasks)} simplified AI tasks") | |
| return tasks | |
| except Exception as e: | |
| print(f"Attempt {attempt + 1} failed for simplified generation: {e}") | |
| if attempt < self.max_retries - 1: | |
| time.sleep(self.retry_delay) | |
| return None | |
| def _generate_basic_tasks(self, risk_level: int) -> Optional[List[Dict]]: | |
| """Secondary fallback: Generate basic tasks with minimal context""" | |
| prompt = self._build_basic_prompt(risk_level) | |
| for attempt in range(self.max_retries): | |
| try: | |
| response = self._call_granite_api(prompt, temperature=0.3) | |
| if response: | |
| tasks = self._extract_and_validate_tasks(response, risk_level, min_tasks=1) | |
| if tasks: | |
| print(f"✓ Generated {len(tasks)} basic AI tasks") | |
| return tasks | |
| except Exception as e: | |
| print(f"Attempt {attempt + 1} failed for basic generation: {e}") | |
| if attempt < self.max_retries - 1: | |
| time.sleep(self.retry_delay) | |
| return None | |
| def _build_comprehensive_prompt(self, user_profile: Dict, mental_health_assessment: str, risk_level: int) -> str: | |
| """Build detailed prompt for comprehensive task generation""" | |
| return f""" | |
| You are a professional wellness coach. Create personalized wellness tasks based on user data. | |
| GUIDELINES: | |
| - Use professional, clear language | |
| - Focus on evidence-based wellness practices | |
| - Ensure all recommendations are safe and appropriate | |
| - Use inclusive, respectful language | |
| USER CONTEXT: | |
| Risk Level: {risk_level}/10 (10 = highest risk) | |
| Assessment: {mental_health_assessment} | |
| USER PROFILE: | |
| - Stress Level: {user_profile.get('Stress_Level', 'Not specified')} | |
| - Sleep: {user_profile.get('Sleep_Hours', 'Not specified')} hours/night, quality: {user_profile.get('Sleep_Quality', 'Not specified')} | |
| - Work: {user_profile.get('Work_Hours', 'Not specified')} hours/week | |
| - Exercise: {user_profile.get('Physical_Activity_Hours', 'Not specified')} hours/week | |
| - Occupation: {user_profile.get('Occupation', 'Not specified')} | |
| - Age: {user_profile.get('Age', 'Not specified')} | |
| - Mood: {user_profile.get('Mood', 'Not specified')} | |
| - Anxiety: {user_profile.get('Anxiety_Frequency', 'Not specified')} | |
| - Energy: {user_profile.get('Energy_Level', 'Not specified')} | |
| TASK REQUIREMENTS: | |
| {self._get_risk_specific_requirements(risk_level)} | |
| Generate 4-6 personalized wellness tasks. Each task must be: | |
| 1. Relevant to their specific situation | |
| 2. Practical and achievable | |
| 3. Evidence-based for mental health improvement | |
| 4. Appropriate for their risk level | |
| 5. Include clear, actionable instructions | |
| MANDATORY JSON FORMAT (return ONLY valid JSON): | |
| [ | |
| {{ | |
| "task_type": "select from: meditation, exercise, sleep_schedule, social_connection, journaling, breathing_exercise, nature_walk, healthy_meal, screen_break, gratitude_practice, professional_help, mindfulness, stress_management, routine_building, creative_activity, relaxation_technique", | |
| "title": "Clear, engaging title", | |
| "description": "Brief description explaining the benefits", | |
| "duration_days": appropriate_number, | |
| "difficulty": "easy/medium/hard", | |
| "instructions": "Step-by-step instructions tailored to their profile", | |
| "completion_criteria": "Clear, measurable success criteria", | |
| "personalization_notes": "Why this task fits their specific situation" | |
| }} | |
| ] | |
| """ | |
| def _build_simplified_prompt(self, user_profile: Dict, risk_level: int) -> str: | |
| """Build simplified prompt when comprehensive generation fails""" | |
| key_factors = [] | |
| if user_profile.get('Stress_Level'): | |
| key_factors.append(f"Stress: {user_profile['Stress_Level']}") | |
| if user_profile.get('Sleep_Hours'): | |
| key_factors.append(f"Sleep: {user_profile['Sleep_Hours']}h") | |
| if user_profile.get('Work_Hours'): | |
| key_factors.append(f"Work: {user_profile['Work_Hours']}h/week") | |
| factors_str = ", ".join(key_factors) if key_factors else "Limited profile data" | |
| return f""" | |
| You are a wellness coach. Create practical wellness tasks. | |
| USER: {factors_str} | |
| Risk Level: {risk_level}/10 | |
| {self._get_risk_specific_requirements(risk_level)} | |
| Generate 3-4 practical wellness tasks as JSON array: | |
| [ | |
| {{ | |
| "task_type": "meditation|exercise|sleep_schedule|journaling|breathing_exercise|professional_help|stress_management", | |
| "title": "Clear task title", | |
| "description": "Brief helpful description", | |
| "duration_days": number, | |
| "difficulty": "easy|medium|hard", | |
| "instructions": "Step-by-step instructions", | |
| "completion_criteria": "How to measure success" | |
| }} | |
| ] | |
| """ | |
| def _build_basic_prompt(self, risk_level: int) -> str: | |
| """Build minimal prompt for basic task generation""" | |
| return f""" | |
| You are a wellness coach. Generate wellness tasks for risk level {risk_level}/10. | |
| {self._get_risk_specific_requirements(risk_level)} | |
| Return 2-3 tasks as JSON: | |
| [ | |
| {{ | |
| "task_type": "breathing_exercise|meditation|professional_help|journaling", | |
| "title": "Task title", | |
| "description": "What this helps with", | |
| "duration_days": 1-7, | |
| "difficulty": "easy|medium", | |
| "instructions": "Clear instructions", | |
| "completion_criteria": "Success measure" | |
| }} | |
| ] | |
| """ | |
| def _get_risk_specific_requirements(self, risk_level: int) -> str: | |
| """Get specific requirements based on risk level""" | |
| if risk_level >= 8: | |
| return """ | |
| CRITICAL PRIORITY: | |
| - MUST include immediate professional help seeking | |
| - Focus on crisis intervention and safety | |
| - Include emergency resources and contacts | |
| - Tasks should provide immediate coping mechanisms | |
| - Maximum task duration: 1-2 days | |
| """ | |
| elif risk_level >= 6: | |
| return """ | |
| HIGH PRIORITY: | |
| - Strongly recommend professional consultation within 1 week | |
| - Include daily anxiety/stress management techniques | |
| - Focus on stabilization and routine building | |
| - Provide structured, manageable activities | |
| - Task duration: 2-7 days | |
| """ | |
| elif risk_level >= 4: | |
| return """ | |
| MODERATE PRIORITY: | |
| - Include both self-care and gradual improvement activities | |
| - Balance mental and physical wellness approaches | |
| - Encourage social connection and support | |
| - Build sustainable, healthy habits | |
| - Task duration: 5-14 days | |
| """ | |
| else: | |
| return """ | |
| MAINTENANCE/PREVENTION: | |
| - Focus on wellness enhancement and prevention | |
| - Include enjoyable, engaging activities | |
| - Support long-term habit building | |
| - Promote overall life satisfaction | |
| - Task duration: 7-21 days | |
| """ | |
| def _extract_and_validate_tasks(self, response: str, risk_level: int, min_tasks: int = 2) -> Optional[List[Dict]]: | |
| """Extract and validate tasks from AI response""" | |
| try: | |
| # Clean the response | |
| cleaned_response = self._clean_json_string(response) | |
| if not cleaned_response: | |
| return None | |
| # Try to parse JSON | |
| try: | |
| tasks = json.loads(cleaned_response) | |
| except json.JSONDecodeError: | |
| # Try to extract individual task objects | |
| tasks = self._extract_individual_tasks(response) | |
| if not tasks: | |
| return None | |
| # Validate task structure | |
| validated_tasks = self._validate_task_structure(tasks, risk_level) | |
| if validated_tasks and len(validated_tasks) >= min_tasks: | |
| return validated_tasks | |
| except Exception as e: | |
| print(f"Error extracting tasks: {e}") | |
| return None | |
| def _clean_json_string(self, json_str: str) -> Optional[str]: | |
| """Clean malformed JSON string""" | |
| try: | |
| # Remove common issues | |
| cleaned = json_str.strip() | |
| # Find JSON array boundaries | |
| start_idx = cleaned.find('[') | |
| end_idx = cleaned.rfind(']') | |
| if start_idx == -1 or end_idx == -1: | |
| return None | |
| cleaned = cleaned[start_idx:end_idx + 1] | |
| # Remove newlines and extra whitespace | |
| cleaned = cleaned.replace('\n', ' ').replace('\r', '') | |
| cleaned = ' '.join(cleaned.split()) | |
| # Fix common JSON issues | |
| cleaned = cleaned.replace(',}', '}') # Remove trailing commas | |
| cleaned = cleaned.replace(',]', ']') | |
| return cleaned | |
| except Exception: | |
| return None | |
| def _extract_individual_tasks(self, response: str) -> Optional[List[Dict]]: | |
| """Extract individual task objects when array parsing fails""" | |
| tasks = [] | |
| # Find individual task objects | |
| import re | |
| task_pattern = r'\{[^{}]*"task_type"[^{}]*\}' | |
| matches = re.findall(task_pattern, response, re.DOTALL) | |
| for match in matches: | |
| try: | |
| # Clean the match | |
| clean_match = match.replace('\n', ' ').replace('\r', '') | |
| clean_match = ' '.join(clean_match.split()) | |
| clean_match = clean_match.replace(',}', '}') | |
| task = json.loads(clean_match) | |
| tasks.append(task) | |
| except: | |
| continue | |
| return tasks if tasks else None | |
| def _validate_task_structure(self, tasks: List[Dict], risk_level: int) -> Optional[List[Dict]]: | |
| """Validate and sanitize task structure""" | |
| if not isinstance(tasks, list): | |
| return None | |
| valid_tasks = [] | |
| required_fields = ['task_type', 'title', 'description', 'duration_days', 'difficulty', 'instructions', 'completion_criteria'] | |
| valid_task_types = { | |
| 'meditation', 'exercise', 'sleep_schedule', 'social_connection', | |
| 'journaling', 'breathing_exercise', 'nature_walk', 'healthy_meal', | |
| 'screen_break', 'gratitude_practice', 'professional_help', | |
| 'mindfulness', 'stress_management', 'routine_building', | |
| 'creative_activity', 'relaxation_technique' | |
| } | |
| # Ensure high-risk users get professional help | |
| has_professional_help = False | |
| for task in tasks: | |
| if not isinstance(task, dict): | |
| continue | |
| # Check required fields | |
| if not all(field in task and str(task[field]).strip() for field in required_fields): | |
| continue | |
| # Validate and clean task | |
| validated_task = self._clean_and_validate_task(task, valid_task_types) | |
| if validated_task: | |
| valid_tasks.append(validated_task) | |
| if validated_task['task_type'] == 'professional_help': | |
| has_professional_help = True | |
| # Ensure high-risk users have professional help task | |
| if risk_level >= 7 and not has_professional_help and len(valid_tasks) > 0: | |
| professional_task = self._generate_professional_help_task(risk_level) | |
| valid_tasks.insert(0, professional_task) | |
| return valid_tasks if len(valid_tasks) >= 1 else None | |
| def _clean_and_validate_task(self, task: Dict, valid_task_types: set) -> Optional[Dict]: | |
| """Clean and validate individual task""" | |
| try: | |
| # Validate task_type | |
| if task.get('task_type') not in valid_task_types: | |
| return None | |
| # Clean and validate duration | |
| try: | |
| duration = int(task.get('duration_days', 7)) | |
| duration = max(1, min(30, duration)) # Clamp between 1-30 days | |
| except (ValueError, TypeError): | |
| duration = 7 | |
| # Validate difficulty | |
| difficulty = task.get('difficulty', 'medium').lower() | |
| if difficulty not in ['easy', 'medium', 'hard']: | |
| difficulty = 'medium' | |
| # Clean text fields | |
| cleaned_task = { | |
| 'task_type': task['task_type'], | |
| 'title': str(task['title']).strip()[:100], | |
| 'description': str(task['description']).strip()[:300], | |
| 'duration_days': duration, | |
| 'difficulty': difficulty, | |
| 'instructions': str(task['instructions']).strip()[:1000], | |
| 'completion_criteria': str(task['completion_criteria']).strip()[:200] | |
| } | |
| # Add optional fields if present | |
| if 'personalization_notes' in task: | |
| cleaned_task['personalization_notes'] = str(task['personalization_notes']).strip()[:200] | |
| return cleaned_task | |
| except Exception as e: | |
| print(f"Error cleaning task: {e}") | |
| return None | |
| def _generate_professional_help_task(self, risk_level: int) -> Dict: | |
| """Generate professional help task for high-risk users""" | |
| urgency = "immediately" if risk_level >= 8 else "within 1-2 days" | |
| return { | |
| "task_type": "professional_help", | |
| "title": "Seek Professional Mental Health Support", | |
| "description": f"Contact a mental health professional {urgency} for assessment and support", | |
| "duration_days": 2, | |
| "difficulty": "medium", | |
| "instructions": f"Contact your healthcare provider, call a mental health helpline, or visit a mental health clinic {urgency}. If in immediate crisis, call emergency services (911) or go to the nearest emergency room.", | |
| "completion_criteria": "Make contact with a mental health professional or crisis support service" | |
| } | |
| def _get_emergency_preset_tasks(self, risk_level: int) -> List[Dict]: | |
| """ONLY used when Granite AI is completely unavailable - minimal preset tasks""" | |
| print("EMERGENCY: Using preset tasks - Granite AI unavailable") | |
| if risk_level >= 7: | |
| return [ | |
| { | |
| "task_type": "professional_help", | |
| "title": "Emergency Professional Help", | |
| "description": "Seek immediate professional mental health support", | |
| "duration_days": 1, | |
| "difficulty": "medium", | |
| "instructions": "Contact emergency mental health services, your doctor, or call a crisis helpline immediately.", | |
| "completion_criteria": "Make contact with professional help" | |
| } | |
| ] | |
| elif risk_level >= 4: | |
| return [ | |
| { | |
| "task_type": "breathing_exercise", | |
| "title": "Daily Breathing Practice", | |
| "description": "Use breathing exercises to manage stress and anxiety", | |
| "duration_days": 7, | |
| "difficulty": "easy", | |
| "instructions": "Practice 4-7-8 breathing: Inhale for 4 counts, hold for 7 counts, exhale for 8 counts. Repeat 4 times, twice daily.", | |
| "completion_criteria": "Complete breathing exercise twice daily for one week" | |
| }, | |
| { | |
| "task_type": "sleep_schedule", | |
| "title": "Improve Sleep Routine", | |
| "description": "Establish a consistent sleep schedule for better rest", | |
| "duration_days": 14, | |
| "difficulty": "medium", | |
| "instructions": "Go to bed and wake up at the same time daily. Create a 30-minute wind-down routine before bed.", | |
| "completion_criteria": "Maintain consistent sleep schedule for 2 weeks" | |
| } | |
| ] | |
| else: | |
| return [ | |
| { | |
| "task_type": "gratitude_practice", | |
| "title": "Daily Gratitude Journal", | |
| "description": "Practice gratitude to boost mood and well-being", | |
| "duration_days": 14, | |
| "difficulty": "easy", | |
| "instructions": "Write down 3 things you're grateful for each morning. Be specific and reflect on why you appreciate them.", | |
| "completion_criteria": "Complete gratitude entries for 14 consecutive days" | |
| }, | |
| { | |
| "task_type": "nature_walk", | |
| "title": "Weekly Nature Walks", | |
| "description": "Connect with nature to reduce stress and improve mood", | |
| "duration_days": 21, | |
| "difficulty": "easy", | |
| "instructions": "Take a 20-30 minute walk in a park, garden, or natural area twice per week. Focus on your surroundings.", | |
| "completion_criteria": "Complete 6 nature walks over 3 weeks" | |
| } | |
| ] | |
| def _call_granite_api(self, prompt: str, temperature: float = 0.3) -> Optional[str]: | |
| """Enhanced API call with better error handling and logging""" | |
| data = { | |
| "model": self.model, | |
| "prompt": prompt, | |
| "stream": False, | |
| "options": { | |
| "temperature": temperature, | |
| "top_p": 0.9, | |
| "num_predict": 2000, # Increased for better task generation | |
| "repeat_penalty": 1.1, | |
| "stop": ["Human:", "Assistant:", "\n\n---"] | |
| } | |
| } | |
| try: | |
| print(f"🤖 Calling Granite API (temp={temperature})...") | |
| response = requests.post(self.base_url, json=data, timeout=90) | |
| if response.status_code == 200: | |
| result = response.json() | |
| ai_response = result.get('response', '').strip() | |
| if ai_response: | |
| print(f"✓ Granite API responded ({len(ai_response)} chars)") | |
| return ai_response | |
| else: | |
| print("⚠ Granite API returned empty response") | |
| return None | |
| else: | |
| print(f"❌ Granite API error: {response.status_code}") | |
| return None | |
| except requests.exceptions.ConnectionError: | |
| print("❌ Cannot connect to Granite (Ollama not running?)") | |
| return None | |
| except requests.exceptions.Timeout: | |
| print("❌ Granite API timeout") | |
| return None | |
| except Exception as e: | |
| print(f"❌ Granite API error: {e}") | |
| return None |