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, PythonInterpreterTool

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

# --- CLASSE ORCHESTRATEUR COMPLÈTE ---
class OrchestratorAgent(MultiStepAgent):
    """
    Agent orchestrateur qui hérite de MultiStepAgent et implémente
    les méthodes requises pour le prompt système et la réponse finale.
    """
    def initialize_system_prompt(self) -> str:
        """Définit le prompt système pour l'agent orchestrateur."""
        return (
            "You are a world-class autonomous agent. Your goal is to fully answer the user's question. "
            "To do so, you have access to a set of tools. "
            "First, think step-by-step and lay out a plan to solve the problem. "
            "Then, execute the plan by calling the tools in the required sequence. "
            "Analyze the results of each tool call. If the plan is not working, reassess and create a new plan. "
            "When you have the final answer, present it clearly to the user."
        )

    def render_final_answer(self, final_context: dict, final_response: str) -> str:
        """
        Formate et retourne la réponse finale de l'agent.
        """
        # Pour l'instant, nous retournons simplement la réponse brute du modèle.
        return final_response

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")
        self.model_manager = ModelManager(self.hf_token)
        self.tools = [search_web, scrape_website, read_file, PythonInterpreterTool()]
        self._init_specialized_agents()
        self.conversation_history = []
        print("✅ UltraAgent initialisé avec succès!")

    def _init_specialized_agents(self):
        """Initialise les agents spécialisés."""
        try:
            self.orchestrator = OrchestratorAgent(
                model=self.model_manager.get_orchestrator(),
                tools=self.tools,
                max_steps=15,
                planning_interval=3
            )
            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."""
        print("\n" + "="*80)
        print(f"🧠 ULTRA-AGENT - Nouvelle question reçue\nQuestion: {question[:200]}{'...' if len(question) > 200 else ''}")
        print("="*80)
        
        try:
            self.conversation_history.append({"role": "user", "content": question, "timestamp": time.time()})
            strategy = self._analyze_question(question)
            print(f"📋 Stratégie sélectionnée: {strategy['type']}")
            response = self._execute_strategy(question, strategy)
            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)
            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."""
        # Cette fonction de scoring simple reste la même
        question_lower = question.lower()
        vision_keywords = ['image', 'photo', 'picture', 'visual', 'voir', 'regarder', 'analyser l\'image', 'screenshot']
        code_keywords = ['code', 'program', 'script', 'algorithm', 'python', 'javascript', 'sql', 'debug', 'programmer']
        if any(kw in question_lower for kw in vision_keywords):
            return {"type": "vision"}
        if any(kw in question_lower for kw in code_keywords):
            return {"type": "code"}
        return {"type": "general"} # Simplification de la stratégie pour le débogage

    def _execute_strategy(self, question: str, strategy: Dict[str, Any]) -> str:
        """Exécute la stratégie déterminée."""
        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)
            else: # general, search_web, file_processing, reasoning
                return self._handle_general_task(question)
        except Exception as e:
            print(f"❌ Erreur dans l'exécution de la stratégie {strategy_type}: {e}")
            raise e # Fait remonter l'erreur pour la gestion centrale

    def _handle_vision_task(self, question: str) -> str:
        """Gère les tâches de vision."""
        print("👁️ Traitement avec le modèle de vision...")
        response = self.model_manager.get_vision_model()(question)
        return f"🔍 Analyse visuelle:\n{response}"

    def _handle_code_task(self, question: str) -> str:
        """Gère les tâches de code."""
        print("💻 Traitement avec l'agent de code...")
        return self.code_agent.run(question)

    def _handle_general_task(self, question: str) -> str:
        """Gère toutes les autres tâches avec l'orchestrateur."""
        print("🎯 Traitement général avec l'orchestrateur...")
        context = self._get_conversation_context()
        enhanced_prompt = f"{context}\nQuestion actuelle: {question}"
        return self.orchestrator.run(enhanced_prompt)

    def _fallback_reasoning(self, question: str, error: str) -> str:
        """Fallback avec le modèle de raisonnement."""
        fallback_prompt = f"An error occurred while processing the question. Question: '{question}'. Error: '{error}'. Provide the best possible answer, explaining the limitations due to the error."
        return self.model_manager.get_reasoning_model()(fallback_prompt)

    def _get_conversation_context(self, max_exchanges: int = 3) -> str:
        """Récupère le contexte des échanges récents."""
        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]
            context_parts.append(f"{role}: {content}")
        return "\n".join(context_parts)


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]}"