test_rag.py

"""
This module provides functions to test the RAG system API endpoints.
It includes tests for uploading documents, querying, and clearing the database.
"""

import os

import requests

BASE_URL = "http://localhost:8000"


def test_upsert(file_path):
    """
    Test the document upload (upsert) endpoint.
    
    Args:
        file_path (str): Path to the PDF file to upload.
    """
    with open(file_path, "rb") as f:
        response = requests.post(f"{BASE_URL}/upsert", files={"file": f})
    print(f"Upsert response: {response.json()}")


def test_query(question):
    """
    Test the query endpoint.
    
    Args:
        question (str): The question to ask the RAG system.
    """
    response = requests.post(f"{BASE_URL}/query", json={"text": question})
    print(f"Query: {question}")
    print(f"Response: {response.json()}")


def clear_db():
    """
    Test the clear database endpoint.
    """
    response = requests.post(f"{BASE_URL}/clear_db")
    print(f"Clear DB response: {response.json()}")


if __name__ == "__main__":
    # Ensure this path is correct and the file exists

    multiple_pdfs = ["./pdfs/1702.04405v2.pdf",
                     "./pdfs/2004.11985v1.pdf",
                    "./pdfs/2404.10699v1.pdf",
                     "./pdfs/DALES_Objects_A_Large_Scale_Benchmark_Dataset_for_Instance_Segmentation_in_Aerial_Lidar.pdf"]
    
    # Optionally, clear the database before testing
    clear_db()

    # Test document upload
    for pdf in multiple_pdfs:
        print("Processing pdf: ", pdf)
        test_upsert(pdf)

    # Test queries
    test_questions = [
        "Explain Dales dataset and DalesObjects datasets?",
        "Can you compare the different datasets that are in the documents?",
        "which dataset is most well tested?"
    ]

    for question in test_questions:
        test_query(question)