app.py

import json
import logging
import os

from dotenv import load_dotenv
from flask import Flask, current_app, jsonify, render_template, request, session

from config import load_runtime_config
from sql_functions import (
    execute_sql_query,
    generate_sql_query,
    get_default_openai_model,
    validate_sql_query,
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

load_dotenv()

SECRET_KEY_PLACEHOLDERS = {
    "",
    "your_secret_key",
    "replace-with-a-random-secret",
}


class NaturalLanguageQueryError(Exception):
    def __init__(self, message, status_code=400):
        super().__init__(message)
        self.status_code = status_code


def get_required_secret_key():
    secret_key = os.getenv("SECRET_KEY", "").strip()
    if secret_key in SECRET_KEY_PLACEHOLDERS:
        raise RuntimeError(
            "SECRET_KEY is required. Set it in .env to a long random value before starting the app."
        )
    return secret_key


def load_table_definitions(file_path="table_definitions.json"):
    with open(file_path, "r", encoding="utf-8") as definitions_file:
        return json.load(definitions_file)


def run_natural_language_query(question, model, table_definitions):
    try:
        result, tokens_used, openai_duration = generate_sql_query(
            question, model, table_definitions
        )
    except RuntimeError as exc:
        logger.error("Azure OpenAI configuration error.")
        raise NaturalLanguageQueryError(str(exc)) from exc

    if "error" in result:
        raise NaturalLanguageQueryError(result["error"])

    sql_query = result["sql"]
    is_safe_query, validation_error = validate_sql_query(sql_query, table_definitions)
    if not is_safe_query:
        logger.warning("Rejected generated SQL: %s", validation_error)
        raise NaturalLanguageQueryError(validation_error)

    try:
        rows, column_names, sql_execution_duration = execute_sql_query(
            sql_query, table_definitions
        )
    except Exception as exc:
        logger.error("SQL execution failed.")
        raise NaturalLanguageQueryError(
            "An error occurred while executing the SQL query.", status_code=500
        ) from exc

    return {
        "sql_query": sql_query,
        "rows": rows,
        "column_names": column_names,
        "tokens_used": tokens_used,
        "openai_duration": openai_duration,
        "sql_execution_duration": sql_execution_duration,
    }


def create_app(test_config=None, table_definitions=None):
    flask_app = Flask(__name__)
    flask_app.config.from_mapping(load_runtime_config())
    flask_app.config["SECRET_KEY"] = get_required_secret_key()
    if test_config:
        flask_app.config.update(test_config)

    definitions = table_definitions
    if definitions is None:
        definitions = load_table_definitions(flask_app.config["TABLE_DEFINITIONS_PATH"])
    flask_app.extensions["table_definitions"] = definitions

    @flask_app.route("/", methods=["GET", "POST"])
    def index():
        recent_questions = session.setdefault("recent_questions", [])
        prefill_question = ""

        if request.method == "POST":
            prefill_question = request.form.get("recent_question") or request.form.get(
                "question"
            )
            if not prefill_question:
                return render_template(
                    "index.html",
                    recent_questions=recent_questions,
                    error="Please enter a question.",
                    prefill_question=prefill_question,
                    table_definitions=definitions,
                )

            model = request.form.get("model") or current_app.config["DEFAULT_MODEL"]
            if not model:
                try:
                    model = get_default_openai_model()
                except RuntimeError as exc:
                    return render_template("error.html", error_message=str(exc))

            if prefill_question not in recent_questions:
                session["recent_questions"] = (recent_questions + [prefill_question])[-10:]

            try:
                query = run_natural_language_query(prefill_question, model, definitions)
            except NaturalLanguageQueryError as exc:
                return render_template(
                    "error.html",
                    error_message=str(exc),
                    submitted_question=prefill_question,
                    recent_questions=session["recent_questions"],
                )

            return render_template(
                "result.html",
                sql_query=query["sql_query"],
                results=query["rows"],
                column_names=query["column_names"],
                num_rows=len(query["rows"]),
                openai_duration=query["openai_duration"],
                sql_execution_duration=query["sql_execution_duration"],
                tokens_used=query["tokens_used"],
                natural_language_question=prefill_question,
            )

        return render_template(
            "index.html",
            recent_questions=recent_questions,
            prefill_question=prefill_question,
            table_definitions=definitions,
        )

    @flask_app.errorhandler(500)
    def internal_error(error):
        logger.error("Server Error: %s", error)
        return render_template(
            "error.html", error_message="An internal error occurred."
        ), 500

    @flask_app.route("/api/natural_language_query", methods=["POST"])
    def api_natural_language_query():
        data = request.get_json(silent=True)
        if not isinstance(data, dict):
            return jsonify({"error": "Request body must be a valid JSON object."}), 400
        if "question" not in data:
            return jsonify({"error": "Missing 'question' field in request"}), 400

        model = current_app.config["DEFAULT_MODEL"]
        if not model:
            try:
                model = get_default_openai_model()
            except RuntimeError as exc:
                return jsonify({"error": str(exc)}), 400

        try:
            query = run_natural_language_query(data["question"], model, definitions)
        except NaturalLanguageQueryError as exc:
            return jsonify({"error": str(exc)}), exc.status_code

        return jsonify(
            {
                "sql_query": query["sql_query"],
                "tokens_used": query["tokens_used"],
                "openai_duration": query["openai_duration"],
                "sql_execution_duration": query["sql_execution_duration"],
                "results": [
                    dict(zip(query["column_names"], row)) for row in query["rows"]
                ],
            }
        ), 200

    return flask_app


app = create_app()


if __name__ == "__main__":
    app.run(
        debug=app.config["DEBUG"],
        host=app.config["HOST"],
        port=app.config["PORT"],
    )