dnd-helpers/main.py

import argparse

from src.rag.manager import RAGManager


def main():
    parser = argparse.ArgumentParser(description="D&D Helpers CLI")
    parser.add_argument(
        "--ingest-pdf",
        type=str,
        help="Path to a PDF file to ingest into the RAG system",
    )
    parser.add_argument(
        "--ingest-file",
        type=str,
        help="Path to a markdown file to ingest into the RAG system",
    )
    parser.add_argument(
        "--ingest-dir",
        type=str,
        help="Path to a directory of markdown files to ingest into the RAG system",
    )

    args = parser.parse_args()

    rag_manager = RAGManager()

    if args.ingest_pdf:
        print(f"Ingesting PDF: {args.ingest_pdf}...")
        rag_manager.ingest_pdf(args.ingest_pdf)
        print("PDF ingestion complete.")

    if args.ingest_file:
        print(f"Ingesting File: {args.ingest_file}...")
        rag_manager.ingest_file(args.ingest_file)
        print("File ingestion complete.")

    if args.ingest_dir:
        print(f"Ingesting Directory: {args.ingest_dir}...")
        rag_manager.ingest_directory(args.ingest_dir)
        print("Directory ingestion complete.")

    if not any([args.ingest_pdf, args.ingest_file, args.ingest_dir]):
        print("Hello from dnd-helpers!")


if __name__ == "__main__":
    main()