Update app.py

app.py

@@ -1,11 +1,11 @@
+import os
 import sqlite3
 import json
 import openai
-from langgraph.graph import StateGraph, START, END
+from langgraph.graph import StateGraph
 from typing import TypedDict, Optional
 
 # Set your OpenAI API key from environment variable
-import os
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
 # Define the state for SQL Execution Workflow
@@ -19,10 +19,6 @@ class SQLExecutionState(TypedDict):
 # Initialize the LangGraph Workflow
 graph = StateGraph(state_schema=SQLExecutionState)
 
-# Add dummy nodes for START and END
-graph.add_node(START, lambda state: state)
-graph.add_node(END, lambda state: state)
-
 # ------------------ 1. Query Understanding Agent ------------------
 def query_understanding_agent(state: SQLExecutionState) -> SQLExecutionState:
     natural_language_query = state["sql_query"]
@@ -129,11 +125,9 @@ def execution_agent(state: SQLExecutionState) -> SQLExecutionState:
 graph.add_node("SQL Execution", execution_agent)
 
 # ------------------ Define Execution Flow ------------------
-graph.add_edge(START, "Query Understanding")
 graph.add_edge("Query Understanding", "Query Validation")
 graph.add_edge("Query Validation", "Query Optimization")
 graph.add_edge("Query Optimization", "SQL Execution")
-graph.add_edge("SQL Execution", END)
 
 compiled_pipeline = graph.compile()