Gonalb/multi_agentic_sql_generator

1from langgraph.graph import StateGraph, START, END2from typing import TypedDict, Optional3 4from agents.table_selection import table_selection_agent5from agents.data_retrieval import sample_data_retrieval_agent6from agents.sql_generation import sql_generation_agent7from agents.validation import query_validation_and_optimization8from agents.execution import execution_agent9from utils.bigquery_utils import init_bigquery_connection10 11# Define the state schema12class SQLExecutionState(TypedDict):13    sql_query: str  # Natural language query14    client: Optional[object]  # BigQuery client15    relevant_tables: Optional[list]  # Tables identified as relevant16    sample_data: Optional[dict]      # Sample data from relevant tables17    generated_sql: Optional[str]     # The actual SQL query (not JSON)18    validation_result: Optional[dict]19    optimized_sql: Optional[str]20    execution_result: Optional[dict]21 22def initialize_client(state: SQLExecutionState) -> SQLExecutionState:23    """Initialize the BigQuery client and add it to the state."""24    client = init_bigquery_connection()25    return {"client": client}26 27def create_workflow():28    """Create and return the workflow graph."""29    # Initialize the LangGraph Workflow30    graph = StateGraph(state_schema=SQLExecutionState)31    32    # Add nodes33    graph.add_node("Initialize Client", initialize_client)34    graph.add_node("Table Selection", table_selection_agent)35    graph.add_node("Sample Data Retrieval", sample_data_retrieval_agent)36    graph.add_node("SQL Generation", sql_generation_agent)37    graph.add_node("Query Validation & Optimization", query_validation_and_optimization)38    graph.add_node("SQL Execution", execution_agent)39    40    # Define execution flow41    graph.add_edge(START, "Initialize Client")42    graph.add_edge("Initialize Client", "Table Selection")43    graph.add_edge("Table Selection", "Sample Data Retrieval")44    graph.add_edge("Sample Data Retrieval", "SQL Generation")45    graph.add_edge("SQL Generation", "Query Validation & Optimization")46    graph.add_edge("Query Validation & Optimization", "SQL Execution")47    graph.add_edge("SQL Execution", END)48    49    # Compile the graph50    return graph.compile()