import uvicorn

from flask import Flask, request, jsonify
from huggingface_hub import InferenceClient

app = Flask(__name__)

API_URL = "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1"
# API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1"

def format_prompt(message, custom_instructions=None):
    prompt = ""
    if custom_instructions:
        prompt += f"[INST] {custom_instructions} [/INST]"
    prompt += f"[INST] {message} [/INST]"
    return prompt

def Mistral7B(prompt, instructions, api, temperature=0.1, max_new_tokens=1, top_p=0.95, repetition_penalty=1.0):
    global API_URL
    try:
        temperature = float(temperature)
        if temperature < 1e-2:
            temperature = 1e-2
        top_p = float(top_p)

        generate_kwargs = dict(
            temperature=temperature,
            max_new_tokens=max_new_tokens,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            do_sample=True,
            seed=69,
        )
        custom_instructions = instructions
        formatted_prompt = format_prompt(prompt, custom_instructions)

        head = {"Authorization": f"Bearer {api}"}
        client = InferenceClient(API_URL, headers=head)
        response = client.text_generation(formatted_prompt, **generate_kwargs)
        return response
    except Exception as e:
        return str(e)

@app.route("/generate-text", methods=["POST"])
def generate_text():
    data = request.json
    prompt = data.get("prompt")
    instructions = data.get("instructions")
    api_key = data.get("api_key")

    if not prompt or not instructions or not api_key:
        return jsonify({"error": "Missing required fields"}), 400

    response = Mistral7B(prompt, instructions, api_key)

    return jsonify({"response": response}), 200