agent.py

import os
import re
import json
import time
from typing import List, Dict, Any, Optional
from pathlib import Path
import tempfile

from smolagents import CodeAgent, MultiStepAgent, AgentError
from smolagents.tools import PythonInterpreterTool

from models import ModelManager
from tools import search_web, scrape_website, read_file

class UltraAgent:
    """
    Agent ultra-puissant avec orchestration multi-modèles.
    Utilise différents modèles spécialisés selon le type de tâche.
    """
    
    def __init__(self, hf_token: Optional[str] = None):
        """Initialise l'UltraAgent et ses composants."""
        print("🚀 Initialisation de l'UltraAgent...")
        
        self.hf_token = hf_token or os.getenv("HF_TOKEN")
        
        # Gestionnaire de modèles
        self.model_manager = ModelManager(self.hf_token)
        
        # Outils disponibles
        self.tools = [
            search_web,
            scrape_website,
            read_file,
            PythonInterpreterTool()
        ]
        
        # Agents spécialisés
        self._init_specialized_agents()
        
        # Historique des interactions
        self.conversation_history = []
        
        print("✅ UltraAgent initialisé avec succès!")

    def _init_specialized_agents(self):
        """Initialise les agents spécialisés comme l'orchestrateur et l'agent de code."""
        try:
            # Agent principal (orchestrateur)
            self.orchestrator = MultiStepAgent(
                model=self.model_manager.get_orchestrator(),
                tools=self.tools,
                max_steps=15,
                planning_interval=3
            )
            
            # Agent de code
            self.code_agent = CodeAgent(
                model=self.model_manager.get_code_agent(),
                tools=[PythonInterpreterTool()],
                additional_authorized_imports=["requests", "pandas", "numpy", "matplotlib", "seaborn", "scipy", "sklearn"]
            )
            
            print("✅ Agents spécialisés initialisés")
            
        except Exception as e:
            print(f"❌ Erreur lors de l'initialisation des agents: {e}")
            raise

    def __call__(self, question: str) -> str:
        """
        Point d'entrée principal de l'agent.
        Analyse la question, détermine la stratégie et exécute la tâche.
        """
        print("\n" + "="*80)
        print("🧠 ULTRA-AGENT - Nouvelle question reçue")
        print(f"Question: {question[:200]}{'...' if len(question) > 200 else ''}")
        print("="*80)
        
        try:
            # Ajoute à l'historique
            self.conversation_history.append({
                "role": "user",
                "content": question,
                "timestamp": time.time()
            })
            
            # Analyse de la question et choix de la stratégie
            strategy = self._analyze_question(question)
            print(f"📋 Stratégie sélectionnée: {strategy['type']}")
            
            # Exécution selon la stratégie
            response = self._execute_strategy(question, strategy)
            
            # Ajoute la réponse à l'historique
            self.conversation_history.append({
                "role": "assistant", 
                "content": response,
                "strategy": strategy,
                "timestamp": time.time()
            })
            
            print(f"✅ Réponse générée avec succès ({len(response)} caractères)")
            return response
            
        except Exception as e:
            error_msg = f"❌ Erreur critique dans l'UltraAgent: {e}"
            print(error_msg)
            
            # Tentative de récupération avec le modèle de raisonnement
            try:
                print("🔄 Tentative de récupération avec le modèle de raisonnement...")
                fallback_response = self._fallback_reasoning(question, str(e))
                return fallback_response
            except Exception as fallback_e:
                final_error = f"Je rencontre des difficultés techniques. Erreur: {str(fallback_e)[:200]}"
                print(f"❌ La récupération a également échoué: {fallback_e}")
                return final_error

    def _analyze_question(self, question: str) -> Dict[str, Any]:
        """
        Analyse la question pour déterminer la meilleure stratégie à adopter.
        """
        question_lower = question.lower()
        
        # Détection de mots-clés pour différents types de tâches
        vision_keywords = ['image', 'photo', 'picture', 'visual', 'voir', 'regarder', 'analyser l\'image', 'screenshot']
        code_keywords = ['code', 'program', 'script', 'algorithm', 'python', 'javascript', 'sql', 'debug', 'programmer']
        search_keywords = ['search', 'find', 'look up', 'recherche', 'chercher', 'trouver', 'internet', 'web', 'actualité']
        file_keywords = ['file', 'document', 'pdf', 'excel', 'csv', 'json', 'read', 'fichier', 'lire', 'ouvrir']
        reasoning_keywords = ['complex', 'analyze', 'think', 'reason', 'solve', 'problem', 'complexe', 'analyser', 'résoudre', 'problème']
        
        # Calcul des scores
        scores = {
            "vision": sum(1 for kw in vision_keywords if kw in question_lower),
            "code": sum(1 for kw in code_keywords if kw in question_lower),
            "search_web": sum(1 for kw in search_keywords if kw in question_lower),
            "file_processing": sum(1 for kw in file_keywords if kw in question_lower),
            "reasoning": sum(1 for kw in reasoning_keywords if kw in question_lower),
        }

        # Détection d'URL pour forcer la recherche web
        if any(domain in question_lower for domain in ['http', 'www', '.com', '.fr', '.org']):
            scores["search_web"] += 2
        
        # Analyse de la complexité
        complexity = len(question.split()) + len(re.findall(r'[.!?]', question))

        # Détermination de la stratégie
        if scores["vision"] > 0:
            return {"type": "vision", "priority": scores["vision"], "complexity": complexity}
        if scores["code"] > 1:
            return {"type": "code", "priority": scores["code"], "complexity": complexity}
        if scores["search_web"] > 0:
            return {"type": "search_web", "priority": scores["search_web"], "complexity": complexity}
        if scores["file_processing"] > 0:
            return {"type": "file_processing", "priority": scores["file_processing"], "complexity": complexity}
        if complexity > 50 or scores["reasoning"] > 0:
            return {"type": "reasoning", "priority": scores["reasoning"], "complexity": complexity}
        
        return {"type": "general", "priority": 0, "complexity": complexity}

    def _execute_strategy(self, question: str, strategy: Dict[str, Any]) -> str:
        """Exécute la stratégie déterminée pour répondre à la question."""
        strategy_type = strategy["type"]
        
        try:
            if strategy_type == "vision":
                return self._handle_vision_task(question)
            elif strategy_type == "code":
                return self._handle_code_task(question)
            elif strategy_type == "search_web":
                return self._handle_web_search_task(question)
            elif strategy_type == "file_processing":
                return self._handle_file_task(question)
            elif strategy_type == "reasoning":
                return self._handle_reasoning_task(question)
            else: # general
                return self._handle_general_task(question)
                
        except Exception as e:
            print(f"❌ Erreur dans l'exécution de la stratégie {strategy_type}: {e}")
            # Fallback vers l'orchestrateur général en cas d'échec
            return self._handle_general_task(question)

    def _handle_vision_task(self, question: str) -> str:
        """Gère les tâches impliquant la vision."""
        print("👁️ Traitement avec le modèle de vision...")
        
        try:
            vision_prompt = f"Analyse visuelle experte requise pour la question suivante : {question}. Si une image est mentionnée ou nécessaire, décris précisément ce que tu dois analyser."
            response = self.model_manager.get_vision_model()(vision_prompt)
            return f"🔍 Analyse visuelle:\n{response}"
        except Exception as e:
            print(f"❌ Erreur vision: {e}")
            return self._handle_general_task(question)

    def _handle_code_task(self, question: str) -> str:
        """Gère les tâches de programmation avec l'agent de code."""
        print("💻 Traitement avec l'agent de code...")
        try:
            response = self.code_agent.run(question)
            return f"💻 Solution de code:\n{response}"
        except Exception as e:
            print(f"❌ Erreur code agent: {e}")
            return self._handle_general_task(question)

    def _handle_web_search_task(self, question: str) -> str:
        """Gère les tâches nécessitant une recherche web."""
        print("🌐 Traitement avec recherche web...")
        try:
            search_prompt = f"Répondez à cette question en utilisant une recherche web approfondie. Utilisez les outils de recherche et de scraping, et citez vos sources. Question : {question}"
            response = self.orchestrator.run(search_prompt)
            return response
        except Exception as e:
            print(f"❌ Erreur web search: {e}")
            return self._handle_general_task(question)

    def _handle_file_task(self, question: str) -> str:
        """Gère les tâches impliquant la lecture de fichiers."""
        print("📁 Traitement de fichiers...")
        try:
            file_prompt = f"Traitez les fichiers nécessaires pour répondre à cette question. Utilisez l'outil `read_file` pour analyser les fichiers mentionnés et extraire les informations pertinentes. Question : {question}"
            response = self.orchestrator.run(file_prompt)
            return response
        except Exception as e:
            print(f"❌ Erreur file processing: {e}")
            return self._handle_general_task(question)

    def _handle_reasoning_task(self, question: str) -> str:
        """Gère les tâches complexes nécessitant un raisonnement approfondi."""
        print("🧠 Traitement avec le modèle de raisonnement...")
        try:
            reasoning_prompt = f"En tant qu'expert en résolution de problèmes complexes, analysez et répondez de manière structurée à la question suivante : {question}. Décomposez le problème, analysez chaque partie et proposez une solution étape par étape."
            response = self.model_manager.get_reasoning_model()(reasoning_prompt)
            return f"🧠 Analyse approfondie:\n{response}"
        except Exception as e:
            print(f"❌ Erreur reasoning: {e}")
            return self._handle_general_task(question)

    def _handle_general_task(self, question: str) -> str:
        """Gère les tâches générales avec l'orchestrateur principal."""
        print("🎯 Traitement général avec l'orchestrateur...")
        try:
            context = self._get_conversation_context()
            enhanced_prompt = f"{context}\nQuestion actuelle: {question}\nInstructions: Analysez et répondez de manière complète, en utilisant les outils si nécessaire."
            response = self.orchestrator.run(enhanced_prompt)
            return response
        except Exception as e:
            print(f"❌ Erreur general task: {e}")
            try:
                # Dernière tentative simple
                simple_response = self.model_manager.get_orchestrator()(question)
                return simple_response
            except Exception as final_e:
                raise final_e # Propage l'erreur au handler principal

    def _fallback_reasoning(self, question: str, error: str) -> str:
        """Utilise le modèle de raisonnement en cas d'erreur pour fournir une réponse de secours."""
        try:
            fallback_prompt = f"Une erreur s'est produite en traitant la question. Question: '{question}'. Erreur: '{error}'. Fournissez la meilleure réponse possible en expliquant les limitations dues à l'erreur."
            response = self.model_manager.get_reasoning_model()(fallback_prompt)
            return f"⚠️ Réponse de récupération:\n{response}"
        except Exception as e:
            raise e

    def _get_conversation_context(self, max_exchanges: int = 3) -> str:
        """Récupère le contexte des échanges récents pour l'injecter dans le prompt."""
        if not self.conversation_history:
            return ""
        
        recent_history = self.conversation_history[-max_exchanges*2:]
        context_parts = ["Contexte de la conversation récente:"]
        for entry in recent_history:
            role = "Utilisateur" if entry["role"] == "user" else "Assistant"
            content = str(entry["content"])[:200] + "..." if len(str(entry["content"])) > 200 else str(entry["content"])
            context_parts.append(f"{role}: {content}")
        
        return "\n".join(context_parts)

    def get_status(self) -> Dict[str, Any]:
        """Retourne le statut actuel de l'agent."""
        return {
            "status": "active",
            "conversation_length": len(self.conversation_history),
            "available_tools": [tool.__name__ for tool in self.tools],
            "models_status": self.model_manager.test_models(),
        }

    def reset_conversation(self):
        """Réinitialise l'historique de la conversation."""
        self.conversation_history = []
        print("🔄 Historique de conversation remis à zéro.")


class BasicAgent:
    """
    Classe de compatibilité pour remplacer l'ancien BasicAgent dans app.py.
    Elle utilise l'UltraAgent en arrière-plan pour toute la logique.
    """
    
    def __init__(self):
        """Initialise le BasicAgent en créant une instance de l'UltraAgent."""
        print("🚀 Initialisation du BasicAgent (powered by UltraAgent)...")
        try:
            # Assurez-vous que le token HF est disponible (via les secrets de l'espace HF)
            if not os.getenv("HF_TOKEN"):
                print("⚠️  Attention: Le token Hugging Face (HF_TOKEN) n'est pas défini. L'initialisation pourrait échouer.")
            
            self.ultra_agent = UltraAgent()
            print("✅ BasicAgent initialisé avec succès!")
        except Exception as e:
            print(f"❌ Erreur critique lors de l'initialisation de l'UltraAgent: {e}")
            self.ultra_agent = None
    
    def __call__(self, question: str) -> str:
        """
        Point d'entrée compatible avec l'interface de app.py.
        Délègue l'appel à l'UltraAgent.
        """
        if self.ultra_agent is None:
            return "Erreur: L'agent n'a pas pu être initialisé correctement. Veuillez vérifier les logs et la configuration (notamment le HF_TOKEN)."
        
        try:
            return self.ultra_agent(question)
        except Exception as e:
            print(f"❌ Erreur lors du traitement par l'UltraAgent: {e}")
            return f"Une erreur s'est produite lors du traitement de votre demande: {str(e)[:200]}"