test_generic_stream.py

import argparse
import sys
from dotenv import load_dotenv
from stream_utils import StreamResponseHandler

# Load environment variables
load_dotenv()

def main():
    # Create a console for rich output
    handler = StreamResponseHandler()
    
    # Parse command line arguments
    parser = argparse.ArgumentParser(description="Test the generic agent streaming API.")
    parser.add_argument("--query", "-q", required=True, help="The query or message to send to the generic agent.")
    
    args = parser.parse_args()
    
    # Check if the server is running
    if not handler.check_server_health():
        sys.exit(1)
    
    # Stream the generic request
    endpoint_url = "http://localhost:8000/test/stream"
    params = {"query": args.query}
    handler.stream_response(endpoint_url, params=params, title="Generic Agent Response")

if __name__ == "__main__":
    main()