-
Notifications
You must be signed in to change notification settings - Fork 22
/
tr_ara_rtx-flask.py
58 lines (48 loc) · 1.9 KB
/
tr_ara_rtx-flask.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
from flask import Flask, request, url_for, jsonify, abort
import requests, sys, threading, time, json
import traceback
ARS_API = 'http://localhost:8000/ars/api'
RTX_URL = 'https://arax.rtx.ai/api/rtx/v1/query'
def setup_app():
DEFAULT_ACTOR = {
'channel': 'general',
'agent': {
'name': 'ara-rtx-agent',
'uri': 'http://localhost:5000'
},
'path': '/rtxquery' # relative to agent's uri
}
# wait for flask to finish intializing before it can accept
# connections from the ars
time.sleep(2)
r = requests.post(ARS_API+'/actors', json=DEFAULT_ACTOR)
if r.status_code != 201 and r.status_code != 302:
app.logger.error('Unable to intialize actor; %s return status %d'
% (r.url, r.status_code))
sys.exit(1)
print('initializing %s...%d\n%s' % (__name__, r.status_code, r.text))
app = Flask(__name__)
threading.Thread(target=setup_app).start()
@app.route('/', methods=['GET'])
def index():
return 'A wrapper ARA for RTX!'
@app.route('/rtxquery', methods=['POST'])
def rtxquery():
data = request.get_json()
if 'model' not in data or data['model'] != 'tr_ars.message':
return abort(400)
app.logger.debug('%s: received message...%s' % (request.url, data['pk']))
mesg = data['fields']
if 'ref' not in mesg or mesg['ref'] != None:
# this is not a head message, so we're not interested
return abort(400)
try:
data = json.loads(mesg['data'])
r = requests.post(RTX_URL, json=data, timeout=60)
return (jsonify(r.text),
r.status_code, # return status code
{'tr_ars.message.status': 'D'}) # set the status of the message
except Exception as e:
app.logger.error("RTX failed: {}".format(traceback.format_exception(type(e), e, e.__traceback__)))
if __name__ == '__main__':
app.run()