test_streaming.py

#!/usr/bin/env python3
"""
Simple script to test streaming from the API
This will show you the word-by-word streaming effect
"""
import httpx
import json
import sys


def test_streaming(question: str, model: str = None):
    """Test streaming with visual feedback"""
    url = "http://localhost:8000/chat"

    payload = {"question": question}
    if model:
        payload["model"] = model

    print(f"\n{'='*60}")
    print(f"Question: {question}")
    print(f"{'='*60}\n")
    print("Response: ", end="", flush=True)

    try:
        with httpx.stream("POST", url, json=payload, timeout=120.0) as response:
            for line in response.iter_lines():
                if line.startswith("data: "):
                    try:
                        data = json.loads(line[6:])  # Remove "data: " prefix

                        if data.get("type") == "content_block_delta":
                            # Print each chunk as it arrives (streaming effect!)
                            content = data.get("content", "")
                            print(content, end="", flush=True)

                        elif data.get("type") == "message_stop":
                            print("\n")
                            print(f"{'='*60}")
                            print("✓ Stream completed")

                        elif data.get("type") == "error":
                            print(f"\n✗ Error: {data.get('error')}")

                    except json.JSONDecodeError:
                        pass

    except httpx.ConnectError:
        print("\n✗ Error: Could not connect to the API server.")
        print("Make sure the server is running on http://localhost:8000")
        sys.exit(1)
    except Exception as e:
        print(f"\n✗ Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    # Test with a question that will generate a longer response
    # so you can see the streaming effect clearly

    print("\n🚀 Testing LLM Streaming API")
    print("Watch how the response appears word by word!\n")

    # Example 1: Short response
    test_streaming("What is Python in one sentence?")

    # Example 2: Longer response to see streaming better
    test_streaming("Explain how machine learning works in simple terms with an example")

    # Example 3: Creative task (generates more tokens)
    test_streaming("Write a haiku about programming")