Create app.py

app.py

@@ -0,0 +1,155 @@
+import gradio as gr
+import sqlite3
+import json
+import openai  # Ensure you have openai==0.27.0 installed
+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.
+openai.api_key = os.getenv("OPENAI_API_KEY")
+
+# Define the state for SQL Execution Workflow
+class SQLExecutionState(TypedDict):
+    sql_query: str
+    structured_metadata: Optional[dict]
+    validation_result: Optional[dict]
+    optimized_sql: Optional[str]
+    execution_result: Optional[dict]
+
+# Initialize the LangGraph Workflow
+graph = StateGraph(state_schema=SQLExecutionState)
+
+# ------------- 1. Query Understanding Agent -------------
+def query_understanding_agent(state: SQLExecutionState) -> SQLExecutionState:
+    natural_language_query = state["sql_query"]
+    prompt = f"""
+    Convert the following natural language query into **structured SQL metadata** based on the database schema.
+    If you cannot generate a query that adheres strictly to the schema, return:
+    {{ "error": "Invalid query: Tables or columns do not match schema" }}
+
+    **Query:** "{natural_language_query}"
+
+    **Database Schema:**
+    - **orders** (order_id, customer_id, order_date, total_amount, status)
+    - **order_items** (order_item_id, order_id, product_id, quantity, subtotal)
+    - **products** (product_id, name, category, price, stock_quantity)
+    - **customers** (customer_id, name, email, phone, address, created_at)
+    - **payments** (payment_id, order_id, payment_date, amount, payment_method, status)
+    
+    **Rules:**
+    - Use only the provided tables.
+    - Ensure correct column names.
+    - Return output strictly in JSON format.
+    - Group by relevant fields when necessary.
+
+    **Example Output Format:**
+    {json.dumps({
+        "operation": "SELECT",
+        "columns": ["customer_id", "SUM(total_amount) AS total_spent"],
+        "table": "orders",
+        "conditions": ["order_date BETWEEN '2024-01-01' AND '2024-12-31'"],
+        "group_by": ["customer_id"],
+        "order_by": ["total_spent DESC"],
+        "limit": 5
+    }, indent=4)}
+
+    **DO NOT return explanations. Only return valid JSON.**
+    """
+    response = openai.ChatCompletion.create(
+        model="gpt-4o-mini",
+        messages=[{"role": "user", "content": prompt}]
+    )
+    try:
+        metadata = json.loads(response["choices"][0]["message"]["content"])
+        return {"structured_metadata": metadata}
+    except json.JSONDecodeError:
+        return {"structured_metadata": {"error": "Invalid JSON response from OpenAI"}}
+
+graph.add_node("Query Understanding", query_understanding_agent)
+
+# ------------- 2. Query Validation Agent -------------
+def query_validation_agent(state: SQLExecutionState) -> SQLExecutionState:
+    sql_metadata = state.get("structured_metadata", {})
+    if "error" in sql_metadata:
+        return {"validation_result": {"error": sql_metadata["error"]}}
+    query = sql_metadata.get("operation", "")
+    restricted_keywords = ["DROP", "DELETE", "TRUNCATE", "ALTER"]
+    if any(keyword in query.upper() for keyword in restricted_keywords):
+        return {"validation_result": {"error": "Potentially harmful SQL operation detected!"}}
+    return {"validation_result": {"valid": True}}
+
+graph.add_node("Query Validation", query_validation_agent)
+
+# ------------- 3. Query Optimization Agent -------------
+def query_optimization_agent(state: SQLExecutionState) -> SQLExecutionState:
+    sql_metadata = state.get("structured_metadata", {})
+    prompt = f"""
+    Optimize the following SQL query for performance while ensuring that the output includes only the required columns and necessary joins.
+    Do not include any extra columns, unnecessary joins, or records that are not required to answer the query.
+
+    Here is the original SQL metadata:
+    {json.dumps(sql_metadata, indent=4)}
+
+    Output only the final optimized SQL query in plain text without any markdown formatting or explanations.
+    """
+    response = openai.ChatCompletion.create(
+        model="gpt-4o-mini",
+        messages=[{"role": "user", "content": prompt}],
+        temperature=0
+    )
+    optimized_query = response["choices"][0]["message"]["content"].strip()
+    if optimized_query.startswith("
+sql"):
+        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 -------------
+def execution_agent(state: SQLExecutionState) -> SQLExecutionState:
+    query = state.get("optimized_sql", "").strip()
+    if not query:
+        return {"execution_result": {"error": "No SQL query to execute."}}
+    try:
+        conn = sqlite3.connect("complex_test_db.sqlite", timeout=20)
+        cursor = conn.cursor()
+        cursor.execute(query)
+        result = cursor.fetchall()
+        cursor.close()
+        conn.close()
+        if not result:
+            return {"execution_result": {"error": "Query executed successfully but returned no results."}}
+        return {"execution_result": result}
+    except sqlite3.Error as e:
+        return {"execution_result": {"error": str(e)}}
+
+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()
+
+# Wrap your multi-agent query execution into a callable function
+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()