sqlgen.py

import time
from groq import Groq
from flask import Flask
from threading import Thread
from utils import GROQ_API_KEY
from postgres_connector import PostgresConnector
from sql_ops import generate_sql, analyze_query_output
from flask_routes import init_routes
from app_state import app_state

# Initialize Groq client
groq_client = Groq(api_key=GROQ_API_KEY)

# Create Flask app
app = Flask(__name__)

# Initialize PostgresConnector
pg = PostgresConnector()

# Initialize routes
init_routes(app, groq_client, pg)

# Flask server runner
def run_flask():
    try:
        print("\nStarting Flask server...")
        app.run(host='0.0.0.0', port=4000)
    except Exception as e:
        print(f"\nError starting Flask server: {str(e)}")
        raise  # Re-raise the exception to see the full traceback

# Main function
def main():
    print("\nWelcome to SQL Generator!")
    print("Please ensure that you have set your `GROQ_API_KEY` environment variable in the `.env` file, or else the agent will not work.")
    print("Your user request to sqlgen will be converted into a SQL query that works in Snowflake.")
    print("When entering your request, be as clear, specific, and concise as possible. This is for best results from the agent.")
    print("\nNOTE: For safety reasons, only SELECT queries are allowed. Any attempts to modify the database will be blocked.")
    
    time.sleep(2)
    print("\nGetting things ready for you in sqlgen...")

    time.sleep(4)
    
    while True:
        print("\n\nPlease go ahead and enter your user request:")
        user_request = input()
    
        print("Generating SQL query...")
        try:
            while True:  # This inner loop handles clarification cycles
                try:
                    sql = generate_sql(user_request, groq_client=groq_client)
                    # print(sql) ### DO NOT UNCOMMENT THIS LINE.
                    print("SQL query generated successfully.")
                    
                    print("Running SQL query...")
                    output = pg.execute_query(sql)
                    print("Analysis of the query output:\n", analyze_query_output(output, groq_client=groq_client))
                    break  # Exit the clarification loop on success
                    
                except Exception as e:
                    error_msg = str(e)
                    if app_state.response_mode == "info_clarification":
                        print("\nI need some clarification to help you better:", error_msg)
                        print("Please provide the requested information:")
                        user_request = input()
                        
                        # Check if we should continue clarification or return to SQL generation
                        response = groq_client.chat.completions.create(
                            model="llama-3.3-70b-versatile",
                            messages=[*app_state.chat_context, {"role": "user", "content": user_request}],
                            temperature=0.5,
                        )
                        
                        if response.choices[0].message.content == "CLARIFICATION_COMPLETE":
                            app_state.set_response_mode("sql_generation")
                            print("\nThanks! Now I'll generate your SQL query...")
                        else:
                            app_state.add_to_chat_context("user", user_request)
                            continue
                    else:
                        print(f"\nError: {error_msg}")
                        print("If you wish, you can try again with a different request.")
                        break
                        
        except Exception as e:
            print(f"\nError: {str(e)}")
            print("If you wish, you can try again with a different request.")

        print("\n\nWould you like to generate and run another SQL query? Note that the agent will remember the previous chat history. (y/n)")
        user_response = input()

        if user_response == "n":
            print("\nExiting CLI interface. API server will continue running in the background.")
            return

if __name__ == "__main__":
    # Start Flask in a separate thread - run the server in the background
    flask_thread = Thread(target=run_flask, daemon=True)
    flask_thread.start()

    time.sleep(2) # This is to prevent the below input from being awaited before the server starts running. Its purpose is to even things out.
    
    print("\nWeb interface is available at http://localhost:4000")
    print("Would you also like to chat with sqlgen using the CLI interface? (y/n):")
    if input() == "y":
        main()
    
    # Keep the program running indefinitely
    print("\nAPI server running. Press Ctrl+C to exit completely.")
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\nShutting down...")
        exit(0)