Update main.py

main.py

@@ -1,7 +1,7 @@
 import asyncio
-import aiohttp
+import websockets
+import json
 from fastapi import FastAPI
-from sse_starlette.sse import EventSourceResponse
 import logging
 
 logging.basicConfig(level=logging.INFO)
@@ -9,35 +9,38 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
-API_URL = "https://b917-160-179-178-105.ngrok-free.app/test-ai-call"
+WEBSOCKET_URL = "ws://b917-160-179-178-105.ngrok-free.app/ws"
 
 call_count = 0
 
-async def event_generator():
+async def websocket_client():
     global call_count
-    async with aiohttp.ClientSession() as session:
-        while True:
-            try:
-                call_count += 1
-                payload = {"count": call_count}
-
-                async with session.post(API_URL, json=payload) as response:
-                    if response.status == 200:
-                        data = await response.json()
-                        logger.info(f"Call {call_count}: API Response: {data}")
-                        yield data
-                    else:
-                        logger.error(f"Call {call_count}: API call failed with status code: {response.status}")
-
-            except Exception as e:
-                logger.error(f"Call {call_count}: Error: {e}")
-
-            await asyncio.sleep(10)
-
-@app.get("/stream")
-async def stream(request):
-    return EventSourceResponse(event_generator())
+    while True:
+        try:
+            async with websockets.connect(WEBSOCKET_URL) as websocket:
+                logger.info("WebSocket connection established")
+                while True:
+                    call_count += 1
+                    payload = {"count": call_count}
+                    await websocket.send(json.dumps(payload))
+                    logger.info(f"Sent: {payload}")
+
+                    response = await websocket.recv()
+                    data = json.loads(response)
+                    logger.info(f"Received: {data}")
+
+                    await asyncio.sleep(1)  # Adjust this delay as needed
+        except websockets.exceptions.ConnectionClosed:
+            logger.error("WebSocket connection closed. Retrying...")
+            await asyncio.sleep(5)
+        except Exception as e:
+            logger.error(f"Error: {e}")
+            await asyncio.sleep(5)
+
+@app.on_event("startup")
+async def startup_event():
+    asyncio.create_task(websocket_client())
 
 @app.get("/")
 async def root():
-    return {"message": "API caller is running", "calls_made": call_count}
+    return {"message": "WebSocket client is running", "calls_made": call_count}