main.py

from sqlalchemy import create_engine
from config import DATABASE_URL

from model_connector import generate_sql_from_nl, format_response
from sql_formatter import execute_sql, extract_db_schema



# Set up the database connection


if __name__ == "__main__":
    engine = create_engine(DATABASE_URL)
    schema_info = extract_db_schema(engine)

    print("Ask a question about the database (type 'quit' to exit):")
    while True:
        user_input = input("User: ")
        if user_input.lower() in ['quit', 'exit']:
            break
        try:
            sql_query = generate_sql_from_nl(user_input, schema_info)
        except Exception as e:
            print("Error generating SQL:", str(e))
            continue
        try:
            results = execute_sql(engine, sql_query)
        except Exception as e:
            print("Error executing SQL:", str(e))
            continue
        answer = format_response(results, sql_query)
        print("Assistant:\n", answer)