queue_flask / app.py
multimodalart's picture
Update app.py
836a3ed
import asyncio
import os
from aiohttp import ClientSession
from flask import Flask, request, jsonify
from werkzeug.exceptions import HTTPException
API_ENDPOINT = os.environ.get("API_ENDPOINT", "")
WORKERS = int(os.environ.get("WORKERS", "1"))
app = Flask(__name__)
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
sem = asyncio.Semaphore(WORKERS)
@app.route('/api', methods=['GET', 'POST'])
def handle_request():
headers = {key: value for key, value in request.headers}
data = request.get_data().decode()
method = request.method
response = loop.run_until_complete(
forward_request(headers, data, method))
return jsonify(response), 200
async def forward_request(headers, data, method):
async with sem, ClientSession() as session:
if method == 'POST':
async with session.post(API_ENDPOINT, headers=headers, data=data) as resp:
return await resp.json()
elif method == 'GET':
async with session.get(API_ENDPOINT, headers=headers) as resp:
return await resp.json()
@app.errorhandler(HTTPException)
def handle_exception(e):
response = e.get_response()
response.data = jsonify({"code": e.code, "name": e.name, "description": e.description})
return response
if __name__ == "__main__":
app.run(host='0.0.0.0', port=5000)