from flask import Flask, render_template, request, jsonify
from huggingface_hub import InferenceClient
import random
import os

keys = [
    os.getenv("KEY1"),
    os.getenv("KEY2"), 
    os.getenv("KEY3"), 
    os.getenv("KEY4"), 
    os.getenv("KEY5"),
    ]


app = Flask(__name__)

INITIAL_PROMPT = os.getenv("SYSTEM")


@app.route('/')
def index():
    return render_template('index.html')

@app.route('/chat', methods=['POST'])
def chat():
    client = InferenceClient(
        os.getenv("MODEL_LUCY"),
        token=random.choice(keys)
    )
    user_input = request.json['message']
    conversation_history = request.json['history']

    prompt = INITIAL_PROMPT + conversation_history + f" {user_input}</user>\n<assistant>"

    response = client.text_generation(prompt, max_new_tokens=512, stop_sequences=["<user>"])

    return jsonify({'response': response})

if __name__ == '__main__':
    app.run(debug=True)