import asyncio
import websockets
import json

async def test_ws():
    uri = "ws://localhost:8000/ws/test-session"
    async with websockets.connect(uri) as ws:
        await ws.send(json.dumps({"query": "What is AI?"}))
        response = await ws.recv()
        print("Response:", response)

asyncio.run(test_ws())