dnd-helpers/tests/test_rag.py

import os

from reportlab.pdfgen import canvas

from src.rag.manager import RAGManager


def create_dummy_phb(pdf_path: str):
    """
    Creates a dummy PDF file to simulate the Player's Handbook for verification.
    """
    print(f"Creating dummy PHB at {pdf_path}...")
    c = canvas.Canvas(pdf_path)

    # Page 1: Fireball
    c.drawString(100, 750, "Fireball")
    c.drawString(
        100,
        730,
        "A bright streak flashes from your pointing finger to a point you choose within range.",
    )
    c.drawString(
        100,
        710,
        "Each creature in a 20-foot-radius sphere centered on that point must make a Dexterity saving throw.",
    )
    c.showPage()

    # Page 2: Grappling
    c.drawString(100, 750, "Grappling")
    c.drawString(
        100,
        730,
        "Grappling is a special option available to any attack that hits with a melee attack.",
    )
    c.drawString(
        100,
        710,
        "A creature is grappled if it is the target of the grapple attack and the attack hits.",
    )
    c.showPage()

    # Page 3: General Rules
    c.drawString(100, 750, "General Rules")
    c.drawString(
        100,
        730,
        "Combat is resolved in rounds, and each round represents 6 seconds of in-game time.",
    )
    c.showPage()

    c.save()
    print(f"Dummy PHB created successfully at {pdf_path}")


def test_rag():
    pdf_path = "data/phb_dummy.pdf"
    create_dummy_phb(pdf_path)

    # Initialize RAG Manager
    rag = RAGManager(persist_dir="data/rag_index_test")

    # Task 2.2: Ingest PDF
    print("\nTesting Ingestion...")
    rag.ingest_pdf(pdf_path)

    # Task 2.3: Retrieve Logic
    print("\nTesting Retrieval...")

    # Test 1: Fireball
    query1 = "What is Fireball?"
    results1 = rag.retrieve(query1)
    print(f"Query: {query1}")
    for res in results1:
        print(f"Source: {res.source} | Snippet: {res.snippet[:100]}...")

    assert len(results1) > 0, "Should have retrieved at least one result for 'Fireball'"
    assert "Fireball" in results1[0].snippet, "The top result should contain 'Fireball'"

    # Test 2: Grappling
    query2 = "How does grappling work?"
    results2 = rag.retrieve(query2)
    print(f"Query: {query2}")
    for res in results2:
        print(f"Source: {res.source} | Snippet: {res.snippet[:100]}...")

    assert len(results2) > 0, (
        "Should have retrieved at least one result for 'Grappling'"
    )
    assert "Grappling" in results2[0].snippet, (
        "The top result should contain 'Grappling'"
    )

    print("\n✅ RAG Verification Successful!")


if __name__ == "__main__":
    test_rag()