RAG/test_client.py

"""
Test client for RAG API
"""
import requests
import json
import sys


def test_health():
    """Test health check endpoint"""
    print("Testing health check...")
    try:
        response = requests.get("http://localhost:8001/health")
        print(f"Status: {response.status_code}")
        print(f"Response: {response.json()}")
        return response.status_code == 200
    except Exception as e:
        print(f"Error: {e}")
        return False


def test_query_stream(query: str):
    """Test streaming query"""
    print(f"\nTesting streaming query: {query}")
    print("-" * 50)
    try:
        url = "http://localhost:8001/query"
        data = {
            "query": query,
            "stream": True,
            "top_k": 5
        }
        
        # Use stream=True and increase timeout for long-running queries
        # Also set stream=True to get streaming response
        response = requests.post(
            url, 
            json=data, 
            stream=True, 
            timeout=(10, 300)  # (connect timeout, read timeout) - 10s connect, 300s read
        )
        
        if response.status_code != 200:
            print(f"Error: {response.status_code}")
            try:
                print(response.text)
            except:
                pass
            return False
        
        print("Response (streaming):")
        # Use iter_content for streaming text response
        # chunk_size=1 to get characters as they arrive
        # decode_unicode=True to handle UTF-8 properly
        try:
            for chunk in response.iter_content(chunk_size=1, decode_unicode=True):
                if chunk:
                    print(chunk, end='', flush=True)
        except requests.exceptions.ChunkedEncodingError as e:
            print(f"\nStreaming error (connection may have closed): {e}")
            return False
        print("\n" + "-" * 50)
        return True
    except requests.exceptions.Timeout as e:
        print(f"Timeout error: {e}")
        print("The query may be taking too long. Try increasing the timeout or checking the server logs.")
        return False
    except Exception as e:
        print(f"Error: {e}")
        return False


def test_query(query: str):
    """Test non-streaming query"""
    print(f"\nTesting non-streaming query: {query}")
    print("-" * 50)
    try:
        url = "http://localhost:8001/query"
        data = {
            "query": query,
            "stream": False,
            "top_k": 5
        }
        
        # Increase timeout for non-streaming queries
        # LLM request_timeout is 120s, so we need at least that plus overhead
        response = requests.post(
            url, 
            json=data, 
            timeout=(10, 300)  # (connect timeout, read timeout) - 10s connect, 300s read
        )
        
        if response.status_code != 200:
            print(f"Error: {response.status_code}")
            try:
                print(response.text)
            except:
                pass
            return False
        
        result = response.json()
        print(f"Response: {result.get('response', 'No response')}")
        print("-" * 50)
        return True
    except requests.exceptions.Timeout as e:
        print(f"Timeout error: {e}")
        print("The query may be taking too long. Try increasing the timeout or checking the server logs.")
        print("You can also try using streaming mode (stream=True) for better progress visibility.")
        return False
    except Exception as e:
        print(f"Error: {e}")
        return False


def test_retrieve(query: str):
    """Test retrieve endpoint (only retrieval, no LLM generation)"""
    print(f"\nTesting retrieve endpoint: {query}")
    print("-" * 50)
    try:
        url = "http://localhost:8001/retrieve"
        data = {
            "query": query,
            "top_k": 5
        }
        
        response = requests.post(
            url, 
            json=data, 
            timeout=(10, 60)  # Retrieval is faster, shorter timeout
        )
        
        if response.status_code != 200:
            print(f"Error: {response.status_code}")
            try:
                print(response.text)
            except:
                pass
            return False
        
        result = response.json()
        print(f"Query: {result.get('query', 'No query')}")
        print(f"Retrieved {result.get('count', 0)} documents:")
        print("-" * 50)
        
        for i, doc in enumerate(result.get('documents', []), 1):
            print(f"\nDocument {i}:")
            print(f"  Score: {doc.get('score', 'N/A')}")
            print(f"  Doc ID: {doc.get('doc_id', 'N/A')}")
            print(f"  Content (first 200 chars): {doc.get('content', '')[:200]}...")
            if doc.get('metadata'):
                print(f"  Metadata keys: {list(doc.get('metadata', {}).keys())}")
        
        print("-" * 50)
        return True
    except requests.exceptions.Timeout as e:
        print(f"Timeout error: {e}")
        return False
    except Exception as e:
        print(f"Error: {e}")
        return False


def test_stats():
    """Test stats endpoint"""
    print("\nTesting stats endpoint...")
    try:
        response = requests.get("http://localhost:8001/stats")
        print(f"Status: {response.status_code}")
        print(f"Response: {json.dumps(response.json(), indent=2, ensure_ascii=False)}")
        return response.status_code == 200
    except Exception as e:
        print(f"Error: {e}")
        return False


def test_sync(full_sync: bool = False):
    """Test sync endpoint"""
    print(f"\nTesting sync endpoint (full_sync={full_sync})...")
    try:
        url = "http://localhost:8001/sync"
        data = {"full_sync": full_sync}
        
        response = requests.post(url, json=data, timeout=300)
        
        if response.status_code != 200:
            print(f"Error: {response.status_code}")
            print(response.text)
            return False
        
        result = response.json()
        print(f"Status: {result.get('status', 'unknown')}")
        print(f"Message: {result.get('message', 'No message')}")
        return True
    except Exception as e:
        print(f"Error: {e}")
        return False


def main():
    """Run all tests"""
    print("=" * 50)
    print("RAG API Test Client")
    print("=" * 50)
    
    # Test health
    if not test_health():
        print("\n❌ Health check failed. Is the server running?")
        sys.exit(1)
    
    # Test stats
    test_stats()
    
    # Test sync (optional)
    if len(sys.argv) > 1 and sys.argv[1] == "--sync":
        test_sync(full_sync=True)
    
    # Test queries
    test_queries = [
        "a07f26ffa7e8234fa097591d776244b3",
        "这个系统是如何工作的？"
    ]
    
    for query in test_queries:
        # # Test retrieve (only retrieval, no LLM)
        test_retrieve(query)
        
        # Test streaming
        # test_query_stream(query)
        
        # Test non-streaming (with LLM generation)
        # test_query(query)
    
    print("\n" + "=" * 50)
    print("Tests completed!")
    print("=" * 50)


if __name__ == "__main__":
    main()