app.py

from flask import Flask, render_template, request, redirect, url_for, flash, send_from_directory
from flask_socketio import SocketIO
import os
import threading
from dotenv import load_dotenv
from werkzeug.utils import secure_filename

# LangChain and agent imports
from typing import Annotated, Literal
from langchain_core.messages import AIMessage, ToolMessage
from pydantic import BaseModel, Field
from typing_extensions import TypedDict
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import AnyMessage, add_messages
from langchain_core.runnables import RunnableLambda, RunnableWithFallbacks
from langgraph.prebuilt import ToolNode
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_core.tools import tool
import traceback

# Load environment variables
load_dotenv()

# Global configuration variables
UPLOAD_FOLDER = os.path.join(os.getcwd(), "uploads")
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
BASE_DIR = os.path.abspath(os.path.dirname(__file__))

# API Keys from .env file
os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")
os.environ["MISTRAL_API_KEY"] = os.getenv("MISTRAL_API_KEY")

# Flask and SocketIO setup
flask_app = Flask(__name__)
flask_app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
# Set secret key for flash messages:
flask_app.config['SECRET_KEY'] = os.getenv("FLASK_SECRET_KEY", "mysecretkey")
socketio = SocketIO(flask_app, cors_allowed_origins="*")

# Global state
agent_app = None
abs_file_path = None

def create_agent_app(db_path: str):
    try:
        from langchain_groq import ChatGroq
        llm = ChatGroq(model="llama3-70b-8192")
    except Exception as e:
        flash(f"[ERROR]: Failed to initialize ChatGroq: {e}", "error")
        raise

    abs_db_path = os.path.abspath(db_path)
    try:
        db_instance = SQLDatabase.from_uri(f"sqlite:///{abs_db_path}")
    except Exception as e:
        flash(f"[ERROR]: Failed to connect to DB: {e}", "error")
        raise
    
    @tool
    def db_query_tool(query: str) -> str:
        """
        Execute a SQL query against the database and return the result.
        If the query is invalid or returns no result, an error message will be returned.
        In case of an error, the user is advised to rewrite the query and try again.
        """
        try:
            result = db_instance.run_no_throw(query)
            return result or "Error: Query failed. Please rewrite your query and try again."
        except Exception as e:
            flash(f"[ERROR]: Exception during query execution: {e}", "error")
            return f"Error: {str(e)}"

    class SubmitFinalAnswer(BaseModel):
        final_answer: str = Field(...)

    class State(TypedDict):
        messages: Annotated[list[AnyMessage], add_messages]

    try:
        query_check_system = """You are a SQL expert with a strong attention to detail.
                                Double check the SQLite query for common mistakes, including:
                                - Using NOT IN with NULL values
                                - Using UNION when UNION ALL should have been used
                                - Using BETWEEN for exclusive ranges
                                - Data type mismatch in predicates
                                - Properly quoting identifiers
                                - Using the correct number of arguments for functions
                                - Casting to the correct data type
                                - Using the proper columns for joins
                                
                                If there are any of the above mistakes, rewrite the query. If there are no mistakes, just reproduce the original query.
                                
                                You will call the appropriate tool to execute the query after running this check.
                             """
        
        query_check = ChatPromptTemplate.from_messages([
            ("system", query_check_system),
            ("placeholder", "{messages}")
        ]) | llm.bind_tools([db_query_tool])

        query_gen_system = """You are a SQL expert with a strong attention to detail.

                            Given an input question, output a syntactically correct SQLite query to run, then look at the results of the query and return the answer.
                            
                            DO NOT call any tool besides SubmitFinalAnswer to submit the final answer.
                            
                            When generating the query:
                            
                            Output the SQL query that answers the input question without a tool call.
                            
                            Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 5 results.
                            You can order the results by a relevant column to return the most interesting examples in the database.
                            Never query for all the columns from a specific table, only ask for the relevant columns given the question.
                            
                            If you get an error while executing a query, rewrite the query and try again.
                            
                            If you get an empty result set, you should try to rewrite the query to get a non-empty result set.
                            NEVER make stuff up if you don't have enough information to answer the query... just say you don't have enough information.
                            
                            If you have enough information to answer the input question, simply invoke the appropriate tool to submit the final answer to the user.
                            
                            DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database. Do not return any sql query except answer.
                           """

        
        query_gen = ChatPromptTemplate.from_messages([
            ("system", query_gen_system),
            ("placeholder", "{messages}")
        ]) | llm.bind_tools([SubmitFinalAnswer])
    except Exception as e:
        flash(f"[ERROR]: Failed to create prompt templates: {e}", "error")
        raise
    
    try:
        toolkit = SQLDatabaseToolkit(db=db_instance, llm=llm)
        tools_instance = toolkit.get_tools()
    except Exception as e:
        flash(f"[ERROR]: Failed to initialize SQL toolkit: {e}", "error")
        raise

    def first_tool_call(state: State):
        return {"messages": [AIMessage(content="", tool_calls=[{"name": "sql_db_list_tables", "args": {}, "id": "tool_abcd123"}])]} 

    def handle_tool_error(state: State):
        tool_calls = state["messages"][-1].tool_calls
        return {"messages": [
            ToolMessage(content="Error occurred. Please revise.", tool_call_id=tc["id"]) for tc in tool_calls
        ]}

    def create_tool_node_with_fallback(tools_list):
        return ToolNode(tools_list).with_fallbacks([RunnableLambda(handle_tool_error)], exception_key="error")

    def query_gen_node(state: State):
        try:
            message = query_gen.invoke(state)
        except Exception as e:
            flash(f"[ERROR]: Exception in query_gen_node: {e}", "error")
            raise
        tool_messages = []
        if message.tool_calls:
            for tc in message.tool_calls:
                if tc["name"] != "SubmitFinalAnswer":
                    tool_messages.append(ToolMessage(
                        content=f"Error: Wrong tool called: {tc['name']}",
                        tool_call_id=tc["id"]
                    ))
        return {"messages": [message] + tool_messages}

    def should_continue(state: State):
        last_message = state["messages"][-1]
        if getattr(last_message, "tool_calls", None):
            return END
        if last_message.content.startswith("Error:"):
            return "query_gen"
        return "correct_query"

    def model_check_query(state: State):
        return {"messages": [query_check.invoke({"messages": [state["messages"][-1]]})]}

    list_tool = next((t for t in tools_instance if t.name == "sql_db_list_tables"), None)
    schema_tool = next((t for t in tools_instance if t.name == "sql_db_schema"), None)
    model_get_schema = llm.bind_tools([schema_tool])

    workflow = StateGraph(State)
    workflow.add_node("first_tool_call", first_tool_call)
    workflow.add_node("list_tables_tool", create_tool_node_with_fallback([list_tool]))
    workflow.add_node("get_schema_tool", create_tool_node_with_fallback([schema_tool]))
    # Fixed unterminated string literal:
    workflow.add_node("model_get_schema", lambda s: {"messages": [model_get_schema.invoke(s["messages"])]})
    workflow.add_node("query_gen", query_gen_node)
    workflow.add_node("correct_query", model_check_query)
    workflow.add_node("execute_query", create_tool_node_with_fallback([db_query_tool]))

    workflow.add_edge(START, "first_tool_call")
    workflow.add_edge("first_tool_call", "list_tables_tool")
    workflow.add_edge("list_tables_tool", "model_get_schema")
    workflow.add_edge("model_get_schema", "get_schema_tool")
    workflow.add_edge("get_schema_tool", "query_gen")
    workflow.add_conditional_edges("query_gen", should_continue)
    workflow.add_edge("correct_query", "execute_query")
    workflow.add_edge("execute_query", "query_gen")

    return workflow.compile()

@flask_app.route("/files/<path:filename>")
def uploaded_file(filename):
    try:
        return send_from_directory(flask_app.config['UPLOAD_FOLDER'], filename)
    except Exception as e:
        flash(f"[ERROR]: Could not send file: {str(e)}", "error")
        return redirect(url_for("index"))

# -------------------------------------------------------------------------
# Helper: run_agent runs the agent with the given prompt.
# -------------------------------------------------------------------------
def run_agent(prompt, socketio):
    global agent_app, abs_file_path
    if not abs_file_path:
        socketio.emit("log", {"message": "[ERROR]: No DB file uploaded."})
        socketio.emit("final", {"message": "No database available. Please upload one and try again."})
        flash("No database available. Please upload one and try again.", "error")
        return
    try:
        # Lazy agent initialization: use the previously uploaded DB.
        if agent_app is None:
            socketio.emit("log", {"message": "[INFO]: Initializing agent for the first time..."})
            agent_app = create_agent_app(abs_file_path)
            socketio.emit("log", {"message": "[INFO]: Agent initialized."})
            flash("Agent initialized.", "info")
        query = {"messages": [("user", prompt)]}
        result = agent_app.invoke(query)
        try:
            result = result["messages"][-1].tool_calls[0]["args"]["final_answer"]
        except Exception as e:
            result = "Query failed or no valid answer found."
            flash("Query failed or no valid answer found.", "warning")
        socketio.emit("final", {"message": result})
    except Exception as e:
        error_message = f"Generation failed: {str(e)}"
        socketio.emit("log", {"message": f"[ERROR]: {error_message}"})
        socketio.emit("final", {"message": "Generation failed."})
        flash(error_message, "error")
        traceback.print_exc()

# -------------------------------------------------------------------------
# Route: index page.
# -------------------------------------------------------------------------
@flask_app.route("/")
def index():
    return render_template("index.html")

# -------------------------------------------------------------------------
# Route: generate (POST) – receives a prompt and runs the agent.
# -------------------------------------------------------------------------
@flask_app.route("/generate", methods=["POST"])
def generate():
    try:
        socketio.emit("log", {"message": "[STEP]: Entering query generation..."})
        data = request.json
        prompt = data.get("prompt", "")
        socketio.emit("log", {"message": f"[INFO]: Received prompt: {prompt}"})
        thread = threading.Thread(target=run_agent, args=(prompt, socketio))
        socketio.emit("log", {"message": f"[INFO]: Starting thread: {thread}"})
        thread.start()
        flash("Query submitted successfully.", "info")
        return "OK", 200
    except Exception as e:
        error_message = f"[ERROR]: {str(e)}"
        socketio.emit("log", {"message": error_message})
        flash(error_message, "error")
        return "ERROR", 500

# -------------------------------------------------------------------------
# Route: upload (GET/POST) – handles uploading the SQLite DB file.
# -------------------------------------------------------------------------
@flask_app.route("/upload", methods=["GET", "POST"])
def upload():
    global abs_file_path, agent_app
    try:
        if request.method == "POST":
            file = request.files.get("file")
            if not file:
                flash("No file uploaded.", "error")
                return "No file uploaded", 400
            filename = secure_filename(file.filename)
            if filename.endswith('.db'):
                db_path = os.path.join(flask_app.config['UPLOAD_FOLDER'], "uploaded.db")
                try:
                    file.save(db_path)
                    abs_file_path = os.path.abspath(db_path)  # Save it here; agent init will occur on first query.
                    agent_app = None  # Reset agent on upload.
                    flash(f"Database file '{filename}' uploaded successfully.", "info")
                    socketio.emit("log", {"message": f"[INFO]: Database file '{filename}' uploaded."})
                    return redirect(url_for("index"))
                except Exception as save_err:
                    flash(f"Error saving file: {save_err}", "error")
                    socketio.emit("log", {"message": f"[ERROR]: Error saving file: {save_err}"})
                    return render_template("upload.html")
            else:
                flash("Only .db files are allowed.", "error")
                return render_template("upload.html")
        return render_template("upload.html")
    except Exception as e:
        error_message = f"[ERROR]: {str(e)}"
        flash(error_message, "error")
        socketio.emit("log", {"message": error_message})
        return render_template("upload.html")

@socketio.on("user_input")
def handle_user_input(data):
    prompt = data.get("message")
    if not prompt:
        socketio.emit("log", {"message": "[ERROR]: Empty prompt."})
        flash("Empty prompt.", "error")
        return
    run_agent(prompt, socketio)

# Expose the Flask app as "app" for Gunicorn
app = flask_app

if __name__ == "__main__":
    socketio.run(app, debug=True)