scripts/validate_rag.py

"""
Automated RAG pipeline validation script.
Tests end-to-end functionality, multi-tenant isolation, and anti-hallucination.
"""
import httpx
import time
import json
from pathlib import Path
from typing import Dict, List, Any, Tuple
import sys

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

BASE_URL = "http://localhost:8000"
TEST_TENANT_A = "tenant_A"
TEST_TENANT_B = "tenant_B"
TEST_USER_A = "user_A"
TEST_USER_B = "user_B"
TEST_KB_A = "kb_A"
TEST_KB_B = "kb_B"

# Test documents
TENANT_A_DOC = Path(__file__).parent.parent / "data" / "test_docs" / "tenant_A_kb.md"
TENANT_B_DOC = Path(__file__).parent.parent / "data" / "test_docs" / "tenant_B_kb.md"

# Test results storage
test_results: List[Dict[str, Any]] = []


def print_header(text: str):
    """Print a formatted header."""
    print("\n" + "=" * 80)
    print(f"  {text}")
    print("=" * 80)


def print_test(test_name: str, passed: bool, reason: str = ""):
    """Print test result."""
    status = "[PASS]" if passed else "[FAIL]"
    print(f"{status} | {test_name}")
    if reason:
        print(f"      └─ {reason}")
    test_results.append({
        "test": test_name,
        "passed": passed,
        "reason": reason
    })


def wait_for_server(max_retries: int = 10, delay: int = 2) -> bool:
    """Wait for the server to be ready."""
    print("Waiting for server to be ready...")
    for i in range(max_retries):
        try:
            response = httpx.get(f"{BASE_URL}/health", timeout=5)
            if response.status_code == 200:
                print("[OK] Server is ready")
                return True
        except Exception:
            pass
        time.sleep(delay)
        print(f"  Retry {i+1}/{max_retries}...")
    print("[FAIL] Server not ready after max retries")
    return False


def upload_document(
    client: httpx.Client,
    file_path: Path,
    tenant_id: str,
    user_id: str,
    kb_id: str
) -> Dict[str, Any]:
    """Upload a document to the knowledge base."""
    try:
        with open(file_path, "rb") as f:
            files = {"file": (file_path.name, f, "text/markdown")}
            data = {
                "tenant_id": tenant_id,
                "user_id": user_id,
                "kb_id": kb_id
            }
            response = client.post(
                f"{BASE_URL}/kb/upload",
                files=files,
                data=data,
                timeout=60
            )
            if response.status_code == 200:
                return {"success": True, "data": response.json()}
            else:
                return {"success": False, "error": response.text}
    except Exception as e:
        return {"success": False, "error": str(e)}


def test_retrieval(
    client: httpx.Client,
    query: str,
    tenant_id: str,
    user_id: str,
    kb_id: str,
    expected_keywords: List[str],
    should_not_contain: List[str] = None,
    top_k: int = 5
) -> Tuple[bool, str]:
    """Test retrieval accuracy."""
    try:
        # Use GET for search endpoint with headers for dev mode auth
        headers = {
            "X-Tenant-Id": tenant_id,
            "X-User-Id": user_id
        }
        response = client.get(
            f"{BASE_URL}/kb/search",
            params={
                "query": query,
                "kb_id": kb_id,
                "top_k": top_k
            },
            headers=headers,
            timeout=30
        )
        
        if response.status_code != 200:
            return False, f"API returned {response.status_code}: {response.text}"
        
        data = response.json()
        results = data.get("results", [])
        
        if not results:
            return False, "No results retrieved"
        
        # Check tenant isolation
        for result in results:
            metadata = result.get("metadata", {})
            result_tenant = metadata.get("tenant_id")
            if result_tenant != tenant_id:
                return False, f"Tenant leak detected! Got tenant_id={result_tenant}, expected {tenant_id}"
        
        # Check for expected keywords
        all_content = " ".join([r.get("content", "") for r in results]).lower()
        found_keywords = [kw for kw in expected_keywords if kw.lower() in all_content]
        
        if not found_keywords:
            return False, f"Expected keywords not found: {expected_keywords}"
        
        # Check for forbidden content
        if should_not_contain:
            for forbidden in should_not_contain:
                if forbidden.lower() in all_content:
                    return False, f"Forbidden content found: {forbidden}"
        
        return True, f"Retrieved {len(results)} results, found keywords: {found_keywords}"
        
    except Exception as e:
        return False, f"Error: {str(e)}"


def test_chat(
    client: httpx.Client,
    question: str,
    tenant_id: str,
    user_id: str,
    kb_id: str,
    expected_keywords: List[str] = None,
    should_refuse: bool = False,
    should_not_contain: List[str] = None
) -> Tuple[bool, str, Dict[str, Any]]:
    """Test full chat endpoint."""
    try:
        # Include headers for dev mode auth
        headers = {
            "X-Tenant-Id": tenant_id,
            "X-User-Id": user_id
        }
        response = client.post(
            f"{BASE_URL}/chat",
            json={
                "tenant_id": tenant_id,
                "user_id": user_id,
                "kb_id": kb_id,
                "question": question
            },
            headers=headers,
            timeout=60
        )
        
        if response.status_code != 200:
            return False, f"API returned {response.status_code}: {response.text}", {}
        
        data = response.json()
        answer = data.get("answer", "").lower()
        citations = data.get("citations", [])
        from_kb = data.get("from_knowledge_base", False)
        confidence = data.get("confidence", 0.0)
        metadata = data.get("metadata", {})
        refused = metadata.get("refused", False)
        
        # Check refusal behavior (STRICT)
        if should_refuse:
            # Check if response explicitly indicates refusal
            refused = data.get("refused", False)
            refusal_keywords = [
                "couldn't find", "don't have", "not available", "contact support",
                "not in the knowledge base", "could not verify", "not enough information",
                "apologize", "couldn't find relevant information"
            ]
            has_refusal_keywords = any(kw in answer for kw in refusal_keywords)
            
            # If answer was generated with citations, it's a FAIL (should have refused)
            if citations and len(citations) > 0:
                return False, (
                    f"Should have refused but generated answer with {len(citations)} citations. "
                    f"Answer: {answer[:300]}"
                ), data
            
            # If confidence is high and answer exists, it's a FAIL
            if confidence >= 0.30 and answer and not has_refusal_keywords:
                return False, (
                    f"Should have refused but generated answer with confidence {confidence:.2f}. "
                    f"Answer: {answer[:300]}"
                ), data
            
            # If not refused and no refusal keywords, it's a FAIL
            if not refused and not has_refusal_keywords:
                return False, (
                    f"Should have refused but didn't. "
                    f"refused={refused}, confidence={confidence:.2f}, citations={len(citations)}. "
                    f"Answer: {answer[:300]}"
                ), data
            
            # If we got here, it properly refused
            return True, f"Properly refused (refused={refused}, confidence={confidence:.2f})", data
        
        # Check for expected keywords
        if expected_keywords:
            found = [kw for kw in expected_keywords if kw.lower() in answer]
            if not found:
                return False, f"Expected keywords not found: {expected_keywords}. Answer: {answer[:200]}", data
        
        # Check citations
        if not should_refuse and from_kb:
            if not citations:
                return False, "Answer claims to be from KB but has no citations", data
        
        # Check for forbidden content
        if should_not_contain:
            for forbidden in should_not_contain:
                if forbidden.lower() in answer:
                    return False, f"Forbidden content found in answer: {forbidden}", data
        
        # Check citation integrity
        if citations and expected_keywords:
            citation_text = " ".join([c.get("excerpt", "") for c in citations]).lower()
            for kw in expected_keywords:
                if kw.lower() in answer and kw.lower() not in citation_text:
                    # This is a warning, not a failure
                    pass
        
        return True, f"Answer generated (confidence: {confidence:.2f}, citations: {len(citations)})", data
        
    except Exception as e:
        return False, f"Error: {str(e)}", {}


def main():
    """Run all validation tests."""
    print_header("RAG Pipeline Validation Suite")
    
    # Check server
    if not wait_for_server():
        print("[FAIL] Cannot proceed without server")
        return
    
    client = httpx.Client(timeout=120.0)
    
    # ========== PHASE 1: Upload Documents ==========
    print_header("Phase 1: Upload Test Documents")
    
    # Upload tenant A doc
    print(f"\n📤 Uploading {TENANT_A_DOC.name} for {TEST_TENANT_A}...")
    result = upload_document(client, TENANT_A_DOC, TEST_TENANT_A, TEST_USER_A, TEST_KB_A)
    if result["success"]:
        print("[OK] Upload successful")
        print("⏳ Waiting for document processing (10 seconds)...")
        time.sleep(10)  # Wait longer for processing (parsing, chunking, embedding)
    else:
        print(f"[FAIL] Upload failed: {result.get('error')}")
        return
    
    # Upload tenant B doc
    print(f"\n📤 Uploading {TENANT_B_DOC.name} for {TEST_TENANT_B}...")
    result = upload_document(client, TENANT_B_DOC, TEST_TENANT_B, TEST_USER_B, TEST_KB_B)
    if result["success"]:
        print("[OK] Upload successful")
        print("⏳ Waiting for document processing (10 seconds)...")
        time.sleep(10)  # Wait longer for processing (parsing, chunking, embedding)
    else:
        print(f"[FAIL] Upload failed: {result.get('error')}")
        return
    
    # ========== PHASE 2: Retrieval Tests ==========
    print_header("Phase 2: Retrieval Accuracy Tests")
    
    # Test 1: Tenant A retrieval
    passed, reason = test_retrieval(
        client,
        "What is the refund window?",
        TEST_TENANT_A,
        TEST_USER_A,
        TEST_KB_A,
        expected_keywords=["7 days"],
        should_not_contain=["30 days"]
    )
    print_test("Retrieval: Tenant A - Refund Window", passed, reason)
    
    # Test 2: Tenant B retrieval
    passed, reason = test_retrieval(
        client,
        "What is the refund window?",
        TEST_TENANT_B,
        TEST_USER_B,
        TEST_KB_B,
        expected_keywords=["30 days"],
        should_not_contain=["7 days"]
    )
    print_test("Retrieval: Tenant B - Refund Window", passed, reason)
    
    # Test 3: Tenant isolation (A should not get B's data)
    passed, reason = test_retrieval(
        client,
        "Starter plan price",
        TEST_TENANT_A,
        TEST_USER_A,
        TEST_KB_A,
        expected_keywords=["499"],
        should_not_contain=["999"]
    )
    print_test("Retrieval: Tenant A - Starter Plan Price (Isolation)", passed, reason)
    
    # Test 4: Tenant isolation (B should not get A's data)
    passed, reason = test_retrieval(
        client,
        "Starter plan price",
        TEST_TENANT_B,
        TEST_USER_B,
        TEST_KB_B,
        expected_keywords=["999"],
        should_not_contain=["499"]
    )
    print_test("Retrieval: Tenant B - Starter Plan Price (Isolation)", passed, reason)
    
    # ========== PHASE 3: Chat Tests ==========
    print_header("Phase 3: Chat Endpoint Tests")
    
    # Test 5: Tenant A chat - refund window
    passed, reason, data = test_chat(
        client,
        "What is the refund window?",
        TEST_TENANT_A,
        TEST_USER_A,
        TEST_KB_A,
        expected_keywords=["7 days"],
        should_not_contain=["30 days"]
    )
    print_test("Chat: Tenant A - Refund Window", passed, reason)
    
    # Test 6: Tenant B chat - refund window
    passed, reason, data = test_chat(
        client,
        "What is the refund window?",
        TEST_TENANT_B,
        TEST_USER_B,
        TEST_KB_B,
        expected_keywords=["30 days"],
        should_not_contain=["7 days"]
    )
    print_test("Chat: Tenant B - Refund Window", passed, reason)
    
    # Test 7: Tenant A chat - Starter plan
    passed, reason, data = test_chat(
        client,
        "What is the Starter plan price?",
        TEST_TENANT_A,
        TEST_USER_A,
        TEST_KB_A,
        expected_keywords=["499"],
        should_not_contain=["999"]
    )
    print_test("Chat: Tenant A - Starter Plan Price", passed, reason)
    
    # Test 8: Tenant B chat - Starter plan
    passed, reason, data = test_chat(
        client,
        "What is the Starter plan price?",
        TEST_TENANT_B,
        TEST_USER_B,
        TEST_KB_B,
        expected_keywords=["999"],
        should_not_contain=["499"]
    )
    print_test("Chat: Tenant B - Starter Plan Price", passed, reason)
    
    # Test 9: Hallucination refusal - out of scope
    passed, reason, data = test_chat(
        client,
        "How to integrate ClientSphere with Shopify?",
        TEST_TENANT_A,
        TEST_USER_A,
        TEST_KB_A,
        should_refuse=True
    )
    print_test("Chat: Hallucination Refusal (Out of Scope)", passed, reason)
    
    # Test 10: Citation integrity
    passed, reason, data = test_chat(
        client,
        "How long do password reset links last?",
        TEST_TENANT_A,
        TEST_USER_A,
        TEST_KB_A,
        expected_keywords=["15"]
    )
    if passed:
        citations = data.get("citations", [])
        if citations:
            print_test("Chat: Citation Integrity", True, f"Found {len(citations)} citations")
        else:
            print_test("Chat: Citation Integrity", False, "No citations provided")
    else:
        print_test("Chat: Citation Integrity", False, reason)
    
    # ========== PHASE 4: Summary ==========
    print_header("Test Summary")
    
    total_tests = len(test_results)
    passed_tests = sum(1 for r in test_results if r["passed"])
    failed_tests = total_tests - passed_tests
    
    print(f"\nTotal Tests: {total_tests}")
    print(f"[PASS] Passed: {passed_tests}")
    print(f"[FAIL] Failed: {failed_tests}")
    print(f"Success Rate: {(passed_tests/total_tests*100):.1f}%")
    
    if failed_tests > 0:
        print("\n[FAIL] Failed Tests:")
        for result in test_results:
            if not result["passed"]:
                print(f"  - {result['test']}: {result['reason']}")
    
    # Final verdict
    print_header("Final Verdict")
    if failed_tests == 0:
        print("[PASS] ALL TESTS PASSED - RAG Pipeline is working correctly")
        return 0
    else:
        print(f"[FAIL] {failed_tests} TEST(S) FAILED - Review issues above")
        return 1


if __name__ == "__main__":
    exit_code = main()
    sys.exit(exit_code)