mindfull_pipeline.py

#!/usr/bin/env python3
"""
Mindfull AI Avatar Chatbot - Main Pipeline
Integrates Ollama (Mistral), F5-TTS, and SadTalker for complete AI avatar experience
"""

import os
import sys
import json
import time
import logging
import asyncio
import subprocess
import threading
from pathlib import Path
from typing import Dict, Any, Optional, Tuple
from dataclasses import dataclass
from datetime import datetime

import requests
import torch
import numpy as np
import soundfile as sf
from PIL import Image

# Import configuration
from mindfull_config import config, logger

@dataclass
class ChatSession:
    """Represents a chat session with the user"""
    session_id: str
    user_name: str = ""
    user_rank: str = ""
    start_time: Optional[datetime] = None
    last_interaction: Optional[datetime] = None
    conversation_history: Optional[list] = None
    
    def __post_init__(self):
        if self.start_time is None:
            self.start_time = datetime.now()
        if self.last_interaction is None:
            self.last_interaction = datetime.now()
        if self.conversation_history is None:
            self.conversation_history = []

class MindfullPipeline:
    """Main pipeline for Mindfull AI Avatar Chatbot"""
    
    def __init__(self):
        self.session = None
        self.is_initialized = False
        self.ollama_available = False
        self.f5tts_available = False
        self.sadtalker_available = False
        
        logger.info("Initializing Mindfull AI Avatar Chatbot...")
        self.initialize_components()
    
    def initialize_components(self) -> bool:
        """Initialize all pipeline components"""
        try:
            # Check Ollama availability
            self.ollama_available = self._check_ollama()
            
            # Check F5-TTS model availability
            self.f5tts_available = self._check_f5tts()
            
            # Check SadTalker availability
            self.sadtalker_available = self._check_sadtalker()
            
            # Validate paths
            path_validation = config.validate_paths()
            
            self.is_initialized = (
                self.ollama_available and 
                self.f5tts_available and 
                path_validation
            )
            
            if self.is_initialized:
                logger.info("✅ Mindfull pipeline initialized successfully!")
                self._log_component_status()
            else:
                logger.error("❌ Failed to initialize Mindfull pipeline")
                self._log_component_status()
            
            return self.is_initialized
            
        except Exception as e:
            logger.error(f"Error during initialization: {str(e)}")
            return False
    
    def _check_ollama(self) -> bool:
        """Check if Ollama is running and model is available"""
        try:
            # Check if Ollama is running
            response = requests.get(config.OLLAMA_MODELS_URL, timeout=5)
            if response.status_code != 200:
                logger.error("Ollama server is not responding")
                return False
            
            # Check if Mindfull model exists
            models = response.json().get('models', [])
            model_names = [model.get('name', '') for model in models]
            
            if config.DEFAULT_MODEL in model_names:
                logger.info(f"✅ Found {config.DEFAULT_MODEL} model")
                return True
            elif config.FALLBACK_MODEL in model_names:
                logger.warning(f"⚠️ Using fallback model: {config.FALLBACK_MODEL}")
                return True
            else:
                logger.error(f"❌ Neither {config.DEFAULT_MODEL} nor {config.FALLBACK_MODEL} found")
                logger.info(f"Available models: {model_names}")
                return False
                
        except requests.exceptions.RequestException as e:
            logger.error(f"Failed to connect to Ollama: {str(e)}")
            return False
    
    def _check_f5tts(self) -> bool:
        """Check if F5-TTS model files are available"""
        required_files = [
            config.F5TTS_MODEL_PATH,
            config.F5TTS_CONFIG_PATH,
            config.DEFAULT_REFERENCE_AUDIO
        ]
        
        missing_files = [f for f in required_files if not f.exists()]
        
        if missing_files:
            logger.error(f"❌ Missing F5-TTS files: {missing_files}")
            return False
        
        logger.info("✅ F5-TTS model files found")
        return True
    
    def _check_sadtalker(self) -> bool:
        """Check if SadTalker is available"""
        if not config.SADTALKER_SCRIPT.exists():
            logger.warning("⚠️ SadTalker script not found - video generation disabled")
            return False
        
        if not config.DEFAULT_AVATAR_IMAGE.exists():
            logger.warning("⚠️ Default avatar image not found")
            return False
            
        logger.info("✅ SadTalker components found")
        return True
    
    def _log_component_status(self):
        """Log the status of all components"""
        logger.info("Component Status:")
        logger.info(f"  Ollama: {'✅' if self.ollama_available else '❌'}")
        logger.info(f"  F5-TTS: {'✅' if self.f5tts_available else '❌'}")
        logger.info(f"  SadTalker: {'✅' if self.sadtalker_available else '❌'}")
    
    def create_session(self) -> str:
        """Create a new chat session"""
        session_id = f"mindfull_session_{int(time.time())}"
        self.session = ChatSession(session_id=session_id)
        logger.info(f"Created new session: {session_id}")
        return session_id
    
    def generate_response(self, user_input: str) -> Tuple[str, str]:
        """
        Generate response using Ollama
        Returns: (response_text, emotion_detected)
        """
        try:
            # Detect emotion from user input
            emotion = self._detect_emotion(user_input)
            
            # Prepare the prompt with context
            prompt = self._prepare_prompt(user_input, emotion)
            
            # Make request to Ollama
            payload = {
                "model": config.DEFAULT_MODEL,
                "prompt": prompt,
                "stream": False,
                **config.MODEL_PARAMS
            }
            
            response = requests.post(
                config.OLLAMA_GENERATE_URL,
                json=payload,
                timeout=config.OLLAMA_TIMEOUT
            )
            
            if response.status_code == 200:
                result = response.json()
                response_text = result.get('response', '').strip()
                
                # Add to conversation history
                if self.session and self.session.conversation_history is not None:
                    self.session.conversation_history.append({
                        'timestamp': datetime.now().isoformat(),
                        'user_input': user_input,
                        'response': response_text,
                        'emotion': emotion
                    })
                    self.session.last_interaction = datetime.now()
                
                logger.info(f"Generated response for emotion '{emotion}': {response_text[:100]}...")
                return response_text, emotion
            else:
                logger.error(f"Ollama request failed: {response.status_code}")
                return config.FALLBACK_MESSAGES["ollama_error"], "neutral"
                
        except Exception as e:
            logger.error(f"Error generating response: {str(e)}")
            return config.FALLBACK_MESSAGES["general_error"], "neutral"
    
    def _detect_emotion(self, text: str) -> str:
        """Detect emotion from user input using keyword matching"""
        text_lower = text.lower()
        
        emotion_scores = {}
        for emotion, keywords in config.EMOTION_KEYWORDS.items():
            score = sum(1 for keyword in keywords if keyword in text_lower)
            if score > 0:
                emotion_scores[emotion] = score
        
        if not emotion_scores:
            return "neutral"
        
        # Return emotion with highest score
        if emotion_scores:
            detected_emotion = max(emotion_scores, key=lambda x: emotion_scores[x])
            logger.info(f"Detected emotion: {detected_emotion} (scores: {emotion_scores})")
            return detected_emotion
    
    def _prepare_prompt(self, user_input: str, emotion: str) -> str:
        """Prepare the prompt for Ollama with context"""
        # Get wellness suggestions for the detected emotion
        suggestions = config.WELLNESS_INTERVENTIONS.get(emotion, 
                      config.WELLNESS_INTERVENTIONS["general"])
        
        context = f"""
User's emotional state: {emotion}
Available wellness suggestions: {', '.join(suggestions[:3])}

User message: {user_input}

Respond as Mindfull, keeping in mind:
1. The user's emotional state is {emotion}
2. Provide appropriate support and suggestions
3. Keep response concise (2-3 sentences) for voice output
4. Be warm and professional
"""
        return context
    
    def generate_audio(self, text: str) -> Optional[str]:
        """
        Generate audio using F5-TTS with fallback methods
        Returns: path to generated audio file or None if failed
        """
        try:
            # Import our simple audio generation
            from simple_audio_gen import generate_audio_simple, generate_audio_minimal

            # Prepare output filename
            timestamp = int(time.time())
            audio_filename = f"mindfull_response_{timestamp}.wav"
            audio_path = config.AUDIO_OUTPUT_DIR / audio_filename

            logger.info(f"Generating audio with F5-TTS: {text[:50]}...")

            # Try F5-TTS, then minimal fallback
            if generate_audio_simple(text, str(audio_path)):
                logger.info(f"✅ F5-TTS audio generated: {audio_path}")
                return str(audio_path)
            elif generate_audio_minimal(text, str(audio_path)):
                logger.info(f"✅ Minimal audio generated: {audio_path}")
                return str(audio_path)
            else:
                logger.error("All audio generation methods failed")
                return None
                
        except Exception as e:
            logger.error(f"Error generating audio: {str(e)}")
            return None
    
    def generate_avatar_video(self, audio_path: str, 
                            avatar_image: Optional[str] = None) -> Optional[str]:
        """
        Generate avatar video using SadTalker with fallback methods
        Returns: path to generated video file or None if failed
        """
        try:
            from simple_video_gen import generate_video_simple
            
            if not audio_path or not Path(audio_path).exists():
                logger.error(f"Audio file not found: {audio_path}")
                return None
            
            # Use default avatar image if none provided
            if avatar_image is None:
                avatar_image = str(config.DEFAULT_AVATAR_IMAGE)
            
            if not Path(avatar_image).exists():
                logger.error(f"Avatar image not found: {avatar_image}")
                return None
            
            # Prepare output filename
            timestamp = int(time.time())
            video_filename = f"mindfull_avatar_{timestamp}.mp4"
            video_path = config.VIDEO_OUTPUT_DIR / video_filename
            
            logger.info(f"Generating avatar video: {video_filename}")
            
            # Use our simple video generation system
            if generate_video_simple(avatar_image, audio_path, str(video_path)):
                logger.info(f"✅ Avatar video generated: {video_path}")
                return str(video_path)  # Return the path on success
            else:
                logger.error("Video generation failed")
                return None
                
        except Exception as e:
            logger.error(f"Error generating avatar video: {str(e)}")
            return None
    
    def process_complete_interaction(self, user_input: str, 
                                   avatar_image: Optional[str] = None) -> Dict[str, Any]:
        """
        Process complete interaction: text -> response -> audio -> video
        Returns: dictionary with all generated content paths and metadata
        """
        start_time = time.time()
        result = {
            "session_id": self.session.session_id if self.session else None,
            "user_input": user_input,
            "timestamp": datetime.now().isoformat(),
            "response_text": None,
            "emotion_detected": None,
            "audio_path": None,
            "video_path": None,
            "processing_time": 0,
            "errors": []
        }
        
        try:
            # Step 1: Generate text response
            logger.info("Step 1: Generating text response...")
            response_text, emotion = self.generate_response(user_input)
            result["response_text"] = response_text
            result["emotion_detected"] = emotion
            
            if not response_text or response_text in config.FALLBACK_MESSAGES.values():
                result["errors"].append("Failed to generate proper response")
                return result
            
            # Step 2: Generate audio
            logger.info("Step 2: Generating audio...")
            audio_path = self.generate_audio(response_text)
            result["audio_path"] = audio_path
            
            if not audio_path:
                result["errors"].append("Failed to generate audio")
            
            # Step 3: Generate avatar video (optional)
            if self.sadtalker_available and audio_path:
                logger.info("Step 3: Generating avatar video...")
                video_path = self.generate_avatar_video(audio_path, avatar_image)
                result["video_path"] = video_path
                
                if not video_path:
                    result["errors"].append("Failed to generate avatar video")
            else:
                logger.info("Step 3: Skipping avatar video generation")
            
            result["processing_time"] = time.time() - start_time
            logger.info(f"✅ Complete interaction processed in {result['processing_time']:.2f}s")
            
        except Exception as e:
            logger.error(f"Error in complete interaction: {str(e)}")
            result["errors"].append(f"Processing error: {str(e)}")
            result["processing_time"] = time.time() - start_time
        
        return result
    
    def get_session_info(self) -> Optional[Dict[str, Any]]:
        """Get current session information"""
        if not self.session or not self.session.start_time or not self.session.last_interaction:
            return None
        
        return {
            "session_id": self.session.session_id,
            "user_name": self.session.user_name,
            "user_rank": self.session.user_rank,
            "start_time": self.session.start_time.isoformat(),
            "last_interaction": self.session.last_interaction.isoformat(),
            "conversation_count": len(self.session.conversation_history or [])
        }
    
    def cleanup_old_files(self, max_age_hours: int = 24):
        """Clean up old generated files"""
        try:
            current_time = time.time()
            cutoff_time = current_time - (max_age_hours * 3600)
            
            for directory in [config.AUDIO_OUTPUT_DIR, config.VIDEO_OUTPUT_DIR, config.TEMP_DIR]:
                if directory.exists():
                    for file_path in directory.iterdir():
                        if file_path.is_file() and file_path.stat().st_mtime < cutoff_time:
                            file_path.unlink()
                            logger.info(f"Cleaned up old file: {file_path}")
                            
        except Exception as e:
            logger.error(f"Error during cleanup: {str(e)}")

def main():
    """Main function for testing the pipeline"""
    print("🤖 Mindfull AI Avatar Chatbot")
    print("=" * 50)
    
    # Initialize pipeline
    pipeline = MindfullPipeline()
    
    if not pipeline.is_initialized:
        print("❌ Failed to initialize pipeline. Please check the logs.")
        return
    
    # Create session
    session_id = pipeline.create_session()
    print(f"Session ID: {session_id}")
    
    # Interactive loop
    print("\n💬 Chat with Mindfull (type 'quit' to exit)")
    print("Note: This will generate audio and video responses")
    print("-" * 50)
    
    while True:
        try:
            user_input = input("\nYou: ").strip()
            
            if user_input.lower() in ['quit', 'exit', 'bye']:
                print("👋 Goodbye! Take care!")
                break
            
            if not user_input:
                continue
            
            print("\n🔄 Processing your message...")
            
            # Process complete interaction
            result = pipeline.process_complete_interaction(user_input)
            
            # Display results
            print(f"\n🤖 Mindfull: {result['response_text']}")
            print(f"😊 Emotion detected: {result['emotion_detected']}")
            print(f"⏱️ Processing time: {result['processing_time']:.2f}s")
            
            if result['audio_path']:
                print(f"🔊 Audio saved: {result['audio_path']}")
            
            if result['video_path']:
                print(f"🎥 Video saved: {result['video_path']}")
            
            if result['errors']:
                print(f"⚠️ Errors: {', '.join(result['errors'])}")
            
        except KeyboardInterrupt:
            print("\n\n👋 Goodbye! Take care!")
            break
        except Exception as e:
            print(f"\n❌ Error: {str(e)}")
    
    # Cleanup
    print("\n🧹 Cleaning up old files...")
    pipeline.cleanup_old_files()
    print("✅ Cleanup complete")

if __name__ == "__main__":
    main()