zenafey's picture
Update main.py
6c2deee
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
import subprocess
subprocess.Popen(["python", "bot.py"])
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'])