agent.py

# Fichier: agent.py (Version Corrigée et Complète)

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)
        
# La classe BasicAgent reste la même, elle ne fait qu'appeler UltraAgent

class BasicAgent:
    def __init__(self):
        try:
            if not os.getenv("HF_TOKEN"):
                print("⚠️ Attention: Le token Hugging Face (HF_TOKEN) n'est pas défini.")
            self.ultra_agent = UltraAgent()
        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:
        if self.ultra_agent is None:
            return "Erreur: L'agent n'a pas pu être initialisé. Vérifiez les logs et la configuration (HF_TOKEN)."
        return self.ultra_agent(question)