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'])