zenafey's picture
Update main.py
beb58a2
raw
history blame
2.45 kB
from flask import Flask, request, Response, jsonify
import os
from rich import print
import json
import requests
from time import sleep
from gradio_client import Client
settings = {
'node': {
'id':os.environ.get('nodeId'),
'models': os.environ.get('nodeModel')
},
'security': {
'passw':os.environ.get('apipassw')
},
'web': {
'port': os.environ.get('webport', 7860),
'host': os.environ.get('webhost', '0.0.0.0'),
'debug': os.environ.get('webdebug', False)
}
}
app = Flask(__name__)
@app.route("/")
def index():
return f'Hi, its a node {settings["node"]["id"]} with {settings["node"]["models"]}'
@app.route('/send', methods=['POST'])
def send():
message = request.json.get('message', 'hello!')
client = Client("https://ysharma-explore-llamav2-with-tgi.hf.space/")
result = client.predict(message, api_name="/chat")
response = {'response': result}
return jsonify(response)
@app.route("/chat/completions", methods=['POST'])
def chat_completions():
streaming = request.json.get('stream', False)
model = request.json.get('model', 'gpt-4')
messages = request.json.get('messages')
response = ''
try:
response = closeai.ChatCompletion.create(model=model, stream=streaming, messages=messages)
except Exception as er:
print(er)
if '429' in str(er):sleep(30)
def errorStream(er):
yield 'data: %s\n\n' % json.dumps({"status":"!=200","error":str(er)}, separators=(',' ':'))
return app.response_class(errorStream(er), mimetype='text/event-stream')
if not streaming:
return {
'model': model,
'result': response["choices"][0]["message"]["content"]
}
def stream():
for token in response:
completion_data = {
'model': model,
'token': token,
'status':200,
'choices': [
{
'delta': {
'content': token
},
'index': 0,
'finish_reason': None
}
]
}
yield 'data: %s\n\n' % json.dumps(completion_data, separators=(',' ':'))
return app.response_class(stream(), mimetype='text/event-stream')
if __name__ == '__main__':
app.run(**settings['web'])