Spaces:
Sleeping
Sleeping
| 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 | |
| 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() | |
| 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. | |
| # ------------------------------------------------------------------------- | |
| def index(): | |
| return render_template("index.html") | |
| # ------------------------------------------------------------------------- | |
| # Route: generate (POST) – receives a prompt and runs the agent. | |
| # ------------------------------------------------------------------------- | |
| 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. | |
| # ------------------------------------------------------------------------- | |
| 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") | |
| 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) | |