Update app.py

app.py

@@ -3,7 +3,8 @@ from flask_socketio import SocketIO
 import threading
 import os
 from dotenv import load_dotenv
-import sqlite3  
+import sqlite3
+from werkzeug.utils import secure_filename
 
 # LangChain and agent imports
 from langchain_community.chat_models.huggingface import ChatHuggingFace  # if needed later
@@ -19,7 +20,7 @@ from langchain import hub
 from langchain_community.tools import DuckDuckGoSearchRun
 
 # Agent requirements and type hints
-from typing import Annotated, Literal, Sequence, TypedDict, Any
+from typing import Annotated, Literal, TypedDict, Any
 from langchain_core.messages import AIMessage, ToolMessage
 from pydantic import BaseModel, Field
 from typing_extensions import TypedDict
@@ -31,284 +32,269 @@ from langgraph.prebuilt import ToolNode
 # Load environment variables
 load_dotenv()
 
-UPLOAD_FOLDER = "uploads/"
-
-# In your .env file, ensure you have:
-# DATABASE_URI=sqlite:///employee.db
+# Global configuration variables
+UPLOAD_FOLDER = os.path.join(os.getcwd(), "uploads")
 BASE_DIR = os.path.abspath(os.path.dirname(__file__))
 DATABASE_URI = f"sqlite:///{os.path.join(BASE_DIR, 'data', 'mydb.db')}"
 print("DATABASE URI:", DATABASE_URI)
 
-
+# API Keys from .env file
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
 os.environ["GROQ_API_KEY"] = GROQ_API_KEY
 os.environ["MISTRAL_API_KEY"] = MISTRAL_API_KEY
 
-# Use ChatGroq LLM (which does not require a Hugging Face API token)
-from langchain_groq import ChatGroq
-from langchain_mistralai.chat_models import ChatMistralAI
-
-###############################################################################
-# Application Factory: create_app()
-#
-# This function sets up the Flask application, SocketIO, routes, and initializes
-# the global agent_app using the default DATABASE_URI. It returns the Flask app.
-###############################################################################
-# --- Application Factory ---
+# Global variables for dynamic agent and DB file path; initially None.
+agent_app = None
 abs_file_path = None
-agent_app= None
 
-def create_app():
+# =============================================================================
+# create_agent_app: Given a database path, initialize the agent workflow.
+# =============================================================================
+def create_agent_app(db_path: str):
+    # Use ChatGroq as our LLM here; you can swap to ChatMistralAI if preferred.
+    from langchain_groq import ChatGroq
+    llm = ChatGroq(model="llama3-70b-8192")
+
+    # -------------------------------------------------------------------------
+    # Define a tool for executing SQL queries.
+    # -------------------------------------------------------------------------
+    @tool
+    def db_query_tool(query: str) -> str:
+        result = db_instance.run_no_throw(query)
+        return result if result else "Error: Query failed. Please rewrite your query and try again."
+
+    # -------------------------------------------------------------------------
+    # Pydantic model for final answer
+    # -------------------------------------------------------------------------
+    class SubmitFinalAnswer(BaseModel):
+        final_answer: str = Field(..., description="The final answer to the user")
+
+    # -------------------------------------------------------------------------
+    # Define state type for our workflow.
+    # -------------------------------------------------------------------------
+    class State(TypedDict):
+        messages: Annotated[list[AnyMessage], add_messages]
+
+    # -------------------------------------------------------------------------
+    # Set up prompt templates (using langchain_core.prompts) for query checking
+    # and query generation.
+    # -------------------------------------------------------------------------
+    from langchain_core.prompts import ChatPromptTemplate
+
+    query_check_system = (
+        "You are a SQL expert with a strong attention to detail.\n"
+        "Double check the SQLite query for common mistakes, including:\n"
+        "- Using NOT IN with NULL values\n"
+        "- Using UNION when UNION ALL should have been used\n"
+        "- Using BETWEEN for exclusive ranges\n"
+        "- Data type mismatch in predicates\n"
+        "- Properly quoting identifiers\n"
+        "- Using the correct number of arguments for functions\n"
+        "- Casting to the correct data type\n"
+        "- Using the proper columns for joins\n\n"
+        "If there are any of the above mistakes, rewrite the query. If there are no mistakes, just reproduce the original query.\n"
+        "You will call the appropriate tool to execute the query after running this check."
+    )
+    query_check_prompt = ChatPromptTemplate.from_messages([
+        ("system", query_check_system),
+        ("placeholder", "{messages}")
+    ])
+    query_check = query_check_prompt | llm.bind_tools([db_query_tool])
+
+    query_gen_system = (
+        "You are a SQL expert with a strong attention to detail.\n\n"
+        "Given an input question, output a syntactically correct SQLite query to run, then look at the results of the query and return the answer.\n\n"
+        "DO NOT call any tool besides SubmitFinalAnswer to submit the final answer.\n\n"
+        "When generating the query:\n"
+        "Output the SQL query that answers the input question without a tool call.\n"
+        "Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 5 results.\n"
+        "You can order the results by a relevant column to return the most interesting examples in the database.\n"
+        "Never query for all the columns from a specific table, only ask for the relevant columns given the question.\n\n"
+        "If you get an error while executing a query, rewrite the query and try again.\n"
+        "If you get an empty result set, you should try to rewrite the query to get a non-empty result set.\n"
+        "NEVER make stuff up if you don't have enough information to answer the query... just say you don't have enough information.\n\n"
+        "If you have enough information to answer the input question, simply invoke the appropriate tool to submit the final answer to the user.\n"
+        "DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database. Do not return any sql query except answer."
+    )
+    query_gen_prompt = ChatPromptTemplate.from_messages([
+        ("system", query_gen_system),
+        ("placeholder", "{messages}")
+    ])
+    query_gen = query_gen_prompt | llm.bind_tools([SubmitFinalAnswer])
+
+    # Update database URI and file path
+    abs_db_path_local = os.path.abspath(db_path)
+    global DATABASE_URI
+    DATABASE_URI = abs_db_path_local
+    db_uri = f"sqlite:///{abs_db_path_local}"
+    print("db_uri", db_uri)
+
+    # Create SQLDatabase connection using langchain utility.
+    from langchain_community.utilities import SQLDatabase
+    db_instance = SQLDatabase.from_uri(db_uri)
+    print("db_instance----->", db_instance)
+
+    # Create SQL toolkit.
+    from langchain_community.agent_toolkits import SQLDatabaseToolkit
+    toolkit_instance = SQLDatabaseToolkit(db=db_instance, llm=llm)
+    tools_instance = toolkit_instance.get_tools()
+
+    # Define workflow nodes and fallback functions.
+    def first_tool_call(state: State) -> dict[str, list[AIMessage]]:
+        return {"messages": [AIMessage(content="", tool_calls=[{"name": "sql_db_list_tables", "args": {}, "id": "tool_abcd123"}])]}
 
-    def create_agent_app(db_path: str):
-        # Construct the SQLite URI from the given file path.
-        # Ensure the db_path is absolute so that SQLAlchemy can locate the file.
-        #llm = ChatMistralAI(model="mistral-large-latest")
-        llm = ChatGroq(model="llama3-70b-8192")
-        
-        @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.
-            """
-            result = db_instance.run_no_throw(query)
-            if not result:
-                return "Error: Query failed. Please rewrite your query and try again."
-            return result
-        
-        # Define a Pydantic model for submitting the final answer
-        class SubmitFinalAnswer(BaseModel):
-            """Submit the final answer to the user based on the query results."""
-            final_answer: str = Field(..., description="The final answer to the user")
-        
-        # Define the state type
-        class State(TypedDict):
-            messages: Annotated[list[AnyMessage], add_messages]
-        
-        # Define prompt templates for query checking and query generation
-        from langchain_core.prompts import ChatPromptTemplate
-        
-        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_prompt = ChatPromptTemplate.from_messages([("system", query_check_system), ("placeholder", "{messages}")])
-        query_check = query_check_prompt | 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_prompt = ChatPromptTemplate.from_messages([("system", query_gen_system), ("placeholder", "{messages}")])
-        query_gen = query_gen_prompt | llm.bind_tools([SubmitFinalAnswer])
-        
-        abs_db_path = os.path.abspath(db_path)
-        global DATABASE_URI
-        DATABASE_URI = abs_db_path    
-        db_uri = f"sqlite:///{abs_db_path}"
-        print("db_uri",db_uri)
-        
-        # Create new SQLDatabase connection using the constructed URI.
-        from langchain_community.utilities import SQLDatabase
-        db_instance = SQLDatabase.from_uri(db_uri)
-        print("db_instance----->",db_instance)
-        print("db_uri----->",db_uri)
-    
-    
-        # Create SQL toolkit and get the tools.
-        from langchain_community.agent_toolkits import SQLDatabaseToolkit
-        toolkit_instance = SQLDatabaseToolkit(db=db_instance, llm=llm)
-        tools_instance = toolkit_instance.get_tools()
-    
-        # Define workflow nodes and fallback functions
-        def first_tool_call(state: State) -> dict[str, list[AIMessage]]:
-            return {"messages": [AIMessage(content="", tool_calls=[{"name": "sql_db_list_tables", "args": {}, "id": "tool_abcd123"}])]}
-    
-        def handle_tool_error(state: State) -> dict:
-            error = state.get("error")
-            tool_calls = state["messages"][-1].tool_calls
-            return {
-                "messages": [
-                    ToolMessage(content=f"Error: {repr(error)}\n please fix your mistakes.", tool_call_id=tc["id"])
-                    for tc in tool_calls
-                ]
-            }
-    
-        def create_tool_node_with_fallback(tools_list: list) -> RunnableWithFallbacks[Any, dict]:
-            return ToolNode(tools_list).with_fallbacks([RunnableLambda(handle_tool_error)], exception_key="error")
-    
-        def query_gen_node(state: State):
-            message = query_gen.invoke(state)
-            # Check for incorrect tool calls
-            tool_messages = []
-            if message.tool_calls:
-                for tc in message.tool_calls:
-                    if tc["name"] != "SubmitFinalAnswer":
-                        tool_messages.append(
-                            ToolMessage(
-                                content=f"Error: The wrong tool was called: {tc['name']}. Please fix your mistakes. Remember to only call SubmitFinalAnswer to submit the final answer. Generated queries should be outputted WITHOUT a tool call.",
-                                tool_call_id=tc["id"],
-                            )
-                        )
-            return {"messages": [message] + tool_messages}
-    
-        def should_continue(state: State) -> Literal[END, "correct_query", "query_gen"]:
-            messages = state["messages"]
-            last_message = messages[-1]
-            if getattr(last_message, "tool_calls", None):
-                return END
-            if last_message.content.startswith("Error:"):
-                return "query_gen"
-            else:
-                return "correct_query"
-    
-        def model_check_query(state: State) -> dict[str, list[AIMessage]]:
-            """Double-check if the query is correct before executing it."""
-            return {"messages": [query_check.invoke({"messages": [state["messages"][-1]]})]}
-    
-        # Get tools for listing tables and fetching schema
-        list_tables_tool = next((tool for tool in tools_instance if tool.name == "sql_db_list_tables"), None)
-        get_schema_tool = next((tool for tool in tools_instance if tool.name == "sql_db_schema"), None)
-    
-        # Define the workflow (state graph)
-        workflow = StateGraph(State)
-        workflow.add_node("first_tool_call", first_tool_call)
-        workflow.add_node("list_tables_tool", create_tool_node_with_fallback([list_tables_tool]))
-        workflow.add_node("get_schema_tool", create_tool_node_with_fallback([get_schema_tool]))
-        model_get_schema = llm.bind_tools([get_schema_tool])
-        workflow.add_node("model_get_schema", lambda state: {"messages": [model_get_schema.invoke(state["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")
-    
-        # Compile and return the agent application workflow.
-        return workflow.compile()
-    
-    # Option: configure static files from uploads folder as well.
+    def handle_tool_error(state: State) -> dict:
+        error = state.get("error")
+        tool_calls = state["messages"][-1].tool_calls
+        return {"messages": [
+            ToolMessage(content=f"Error: {repr(error)}. Please fix your mistakes.", tool_call_id=tc["id"])
+            for tc in tool_calls
+        ]}
+
+    def create_tool_node_with_fallback(tools_list: list) -> RunnableWithFallbacks[Any, dict]:
+        return ToolNode(tools_list).with_fallbacks([RunnableLambda(handle_tool_error)], exception_key="error")
+
+    def query_gen_node(state: State):
+        message = query_gen.invoke(state)
+        tool_messages = []
+        if message.tool_calls:
+            for tc in message.tool_calls:
+                if tc["name"] != "SubmitFinalAnswer":
+                    tool_messages.append(ToolMessage(
+                        content=f"Error: The wrong tool was called: {tc['name']}. Please fix your mistakes.",
+                        tool_call_id=tc["id"]
+                    ))
+        return {"messages": [message] + tool_messages}
+
+    def should_continue(state: State) -> Literal[END, "correct_query", "query_gen"]:
+        messages = state["messages"]
+        last_message = 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) -> dict[str, list[AIMessage]]:
+        return {"messages": [query_check.invoke({"messages": [state["messages"][-1]]})]}
+
+    # Get table listing and schema tools.
+    list_tables_tool = next((tool for tool in tools_instance if tool.name == "sql_db_list_tables"), None)
+    get_schema_tool = next((tool for tool in tools_instance if tool.name == "sql_db_schema"), None)
+
+    workflow = StateGraph(State)
+    workflow.add_node("first_tool_call", first_tool_call)
+    workflow.add_node("list_tables_tool", create_tool_node_with_fallback([list_tables_tool]))
+    workflow.add_node("get_schema_tool", create_tool_node_with_fallback([get_schema_tool]))
+    model_get_schema = llm.bind_tools([get_schema_tool])
+    workflow.add_node("model_get_schema", lambda state: {"messages": [model_get_schema.invoke(state["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 compiled workflow
+    return workflow.compile()
+
+# =============================================================================
+# create_app: The application factory.
+# =============================================================================
+def create_app():
+    # Configure static folder for uploads.
     flask_app = Flask(__name__, static_url_path='/uploads', static_folder='uploads')
     socketio = SocketIO(flask_app, cors_allowed_origins="*")
-    
-    # Set up uploads folder
-    UPLOAD_FOLDER_LOCAL = os.path.join(os.getcwd(), "uploads")
-    if not os.path.exists(UPLOAD_FOLDER_LOCAL):
-        os.makedirs(UPLOAD_FOLDER_LOCAL)
-    flask_app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER_LOCAL
-
-    # Static route: option if you want a custom route to serve files:
+
+    # Ensure uploads folder exists.
+    if not os.path.exists(UPLOAD_FOLDER):
+        os.makedirs(UPLOAD_FOLDER)
+    flask_app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+
+    # Serve uploaded files via a custom route.
     @flask_app.route("/files/<path:filename>")
     def uploaded_file(filename):
         return send_from_directory(flask_app.config['UPLOAD_FOLDER'], filename)
 
-    # Helper function to run the agent; uses the global agent_app.
+    # -------------------------------------------------------------------------
+    # Helper: run_agent runs the agent with the given prompt.
+    # -------------------------------------------------------------------------
     def run_agent(prompt, socketio):
         global agent_app
         if agent_app is None:
-            socketio.emit("log", {"message": "[ERROR]: No database has been uploaded. Please upload a database file first."})
-            socketio.emit("final", {"message": "No database available. Upload a database and try again."})
+            socketio.emit("log", {"message": "[ERROR]: No database has been uploaded. Upload a database file first."})
+            socketio.emit("final", {"message": "No database available. Upload one and try again."})
             return
         try:
             query = {"messages": [("user", prompt)]}
-            agent_app = create_agent_app(abs_file_path)
             result = agent_app.invoke(query)
             try:
                 result = result["messages"][-1].tool_calls[0]["args"]["final_answer"]
             except Exception:
                 result = "Query failed or no valid answer found."
-    
             print("final_answer------>", result)
-            socketio.emit("final", {"message": f"{result}"})
-            
+            socketio.emit("final", {"message": result})
         except Exception as e:
             print(f"[ERROR]: {str(e)}")
             socketio.emit("log", {"message": f"[ERROR]: {str(e)}"})
             socketio.emit("final", {"message": "Generation failed."})
 
+    # -------------------------------------------------------------------------
+    # Route: index page
+    # -------------------------------------------------------------------------
     @flask_app.route("/")
     def index():
         return render_template("index.html")
 
+    # -------------------------------------------------------------------------
+    # Route: generate (POST) – receives a prompt, runs the agent.
+    # -------------------------------------------------------------------------
     @flask_app.route("/generate", methods=["POST"])
     def generate():
         try:
             socketio.emit("log", {"message": "[STEP]: Entering query_gen..."})
             data = request.json
             prompt = data.get("prompt", "")
-            socketio.emit("log", {"message": f"[INFO]: Received prompt: {prompt}\n"})
-            # Run the agent in a separate thread
+            socketio.emit("log", {"message": f"[INFO]: Received prompt: {prompt}"})
             thread = threading.Thread(target=run_agent, args=(prompt, socketio))
-            socketio.emit("log", {"message": f"[INFO]: thread info: {thread}\n"})
-            socketio.emit("log", {"message": f"[INFO]: DB PATH: {abs_file_path}\n"})
+            socketio.emit("log", {"message": f"[INFO]: Starting thread: {thread}"})
             thread.start()
             return "OK", 200
         except Exception as e:
             print(f"[ERROR]: {str(e)}")
             socketio.emit("log", {"message": f"[ERROR]: {str(e)}"})
+            return "ERROR", 500
 
-    @flask_app.route("/upload", methods=["POST", "GET"])
+    # -------------------------------------------------------------------------
+    # 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 request.method == "POST":
+                file = request.files.get("file")
                 if not file:
                     print("No file uploaded")
                     return "No file uploaded", 400
-                if file and file.filename.endswith('.db'):
-                    # Save file using flask_app.config
-                    db_path = os.path.join(flask_app.config['UPLOAD_FOLDER'], 'uploaded.db')
-                    socketio.emit("log", {"message": f"[INFO]: Saving file to: {db_path}\n"})
+                # Secure the filename to avoid path traversal issues.
+                filename = secure_filename(file.filename)
+                if filename.endswith('.db'):
+                    db_path = os.path.join(flask_app.config['UPLOAD_FOLDER'], "uploaded.db")
                     print("Saving file to:", db_path)
                     file.save(db_path)
-                    
-                    # Reinitialize the agent_app with the new database file
-                    global abs_file_path
                     abs_file_path = os.path.abspath(db_path)
-                    global agent_app
                     agent_app = create_agent_app(abs_file_path)
-                    
-                    print(f"[INFO_PRINT]: Database file '{file.filename}' uploaded and loaded.")
-                    socketio.emit("log", {"message": f"[INFO]: Database file '{file.filename}' uploaded and loaded."})
+                    print(f"[INFO]: Database file '{filename}' uploaded and loaded.")
+                    socketio.emit("log", {"message": f"[INFO]: Database file '{filename}' uploaded and loaded."})
                     return redirect(url_for("index"))
-            # For GET, render upload form:
             return render_template("upload.html")
         except Exception as e:
             print(f"[ERROR]: {str(e)}")
@@ -317,7 +303,9 @@ def create_app():
 
     return flask_app, socketio
 
+# =============================================================================
 # Create the app for Gunicorn compatibility.
+# =============================================================================
 app, socketio_instance = create_app()
 
 if __name__ == "__main__":