MooseCP/tools/browse_wikipedia.py

import httpx
import json
import re
import asyncio
from typing import Any, Dict, List, Optional
from .utils import ToolError
import config

API_URL = "https://en.wikipedia.org/w/api.php"

def strip_html(text: str) -> str:
    """Removes HTML tags from a string using regex to provide clean text to the LLM."""
    return re.sub(r'<[^>]*>', '', text)

def clean_wikitext(text: str) -> str:
    """
    Removes the most distracting elements of raw Wikitext:
    1. HTML comments (<!-- ... -->)
    2. Citations (<ref /> and <ref>...</ref>)
    """
    # Remove HTML comments
    text = re.sub(r'<!--.*?-->', '', text, flags=re.DOTALL)
    # Remove self-closing citations FIRST to prevent them being seen as opening tags
    text = re.sub(r'<ref[^>]*/>', '', text)
    # Remove paired citations
    text = re.sub(r'<ref[^>]*>.*?</ref>', '', text, flags=re.DOTALL)
    return text

async def _make_request(params: Dict[str, Any]) -> Dict[str, Any]:
    """
    Asynchronous helper to make the API request using httpx.
    A custom User-Agent is used to avoid bot detection.
    """
    headers = {
        "User-Agent": config.USER_AGENT
    }
    
    async with httpx.AsyncClient(headers=headers, timeout=15.0) as client:
        response = await client.get(API_URL, params=params)
        response.raise_for_status()
        return response.json()

async def _search(title: str, limit: int = 5, fallback: bool = False) -> str:
    """
    Searches Wikipedia for a title.
    fallback=True: Used when a direct page request fails, providing 'Article not found' header.
    """
    params = {
        "action": "query",
        "list": "search",
        "srsearch": title,
        "format": "json",
        "srlimit": limit
    }
    # Directly await the async request
    data = await _make_request(params)
    
    search_results = data.get("query", {}).get("search", [])
    if not search_results:
        return f"No Wikipedia results found for '{title}'."
    
    if fallback:
        header = f"Article not found. You MUST select one of the following articles and submit another query:"
    else:
        header = f"Search results for '{title}':"
        
    lines = [header]
    for i, res in enumerate(search_results, 1):
        title_res = res.get("title")
        snippet = strip_html(res.get("snippet", ""))
        # Brackets around title help the AI identify the exact string for subsequent calls
        lines.append(f"{i}. [{title_res}] - Snippet: {snippet}")
    
    if fallback:
        lines.append("\nUse the title in brackets [ ] to explore the selected page.")
        
    return "\n".join(lines)

async def _fetch_content(title: str, section_index: int) -> tuple[Optional[str], Optional[str]]:
    """
    Retrieves raw Wikitext for a specific section.
    section_index=0 returns the lead section.
    Returns a tuple of (content, final_title).
    """
    params = {
        "action": "query",
        "prop": "revisions",
        "rvprop": "content",
        "rvsection": section_index,
        "titles": title,
        "redirects": 1,
        "format": "json"
    }
    data = await _make_request(params)
    
    pages = data.get("query", {}).get("pages", {})
    if not pages:
        return None, None
        
    page_id = next(iter(pages))
    page = pages[page_id]
    final_title = page.get("title")
    
    if "missing" in page:
        return None, None
        
    revisions = page.get("revisions", [])
    if not revisions:
        return None, final_title
        
    content = revisions[0].get("*")
    return (clean_wikitext(content) if content else None), final_title

async def _fetch_toc(title: str) -> tuple[Optional[str], Optional[str]]:
    """
    Retrieves the Table of Contents data and formats it hierarchically.
    Returns a tuple of (toc, final_title).
    """
    params = {
        "action": "parse",
        "page": title,
        "prop": "tocdata",
        "format": "json",
        "redirects": 1
    }
    data = await _make_request(params)
    
    parse_data = data.get("parse")
    if not parse_data:
        return None, None
        
    final_title = parse_data.get("title")
    toc_data = parse_data.get("tocdata", {})
    sections = toc_data.get("sections", [])
    if not sections:
        return "No table of contents found for this page.", final_title
    
    lines = [f"Table of Contents for \"{final_title}\":"]
    for s in sections:
        level = s.get("tocLevel", 1)
        index = s.get("index")
        line = s.get("line")
        indent = "  " * (level - 1)
        lines.append(f"{indent}[{index}] {line}")
        
    return "\n".join(lines), final_title

async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
    """
    Main handler for the browse_wikipedia tool.
    Supports modes: summary, toc, section, and search.
    """
    title = args.get("title")
    if not title:
        raise ToolError("Missing required parameter: 'title'")
        
    mode = args.get("mode", "summary")
    section_index = args.get("section_index")
    search_limit = args.get("search_limit", 5)
    
    final_title = None
    
    if mode == "search":
        result = await _search(title, search_limit, fallback=False)
    elif mode == "summary":
        # Lead section + ToC is the default 'summary' to guide the AI's next steps
        content, title_from_content = await _fetch_content(title, 0)
        final_title = title_from_content
        if content is None:
            result = await _search(title, search_limit, fallback=True)
        else:
            toc, title_from_toc = await _fetch_toc(title)
            if title_from_toc:
                final_title = title_from_toc
            result = f"{content}\n\n---\n\n{toc}"
    elif mode == "toc":
        toc, title_from_toc = await _fetch_toc(title)
        final_title = title_from_toc
        if toc is None:
            result = await _search(title, search_limit, fallback=True)
        else:
            result = toc
    elif mode == "section":
        if section_index is None:
            raise ToolError("Missing required parameter 'section_index' for mode='section'")
        content, title_from_content = await _fetch_content(title, int(section_index))
        final_title = title_from_content
        if content is None:
            result = await _search(title, search_limit, fallback=True)
        else:
            result = content
    else:
        raise ToolError(f"Invalid mode '{mode}'. Supported modes: summary, toc, section, search")
    
    if final_title and final_title != title:
        result = f"Redirected to \"{final_title}\"\n\n{result}"
            
    return [{"type": "text", "text": result}]