Jamiiwej2903 commited on
Commit
77006af
·
verified ·
1 Parent(s): 6e9bf3d

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +30 -27
main.py CHANGED
@@ -1,7 +1,7 @@
1
  import asyncio
2
- import aiohttp
 
3
  from fastapi import FastAPI
4
- from sse_starlette.sse import EventSourceResponse
5
  import logging
6
 
7
  logging.basicConfig(level=logging.INFO)
@@ -9,35 +9,38 @@ logger = logging.getLogger(__name__)
9
 
10
  app = FastAPI()
11
 
12
- API_URL = "https://b917-160-179-178-105.ngrok-free.app/test-ai-call"
13
 
14
  call_count = 0
15
 
16
- async def event_generator():
17
  global call_count
18
- async with aiohttp.ClientSession() as session:
19
- while True:
20
- try:
21
- call_count += 1
22
- payload = {"count": call_count}
23
-
24
- async with session.post(API_URL, json=payload) as response:
25
- if response.status == 200:
26
- data = await response.json()
27
- logger.info(f"Call {call_count}: API Response: {data}")
28
- yield data
29
- else:
30
- logger.error(f"Call {call_count}: API call failed with status code: {response.status}")
31
-
32
- except Exception as e:
33
- logger.error(f"Call {call_count}: Error: {e}")
34
-
35
- await asyncio.sleep(10)
36
-
37
- @app.get("/stream")
38
- async def stream(request):
39
- return EventSourceResponse(event_generator())
 
 
 
40
 
41
  @app.get("/")
42
  async def root():
43
- return {"message": "API caller is running", "calls_made": call_count}
 
1
  import asyncio
2
+ import websockets
3
+ import json
4
  from fastapi import FastAPI
 
5
  import logging
6
 
7
  logging.basicConfig(level=logging.INFO)
 
9
 
10
  app = FastAPI()
11
 
12
+ WEBSOCKET_URL = "ws://b917-160-179-178-105.ngrok-free.app/ws"
13
 
14
  call_count = 0
15
 
16
+ async def websocket_client():
17
  global call_count
18
+ while True:
19
+ try:
20
+ async with websockets.connect(WEBSOCKET_URL) as websocket:
21
+ logger.info("WebSocket connection established")
22
+ while True:
23
+ call_count += 1
24
+ payload = {"count": call_count}
25
+ await websocket.send(json.dumps(payload))
26
+ logger.info(f"Sent: {payload}")
27
+
28
+ response = await websocket.recv()
29
+ data = json.loads(response)
30
+ logger.info(f"Received: {data}")
31
+
32
+ await asyncio.sleep(1) # Adjust this delay as needed
33
+ except websockets.exceptions.ConnectionClosed:
34
+ logger.error("WebSocket connection closed. Retrying...")
35
+ await asyncio.sleep(5)
36
+ except Exception as e:
37
+ logger.error(f"Error: {e}")
38
+ await asyncio.sleep(5)
39
+
40
+ @app.on_event("startup")
41
+ async def startup_event():
42
+ asyncio.create_task(websocket_client())
43
 
44
  @app.get("/")
45
  async def root():
46
+ return {"message": "WebSocket client is running", "calls_made": call_count}