from flask import Flask, request, jsonify
from huggingface_hub import InferenceClient
from hashlib import sha256
import json
import concurrent.futures

app = Flask(__name__)

hashed_api_key = "6b60036616b90f3b70de654faecc6fd34feb22ab7c8c44628472c772c1a45a41"
client = InferenceClient("mistralai/Mistral-Small-24B-Instruct-2501")

def generate_response(params):
    try:
        response = client.chat_completion(**params)
        # Extract the content from the response
        message_content = response.choices[0].message.content
        return message_content
    except Exception as e:
        return f"An error occurred: {str(e)}"

@app.route('/respond', methods=['POST'])
def respond():
    params = request.json  # Уже Python-объект
    api_key = params.pop("api_key", None)
    if sha256(api_key.encode('utf-8')).hexdigest() != hashed_api_key:
        return jsonify({"error": "Invalid API key"}), 403

    with concurrent.futures.ThreadPoolExecutor() as executor:
        future = executor.submit(generate_response, params)
        result = future.result()
        return jsonify({"response": result})


if __name__ == "__main__":
    print('Ok')
    app.run('0.0.0.0', port=7860)