agent_websearch.py

import wikipedia
import requests
from bs4 import BeautifulSoup
from typing import List, Dict, Any
import os

from ddgs import DDGS
from dotenv import load_dotenv
from smolagents import CodeAgent, LiteLLMModel, tool, DuckDuckGoSearchTool


# Configuration de Wikipedia en français
wikipedia.set_lang("fr")

# Initialisation de l'outil de recherche DuckDuckGo
search_tool = DuckDuckGoSearchTool(5)

load_dotenv()  # Charge les variables d'environnement

@tool
def search_internet(query: str) -> List[Dict[str, str]]:
    """
    Recherche des informations sur internet via DuckDuckGo.

    Args:
        query: Terme de recherche

    Returns:
        Liste de dictionnaires contenant les résultats de recherche
    """
    try:
        # Utilisation de DuckDuckGoSearchTool
        results = search_tool(query)
        return results
    except Exception as e:
        return [{'error': f'Erreur lors de la recherche: {str(e)}'}]


@tool
def get_wikipedia_page(title: str, lang: str = "fr") -> Dict[str, Any]:
    """
    Récupère une page Wikipedia complète avec son contenu et ses tableaux sous forme d'un dictionnaire.

    Args:
        title: Titre de la page Wikipedia
        lang: Langue de Wikipedia (par défaut: français)

    Returns:
        Dictionnaire contenant les informations de la page et ses tableaux
    """
    try:
        wikipedia.set_lang(lang)
        page = wikipedia.page(title)

        # Récupération du HTML de la page pour extraire les tableaux
        response = requests.get(page.url)
        soup = BeautifulSoup(response.content, 'html.parser')

        # Extraction des tableaux avec leurs titres
        tables = []
        wiki_tables = soup.find_all('table', class_='wikitable')

        for i, table in enumerate(wiki_tables):
            table_data = []
            headers = []
            table_title = ""

            # Recherche du titre du tableau
            # Chercher dans les éléments précédents pour trouver un titre
            current_element = table.find_previous_sibling()
            title_found = False

            while current_element and not title_found:
                if current_element.name in ['h1', 'h2', 'h3', 'h4', 'h5', 'h6']:
                    table_title = current_element.get_text(strip=True)
                    title_found = True
                elif current_element.name == 'p':
                    # Parfois le titre est dans un paragraphe juste avant
                    text = current_element.get_text(strip=True)
                    if len(text) < 100 and text.endswith(':'):
                        table_title = text.rstrip(':')
                        title_found = True
                elif current_element.name == 'caption':
                    # Certains tableaux ont une balise caption
                    table_title = current_element.get_text(strip=True)
                    title_found = True
                current_element = current_element.find_previous_sibling()

            # Vérifier si le tableau a une caption directe
            caption = table.find('caption')
            if caption and not table_title:
                table_title = caption.get_text(strip=True)

            # Si aucun titre trouvé, chercher dans les éléments suivants
            if not table_title:
                # Regarder les éléments th avec colspan ou les premiers éléments
                first_row = table.find('tr')
                if first_row:
                    first_cell = first_row.find(['th', 'td'])
                    if first_cell and first_cell.get('colspan'):
                        potential_title = first_cell.get_text(strip=True)
                        if len(potential_title) < 100:
                            table_title = potential_title

            # Valeur par défaut si aucun titre trouvé
            if not table_title:
                table_title = f"Tableau {i + 1}"

            # Extraction des en-têtes
            header_row = table.find('tr')
            if header_row:
                for th in header_row.find_all(['th', 'td']):
                    headers.append(th.get_text(strip=True))

            # Extraction des données
            rows = table.find_all('tr')
            data_rows = []

            # Déterminer quelle ligne contient les données (pas les en-têtes)
            start_row = 1 if headers else 0

            for row in rows[start_row:]:
                cells = row.find_all(['td', 'th'])
                if cells:
                    row_data = {}
                    for j, cell in enumerate(cells):
                        header = headers[j] if j < len(headers) else f'Colonne_{j + 1}'
                        cell_text = cell.get_text(strip=True)
                        row_data[header] = cell_text
                    if any(row_data.values()):  # Ignorer les lignes vides
                        data_rows.append(row_data)

            tables.append({
                'table_index': i,
                'table_title': table_title,
                'headers': headers,
                'data': data_rows,
                'row_count': len(data_rows),
                'column_count': len(headers)
            })

        return {
            'title': page.title,
            'url': page.url,
            'summary': page.summary,
            'content': page.content,
            'links': page.links[:20],  # Limite aux 20 premiers liens
            'categories': page.categories[:10],  # Limite aux 10 premières catégories
            'images': page.images[:5] if page.images else [],  # Limite aux 5 premières images
            'tables': tables,
            'table_count': len(tables),
            'page_length': len(page.content),
            'has_tables': len(tables) > 0
        }
    except wikipedia.exceptions.DisambiguationError as e:
        return {
            'error': 'Ambiguité détectée',
            'options': e.options[:10],  # Limite aux 10 premières options
            'type': 'disambiguation'
        }
    except wikipedia.exceptions.PageError:
        return {
            'error': f'Page "{title}" non trouvée',
            'type': 'page_not_found'
        }
    except Exception as e:
        return {
            'error': f'Erreur: {str(e)}',
            'type': 'general_error'
        }


# Configuration du modèle LiteLLM
def create_wikipedia_agent():
    """
    Crée un agent Wikipedia avec tous les outils nécessaires.
    """
    # Initialisation du modèle (ajustez selon votre configuration)
    model = LiteLLMModel(model_id="gemini/gemini-2.5-flash", api_key=os.getenv("GOOGLE_API_KEY"))

    system_prompt = """
    Tu es un agent de recherche d'information, tu es capable de trouver des informations sur Wikipédia ou par recherche internet.
    Tu peux exécuter des recherches, récupérer des pages Wikipedia, extraire des tableaux.
    Pour trouver un article, tu peux rechercher sur internet avec le plus de détails puis trouver le premier lien Wikipedia pertinent.
    Si tu ne trouves pas d'outils spécifiques à une tâche, tu peux utiliser la recherche internet pour trouver des informations pertinentes.
    Tu dois toujours fournir des réponses claires et concises, en citant les sources lorsque c'est possible.
    """
    # Création de l'agent avec les outils
    return CodeAgent(
        name="websearch_agent",
        tools=[
            search_internet,
            get_wikipedia_page,
        ],
        model=model,
        description=system_prompt,
        additional_authorized_imports=["wikipedia", "requests", "bs4", "pandas"]
    )