tools/wikipedia_rag/run.py

import os
import argparse
from dotenv import load_dotenv
from tool import WikipediaRAGTool

def main():
    # Load environment variables
    load_dotenv()
    
    # Set up argument parser
    parser = argparse.ArgumentParser(description='Run Wikipedia RAG Tool')
    parser.add_argument('--query', type=str, required=True, help='Search query for Wikipedia articles')
    parser.add_argument('--dataset-path', type=str, default='wikipedia-structured-contents',
                      help='Path to the Wikipedia dataset')
    args = parser.parse_args()
    
    # Initialize the tool
    tool = WikipediaRAGTool(dataset_path=args.dataset_path)
    
    # Run the query
    print(f"\nQuery: {args.query}")
    print("-" * 50)
    result = tool.forward(args.query)
    print(f"Result: {result}")
    print("-" * 50)

if __name__ == "__main__":
    main()