app.py

from flask import Flask, request, jsonify, render_template, send_from_directory
from werkzeug.utils import secure_filename
import os

from env import device
from reasoner import reasoner
from compress import compress_model, get_model_metadata

app = Flask(__name__, template_folder=".", static_folder="static")

UPLOAD_FOLDER = "uploads"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER

model_context = {
    "path": None,
    "name": None,
    "metadata": None
}

reasoner = reasoner()


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


@app.route("/predict", methods=["POST"])
def predict():
    data = request.get_json()
    user_input = data.get("data", [None])[0]
    if not user_input:
        return jsonify({"error": "No input provided"}), 400

    model_context['device'] = device.type
    response = reasoner.ask(user_input, model_info=model_context)

    return jsonify({"data": [response]})


@app.route("/upload", methods=["POST"])
def upload_model():
    file = request.files["file"]
    filename = secure_filename(file.filename)
    path = os.path.join(app.config["UPLOAD_FOLDER"], filename)
    file.save(path)

    model_context["path"] = path
    model_context["name"] = filename
    model_context["metadata"] = get_model_metadata(path)

    return jsonify({"message": f"Model '{filename}' uploaded", "path": path})

@app.route("/compress", methods=["POST"])
def compress():
    if not model_context.get("path"):
        return jsonify({"error": "Please upload a model file first."}), 400

    data = request.get_json()
    mode = data.get("dtype", "auto")

    try:
        result = compress_model(model_path=model_context["path"], mode=mode)
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": f"Failed to compress model: {e}"}), 500

@app.route("/download/<path:filename>")
def download_file(filename):
    try:
        return send_from_directory(
            directory='.',
            path=filename,
            as_attachment=True
        )
    except FileNotFoundError:
        return "Error: File not found.", 404


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=True)