utils/wiki_tool.py

"""
Wikipedia Tool - Tìm kiếm thông tin trên Wikipedia
"""

import wikipedia
from typing import Dict, Any, Optional
import os
from huggingface_hub import InferenceClient

# Khởi tạo client AI để optimize query
HF_TOKEN = os.environ.get("HF_TOKEN")
if HF_TOKEN:
    ai_client = InferenceClient(provider="auto", api_key=HF_TOKEN)
else:
    ai_client = None

def search_wikipedia(query: str, lang: str = "en", summary_length: int = 2000) -> Dict[str, Any]:
    """
    Main function: Tìm kiếm thông tin trên Wikipedia
    
    Args:
        query: Từ khóa tìm kiếm
        lang: Ngôn ngữ Wikipedia (default: "en")
        summary_length: Độ dài tóm tắt tối đa
        
    Returns:
        Dict chứa thông tin Wikipedia
    """
    try:
        # Set language
        wikipedia.set_lang(lang)
        
        # Tìm kiếm trực tiếp
        try:
            page = wikipedia.page(query)
            
            return {
                "success": True,
                "title": page.title,
                "summary": page.summary[:summary_length] if page.summary else "No summary available",
                "url": page.url,
                "categories": page.categories[:5] if hasattr(page, 'categories') else [],
                "query_used": query,
                "search_method": "direct"
            }
            
        except wikipedia.DisambiguationError as e:
            # Nếu có nhiều kết quả, lấy kết quả đầu tiên
            if e.options:
                page = wikipedia.page(e.options[0])
                return {
                    "success": True,
                    "title": page.title,
                    "summary": page.summary[:summary_length] if page.summary else "No summary available",
                    "url": page.url,
                    "categories": page.categories[:5] if hasattr(page, 'categories') else [],
                    "query_used": e.options[0],
                    "search_method": "disambiguation_first",
                    "other_options": e.options[1:5]  # Lấy 4 option khác
                }
            else:
                raise
                
        except wikipedia.PageError:
            # Nếu không tìm thấy trang, thử search
            search_results = wikipedia.search(query, results=5)
            
            if search_results:
                # Thử lấy kết quả đầu tiên
                page = wikipedia.page(search_results[0])
                return {
                    "success": True,
                    "title": page.title,
                    "summary": page.summary[:summary_length] if page.summary else "No summary available",
                    "url": page.url,
                    "categories": page.categories[:5] if hasattr(page, 'categories') else [],
                    "query_used": search_results[0],
                    "search_method": "search_first",
                    "other_results": search_results[1:5]
                }
            else:
                return {
                    "success": False,
                    "error": f"No Wikipedia results found for: {query}",
                    "query_used": query,
                    "search_method": "search_failed"
                }
                
    except Exception as e:
        return {
            "success": False,
            "error": f"Wikipedia search error: {str(e)}",
            "query_used": query,
            "search_method": "error"
        }

def extract_search_query_from_question(question: str) -> str:
    """
    Trích xuất từ khóa tìm kiếm từ câu hỏi
    """
    # Loại bỏ các từ nghi vấn phổ biến
    question_words = ["who", "what", "when", "where", "why", "how", "which", "whose"]
    stop_words = ["is", "are", "was", "were", "the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "by"]
    
    # Chuyển thành lowercase và tách từ
    words = question.lower().replace("?", "").split()
    
    # Loại bỏ question words và stop words
    filtered_words = [
        word for word in words 
        if word not in question_words and word not in stop_words and len(word) > 2
    ]
    
    # Ghép lại thành query
    if filtered_words:
        return " ".join(filtered_words[:4])  # Lấy tối đa 4 từ
    else:
        # Fallback: lấy toàn bộ câu hỏi
        return question.replace("?", "").strip()

def optimize_wiki_query_with_ai(question: str) -> str:
    """
    Dùng AI để sinh ra key search Wikipedia tốt nhất từ câu hỏi
    """
    if not ai_client:
        return extract_search_query_from_question(question)
    prompt = f"""
Given the following question, extract the best possible Wikipedia search query (a short phrase or entity name, not a full sentence). Only output the search query, nothing else.

Question: {question}
"""
    try:
        completion = ai_client.chat.completions.create(
            model="Qwen/Qwen3-8B",
            messages=[{"role": "user", "content": prompt + "\n/no_thinking"}],
            max_tokens=32
        )
        query = completion.choices[0].message.content.strip()
        # Nếu AI trả về rỗng, fallback
        if not query:
            return extract_search_query_from_question(question)
        return query
    except Exception as e:
        print(f"[WikiTool] AI optimize query failed: {e}")
        return extract_search_query_from_question(question)

def search_wikipedia_from_question(question: str, lang: str = "en") -> Dict[str, Any]:
    """
    Tìm kiếm Wikipedia từ câu hỏi, tự động optimize query bằng AI
    """
    # Optimize query bằng AI
    search_query = optimize_wiki_query_with_ai(question)
    print(f"🔍 Wikipedia search query (AI optimized): '{search_query}' from question: '{question[:50]}...'")
    # Tìm kiếm
    result = search_wikipedia(search_query, lang)
    result["original_question"] = question
    result["extracted_query"] = search_query
    return result

def get_multiple_wikipedia_results(query: str, lang: str = "en", num_results: int = 3) -> Dict[str, Any]:
    """
    Lấy nhiều kết quả Wikipedia
    """
    try:
        wikipedia.set_lang(lang)
        search_results = wikipedia.search(query, results=num_results)
        
        results = []
        for result_title in search_results:
            try:
                page = wikipedia.page(result_title)
                results.append({
                    "title": page.title,
                    "summary": page.summary[:500] if page.summary else "No summary",
                    "url": page.url
                })
            except:
                continue
        
        return {
            "success": True,
            "query": query,
            "results": results,
            "total_found": len(results)
        }
        
    except Exception as e:
        return {
            "success": False,
            "error": f"Error getting multiple results: {str(e)}",
            "query": query
        }

# Test function
if __name__ == "__main__":
    # Test direct search
    result1 = search_wikipedia("Mercedes Sosa")
    print("Direct search result:", result1["title"] if result1["success"] else result1["error"])
    
    # Test question-based search
    result2 = search_wikipedia_from_question("Who was Mercedes Sosa?")
    print("Question-based result:", result2["title"] if result2["success"] else result2["error"])
    
    # Test multiple results
    result3 = get_multiple_wikipedia_results("Python programming", num_results=2)
    print("Multiple results:", len(result3["results"]) if result3["success"] else result3["error"])