multimodalart HF Staff commited on
Commit
836a3ed
·
1 Parent(s): 1a71614

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -39
app.py CHANGED
@@ -1,57 +1,45 @@
1
- from flask import Flask, request, jsonify
2
- from queue import Queue
3
- from threading import Thread
4
  import os
5
- import time
6
- import requests
7
-
8
- app = Flask(__name__)
9
 
10
- queue = Queue()
11
  API_ENDPOINT = os.environ.get("API_ENDPOINT", "")
12
  WORKERS = int(os.environ.get("WORKERS", "1"))
13
 
14
- def worker():
15
- while True:
16
- req = queue.get()
17
- if req is None:
18
- break
19
 
20
- # Forward the request to the API endpoint
21
- response = requests.request(req['method'], API_ENDPOINT, headers=req['headers'], data=req['data'])
22
- req['result'] = response.json()
23
- queue.task_done()
24
 
25
  @app.route('/api', methods=['GET', 'POST'])
26
  def handle_request():
27
  headers = {key: value for key, value in request.headers}
28
- data = request.get_data()
29
  method = request.method
30
 
31
- # Add the request to the queue
32
- req = {'headers': headers, 'data': data, 'method': method, 'result': None}
33
- queue.put(req)
 
34
 
35
- # Wait until the request has been processed
36
- while req['result'] is None:
37
- time.sleep(0.1)
 
 
 
 
 
38
 
39
- return jsonify(req['result']), 200
 
 
 
 
 
40
 
41
 
42
  if __name__ == "__main__":
43
- # Start the worker threads
44
- threads = []
45
- for i in range(WORKERS):
46
- t = Thread(target=worker)
47
- t.start()
48
- threads.append(t)
49
-
50
- # Start the Flask app
51
  app.run(host='0.0.0.0', port=5000)
52
-
53
- # Stop the worker threads
54
- for i in range(WORKERS):
55
- queue.put(None)
56
- for t in threads:
57
- t.join()
 
1
+ import asyncio
 
 
2
  import os
3
+ from aiohttp import ClientSession
4
+ from flask import Flask, request, jsonify
5
+ from werkzeug.exceptions import HTTPException
 
6
 
 
7
  API_ENDPOINT = os.environ.get("API_ENDPOINT", "")
8
  WORKERS = int(os.environ.get("WORKERS", "1"))
9
 
10
+ app = Flask(__name__)
11
+ loop = asyncio.new_event_loop()
12
+ asyncio.set_event_loop(loop)
13
+ sem = asyncio.Semaphore(WORKERS)
 
14
 
 
 
 
 
15
 
16
  @app.route('/api', methods=['GET', 'POST'])
17
  def handle_request():
18
  headers = {key: value for key, value in request.headers}
19
+ data = request.get_data().decode()
20
  method = request.method
21
 
22
+ response = loop.run_until_complete(
23
+ forward_request(headers, data, method))
24
+ return jsonify(response), 200
25
+
26
 
27
+ async def forward_request(headers, data, method):
28
+ async with sem, ClientSession() as session:
29
+ if method == 'POST':
30
+ async with session.post(API_ENDPOINT, headers=headers, data=data) as resp:
31
+ return await resp.json()
32
+ elif method == 'GET':
33
+ async with session.get(API_ENDPOINT, headers=headers) as resp:
34
+ return await resp.json()
35
 
36
+
37
+ @app.errorhandler(HTTPException)
38
+ def handle_exception(e):
39
+ response = e.get_response()
40
+ response.data = jsonify({"code": e.code, "name": e.name, "description": e.description})
41
+ return response
42
 
43
 
44
  if __name__ == "__main__":
 
 
 
 
 
 
 
 
45
  app.run(host='0.0.0.0', port=5000)