scripts/demo_websocket_flow.py

import asyncio
import httpx
import websockets
import json
import sys

API_URL = "http://localhost:8001/api/v1"
WS_URL = "ws://localhost:8001/api/v1"

async def demo_generation():
    print("🎵 AudioForge Real-Time Generation Demo 🎵")
    print("------------------------------------------")

    # 1. Create Generation Request
    prompt = "A fast-paced techno track with a thumping bassline"
    print(f"\n[1] Sending POST request...")
    print(f"    Prompt: '{prompt}'")
    
    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(
                f"{API_URL}/generations/",
                json={"prompt": prompt, "duration": 5}  # Short duration for demo
            )
            response.raise_for_status()
            data = response.json()
            gen_id = data["id"]
            print(f"    ✅ Request Accepted! Generation ID: {gen_id}")
        except Exception as e:
            print(f"    ❌ Failed to create generation: {e}")
            return

    # 2. Connect to WebSocket
    print(f"\n[2] Connecting to WebSocket for updates...")
    ws_endpoint = f"{WS_URL}/ws/generations/{gen_id}"
    
    try:
        async with websockets.connect(ws_endpoint) as websocket:
            print("    ✅ Connected! Waiting for real-time updates...\n")
            print("    [STATUS]      [PROGRESS]    [MESSAGE]")
            print("    -------------------------------------")
            
            async for message in websocket:
                data = json.loads(message)
                
                # Format output nicely
                status = data.get("status", "unknown").upper()
                progress = f"{data.get('progress', 0)}%"
                msg = data.get("message", "")
                
                print(f"    {status:<13} {progress:<13} {msg}")
                
                if data.get("status") in ["completed", "failed"]:
                    if data.get("status") == "completed":
                        print(f"\n    🎉 Generation Complete! Audio URL: {data.get('audio_url')}")
                    else:
                        print(f"\n    ❌ Generation Failed: {data.get('error')}")
                    break
                    
    except Exception as e:
        print(f"\n    ❌ WebSocket Error: {e}")

if __name__ == "__main__":
    try:
        asyncio.run(demo_generation())
    except KeyboardInterrupt:
        print("\nDemo cancelled.")