zenafey's picture
Update main.py
2a97875
raw
history blame
2.45 kB
from flask import Flask, request, Response
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")
print('pon')
return app.response_class({'response': result})
@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'])