Gonalb/multi_agentic_sql_generator
1
1import chainlit as cl2import pandas as pd3import time4from typing import Dict, Any5 6from agents.table_selection import table_selection_agent7from agents.data_retrieval import sample_data_retrieval_agent8from agents.sql_generation import sql_generation_agent9from agents.validation import query_validation_and_optimization10from agents.execution import execution_agent11from utils.bigquery_utils import init_bigquery_connection12from utils.feedback_utils import save_feedback_to_bigquery13 14@cl.on_chat_start15async def on_chat_start():16 """Initialize the chat session."""17 # Initialize BigQuery client18 client = init_bigquery_connection()19 20 # Store the client in the user session21 cl.user_session.set("client", client)22 23 # Send a welcome message24 await cl.Message(25 content="👋 Welcome to the Natural Language to SQL Query Assistant! Ask me any question about your e-commerce data.",26 author="SQL Assistant"27 ).send()28 29 # Add some example questions without using actions30 await cl.Message(31 content="Here are some example questions you can ask:",32 author="SQL Assistant"33 ).send()34 35 examples = [36 "What are the top 5 products by revenue?",37 "How many orders were placed in the last month?",38 "Which customers spent the most in 2023?",39 "What is the average order value by product category?"40 ]41 42 # Display all examples in a single message43 examples_text = "\n\n".join([f"• {example}" for example in examples])44 examples_text += "\n\n(You can copy and paste any of these examples to try them out)"45 46 await cl.Message(47 content=examples_text,48 author="SQL Assistant"49 ).send()50 51@cl.on_message52async def on_message(message: cl.Message):53 """Handle user messages."""54 query = message.content55 56 # Check if we're in "awaiting feedback" mode57 awaiting_feedback = cl.user_session.get("awaiting_feedback", False)58 if awaiting_feedback:59 client = cl.user_session.get("client")60 original_query = cl.user_session.get("original_query")61 generated_sql = cl.user_session.get("generated_sql")62 optimized_sql = cl.user_session.get("optimized_sql")63 64 # Save the detailed feedback65 feedback_details = f"negative: {query}"66 success = save_feedback_to_bigquery(67 client, 68 original_query, 69 generated_sql, 70 optimized_sql, 71 feedback_details72 )73 74 # Reset the awaiting feedback flag75 cl.user_session.set("awaiting_feedback", False)76 77 if success:78 await cl.Message(content="Thanks for your detailed feedback! I've saved it to improve future responses.", author="SQL Assistant").send()79 else:80 await cl.Message(content="Thanks for your feedback! (Note: There was an issue saving it to the database)", author="SQL Assistant").send()81 return82 83 # If not in feedback mode, process as a regular query84 # Get the BigQuery client from the user session85 client = cl.user_session.get("client")86 87 # Store the original query in the user session for feedback88 cl.user_session.set("original_query", query)89 90 # Send a thinking message91 thinking_msg = await cl.Message(content="🤔 Thinking...", author="SQL Assistant").send()92 93 try:94 # Step 1: Analyze relevant tables95 thinking_msg.content = "🔍 Analyzing relevant tables..."96 await thinking_msg.update()97 98 # Initialize the state with the query99 state = {"sql_query": query, "client": client}100 tables_state = table_selection_agent(state)101 relevant_tables = tables_state.get("relevant_tables", [])102 103 # Send the tables analysis with a slight delay for better UX104 await cl.sleep(1)105 if relevant_tables:106 tables_text = "I've identified these relevant tables for your query:\n\n"107 tables_text += "\n".join([f"- `{table}`" for table in relevant_tables])108 await cl.Message(content=tables_text, author="SQL Assistant").send()109 110 # Step 2: Retrieve sample data111 thinking_msg.content = "📊 Retrieving sample data..."112 await thinking_msg.update()113 await cl.sleep(1)114 115 # Update state with relevant tables and get sample data116 state.update(tables_state)117 sample_data_state = sample_data_retrieval_agent(state)118 119 # Step 3: Generate SQL120 thinking_msg.content = "💻 Generating SQL query..."121 await thinking_msg.update()122 await cl.sleep(1)123 124 # Update state with sample data and generate SQL125 state.update(sample_data_state)126 sql_state = sql_generation_agent(state)127 generated_sql = sql_state.get("generated_sql", "No SQL generated")128 129 # Store the generated SQL in the user session130 cl.user_session.set("generated_sql", generated_sql)131 132 # Send the generated SQL133 await cl.Message(134 content=f"Here's the SQL query I generated:\n\n```sql\n{generated_sql}\n```",135 author="SQL Assistant"136 ).send()137 138 # Step 4: Optimize SQL139 thinking_msg.content = "🔧 Optimizing the query..."140 await thinking_msg.update()141 await cl.sleep(1)142 143 # Update state with generated SQL and optimize144 state.update(sql_state)145 optimization_state = query_validation_and_optimization(state)146 optimized_sql = optimization_state.get("optimized_sql", "No optimized SQL")147 148 # Store the optimized SQL in the user session149 cl.user_session.set("optimized_sql", optimized_sql)150 151 # Send the optimized SQL152 await cl.Message(153 content=f"Here's the optimized version of the query:\n\n```sql\n{optimized_sql}\n```",154 author="SQL Assistant"155 ).send()156 157 # Step 5: Execute query158 thinking_msg.content = "⚙️ Executing query..."159 await thinking_msg.update()160 await cl.sleep(1)161 162 # Update state with optimized SQL and execute163 state.update(optimization_state)164 execution_state = execution_agent(state)165 execution_result = execution_state.get("execution_result", {})166 167 # Format and send the results168 if isinstance(execution_result, dict) and "error" in execution_result:169 error_msg = execution_result.get("error", "Unknown error occurred")170 await cl.Message(171 content=f"❌ Error executing query: {error_msg}",172 author="SQL Assistant"173 ).send()174 elif not execution_result:175 await cl.Message(176 content="✅ Query executed successfully but returned no results.",177 author="SQL Assistant"178 ).send()179 else:180 try:181 # Convert results to DataFrame for better display182 if isinstance(execution_result[0], tuple):183 # Try to get column names from BigQuery schema184 try:185 # Get the schema from the query job186 query_job = client.query(optimized_sql)187 schema = query_job.result().schema188 column_names = [field.name for field in schema]189 190 # Use these column names for the DataFrame191 df = pd.DataFrame(execution_result, columns=column_names)192 except Exception:193 # Fallback to generic column names194 columns = [f"Column_{i}" for i in range(len(execution_result[0]))]195 df = pd.DataFrame(execution_result, columns=columns)196 else:197 df = pd.DataFrame(execution_result)198 199 # Display the DataFrame as a table200 await cl.Message(201 content="✅ Query executed successfully! Here are the results:",202 author="SQL Assistant"203 ).send()204 205 # Send the DataFrame as an element206 elements = [cl.Dataframe(data=df)]207 await cl.Message(content="", elements=elements, author="SQL Assistant").send()208 209 # Also provide a summary of the results with feedback buttons210 num_rows = len(df)211 num_cols = len(df.columns)212 213 # Ask for feedback using AskActionMessage214 res = await cl.AskActionMessage(215 content=f"The query returned {num_rows} rows and {num_cols} columns.\n\nWas this result helpful?",216 actions=[217 cl.Action(name="feedback_positive", payload={"value": "positive"}, label="👍 Good results"),218 cl.Action(name="feedback_negative", payload={"value": "negative"}, label="👎 Not what I wanted")219 ],220 ).send()221 222 if res:223 feedback_value = res.get("payload", {}).get("value")224 225 client = cl.user_session.get("client")226 original_query = cl.user_session.get("original_query")227 generated_sql = cl.user_session.get("generated_sql")228 optimized_sql = cl.user_session.get("optimized_sql")229 230 if feedback_value == "positive":231 # Handle positive feedback232 success = save_feedback_to_bigquery(233 client, 234 original_query, 235 generated_sql, 236 optimized_sql, 237 "positive"238 )239 240 if success:241 await cl.Message(content="Thanks for your positive feedback! I've saved it to improve future responses.", author="SQL Assistant").send()242 else:243 await cl.Message(content="Thanks for your feedback! (Note: There was an issue saving it to the database)", author="SQL Assistant").send()244 245 elif feedback_value == "negative":246 # For negative feedback, just ask for text input247 await cl.Message(content="I'm sorry the results weren't what you expected. Please type your feedback about what was wrong.", author="SQL Assistant").send()248 249 # Set flag to indicate we're awaiting detailed feedback250 cl.user_session.set("awaiting_feedback", True)251 252 # Save initial negative feedback253 save_feedback_to_bigquery(254 client, 255 original_query, 256 generated_sql, 257 optimized_sql, 258 "negative"259 )260 261 except Exception as e:262 await cl.Message(263 content=f"❌ Error formatting results: {str(e)}",264 author="SQL Assistant"265 ).send()266 267 except Exception as e:268 # Handle any errors269 thinking_msg.content = f"❌ Error: {str(e)}"270 await thinking_msg.update()271 272 await cl.Message(273 content=f"I encountered an error while processing your query: {str(e)}",274 author="SQL Assistant"275 ).send()276 277# Callback handlers for actions278@cl.action_callback("feedback_positive")279async def on_feedback_positive(action):280 """Handle positive feedback."""281 client = cl.user_session.get("client")282 original_query = cl.user_session.get("original_query")283 generated_sql = cl.user_session.get("generated_sql")284 optimized_sql = cl.user_session.get("optimized_sql")285 286 # Handle positive feedback287 success = save_feedback_to_bigquery(288 client, 289 original_query, 290 generated_sql, 291 optimized_sql, 292 "positive"293 )294 295 if success:296 await cl.Message(content="Thanks for your positive feedback! I've saved it to improve future responses.", author="SQL Assistant").send()297 else:298 await cl.Message(content="Thanks for your feedback! (Note: There was an issue saving it to the database)", author="SQL Assistant").send()299 300@cl.action_callback("feedback_negative")301async def on_feedback_negative(action):302 """Handle negative feedback."""303 # Ask for more detailed feedback304 await cl.Message(content="I'm sorry the results weren't what you expected. Please type your feedback about what was wrong.", author="SQL Assistant").send()305 306 # Set flag to indicate we're awaiting detailed feedback307 cl.user_session.set("awaiting_feedback", True)308 309 client = cl.user_session.get("client")310 original_query = cl.user_session.get("original_query")311 generated_sql = cl.user_session.get("generated_sql")312 optimized_sql = cl.user_session.get("optimized_sql")313 314 # Save initial negative feedback315 save_feedback_to_bigquery(316 client, 317 original_query, 318 generated_sql, 319 optimized_sql, 320 "negative"321 )322 323# This is needed for Chainlit to run properly324if __name__ == "__main__":325 # Note: Chainlit uses its own CLI command to run the app326 # You'll run this with: chainlit run new_app.py -w327 pass