zenafey's picture
Update main.py
94db756
raw
history blame
2.46 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, # str in 'Message' Textbox component
api_name="/chat")
return {'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'])