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()