WebSocket API — Real-Time Streaming Chat via /ws/chat

import asyncio
import json
import websockets

WS_URL = "wss://api.example.com/ws/chat"
TOKEN  = "eyJraWQ..."  # JWT from AWS Cognito

async def chat_session():
    async with websockets.connect(WS_URL) as ws:
        # Step 1: Authenticate
        await ws.send(json.dumps({"type": "auth", "token": TOKEN}))
        auth_response = json.loads(await ws.recv())

        if auth_response["type"] != "auth_success":
            raise RuntimeError(f"Auth failed: {auth_response}")

        user_id = auth_response["user_id"]
        print(f"Authenticated as {user_id}")

        session_id = "session-py-001"
        message    = "¿Cuáles son los ingresos por línea de producto en el año?"

        # Step 2: Send chat message
        await ws.send(json.dumps({
            "type": "chat",
            "message": message,
            "session_id": session_id,
            "confirmed": False,
        }))

        # Step 3: Process the event stream
        final_data = None
        async for raw in ws:
            event = json.loads(raw)
            event_type = event.get("type")

            if event_type == "node_start":
                print(f"  ⏳  {event['label']}")

            elif event_type == "node_end":
                print(f"  ✅  {event['node']} done")

            elif event_type == "confirmation_required":
                print(f"\n  HITL: {event['message']}")
                answer = input("  Confirm? (y/n): ").strip().lower()
                if answer == "y":
                    await ws.send(json.dumps({
                        "type": "chat",
                        "message": message,
                        "session_id": event["session_id"],
                        "confirmed": True,
                    }))
                else:
                    print("  Cancelled.")
                    break

            elif event_type == "response":
                final_data = event["data"]
                print(f"\n{final_data['response']}")
                if final_data.get("sql"):
                    print(f"\nSQL: {final_data['sql']}")
                if final_data.get("kpis"):
                    print(f"KPIs: {final_data['kpis']}")
                break  # Exit after first response; loop for multi-turn

            elif event_type == "error":
                print(f"  Error [{event['error_code']}]: {event['message']}")
                break

        # Step 4: Submit feedback (optional)
        if final_data and final_data.get("query_log_id"):
            await ws.send(json.dumps({
                "type": "feedback",
                "query_id": final_data["query_log_id"],
                "timestamp": final_data["query_timestamp"],
                "feedback": "up",
            }))
            ack = json.loads(await ws.recv())
            print(f"Feedback acknowledged: {ack}")

asyncio.run(chat_session())