rag_test.py

import os
import shutil
from dotenv import load_dotenv
from agent import upload_file, vector_store

# Load environment variables
load_dotenv(".env", override=True)
load_dotenv("../.env", override=False)

def test_rag():
    print("--- Starting RAG Test ---")
    
    # Create a dummy test file
    test_file = "rag_test_doc.txt"
    with open(test_file, "w") as f:
        f.write("The secret code is 12345. This is a confidential document for Project X.")
    
    conversation_id_1 = "chat_1"
    conversation_id_2 = "chat_2"
    
    try:
        # 1. Upload to Chat 1
        print(f"\n1. Uploading to {conversation_id_1}...")
        upload_file(test_file, conversation_id_1)
        
        # 2. Retrieve from Chat 1 (Should find it)
        print(f"\n2. Searching in {conversation_id_1}...")
        results_1 = vector_store.similarity_search("secret code", k=1, filter={"conversation_id": conversation_id_1})
        if results_1:
            print(f"✅ Found: {results_1[0].page_content}")
        else:
            print("❌ Not found (Unexpected)")
            
        # 3. Retrieve from Chat 2 (Should NOT find it)
        print(f"\n3. Searching in {conversation_id_2}...")
        results_2 = vector_store.similarity_search("secret code", k=1, filter={"conversation_id": conversation_id_2})
        if not results_2:
            print("✅ Not found (Expected)")
        else:
            print(f"❌ Found (Unexpected): {results_2[0].page_content}")
            
    except Exception as e:
        print(f"Test Failed: {e}")
    finally:
        # Cleanup
        if os.path.exists(test_file):
            os.remove(test_file)

if __name__ == "__main__":
    test_rag()