Update app.py

app.py

@@ -1,12 +1,11 @@
-import gradio as gr
 import sqlite3
-import json,os
-import openai  # Ensure you have openai==0.27.0 installed
+import json
+import openai
 from langgraph.graph import StateGraph, START, END
 from typing import TypedDict, Optional
 
-# --- Set your OpenAI API key securely ---
-# Remove hardcoding and rely on environment variables in your Space.
+# 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
@@ -20,7 +19,11 @@ class SQLExecutionState(TypedDict):
 # Initialize the LangGraph Workflow
 graph = StateGraph(state_schema=SQLExecutionState)
 
-# ------------- 1. Query Understanding Agent -------------
+# 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"]
     prompt = f"""
@@ -68,7 +71,7 @@ def query_understanding_agent(state: SQLExecutionState) -> SQLExecutionState:
 
 graph.add_node("Query Understanding", query_understanding_agent)
 
-# ------------- 2. Query Validation Agent -------------
+# ------------------ 2. Query Validation Agent ------------------
 def query_validation_agent(state: SQLExecutionState) -> SQLExecutionState:
     sql_metadata = state.get("structured_metadata", {})
     if "error" in sql_metadata:
@@ -81,7 +84,7 @@ def query_validation_agent(state: SQLExecutionState) -> SQLExecutionState:
 
 graph.add_node("Query Validation", query_validation_agent)
 
-# ------------- 3. Query Optimization Agent -------------
+# ------------------ 3. Query Optimization Agent ------------------
 def query_optimization_agent(state: SQLExecutionState) -> SQLExecutionState:
     sql_metadata = state.get("structured_metadata", {})
     prompt = f"""
@@ -103,10 +106,9 @@ def query_optimization_agent(state: SQLExecutionState) -> SQLExecutionState:
         optimized_query = optimized_query.replace("```sql", "").replace("```", "").strip()
     return {"optimized_sql": optimized_query}
 
-
 graph.add_node("Query Optimization", query_optimization_agent)
 
-# ------------- 4. SQL Execution Agent -------------
+# ------------------ 4. SQL Execution Agent ------------------
 def execution_agent(state: SQLExecutionState) -> SQLExecutionState:
     query = state.get("optimized_sql", "").strip()
     if not query:
@@ -126,28 +128,20 @@ def execution_agent(state: SQLExecutionState) -> SQLExecutionState:
 
 graph.add_node("SQL Execution", execution_agent)
 
-# Define Execution Flow
-graph.add_edge("START", "Query Understanding")
+# ------------------ 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")
+graph.add_edge("SQL Execution", END)
 
 compiled_pipeline = graph.compile()
 
-# Wrap your multi-agent query execution into a callable function
+# ------------------ Example Execution ------------------
 def run_multi_agent_query(natural_language_query):
     result = compiled_pipeline.invoke({"sql_query": natural_language_query})
     return json.dumps(result.get("execution_result", {}), indent=2)
 
-# Gradio Interface
-iface = gr.Interface(
-    fn=run_multi_agent_query,
-    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter your natural language SQL query here..."),
-    outputs="text",
-    title="Multi-Agent SQL Generator",
-    description="Enter a natural language query to generate and execute SQL."
-)
-
 if __name__ == "__main__":
-    iface.launch()
+    # For testing outside of Hugging Face Spaces
+    print(run_multi_agent_query("Find the email_id of the top 5 customers who spent the most in 2024."))