tools/wikipedia_rag/run.py

import argparse
import os

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()