import gradio as gr import requests import uuid import time import urllib3 import os import pandas as pd import base64 from dotenv import load_dotenv load_dotenv() # Disable SSL warnings for corporate networks urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) # --- CONFIGURATION --- BASE_URL = os.getenv("BASE_URL") BEARER_TOKEN = os.getenv("BEARER_TOKEN") HEADERS = { "Authorization": f"Bearer {BEARER_TOKEN}", "Content-Type": "application/json" } def process_zbb_query(user_message, history, session_id, last_query, raw_sql): """ Main chat handler. Format: List of Dictionaries [{"role": "user", "content": ...}] """ # 1. Update Last Query State & Session last_query = user_message if not session_id: session_id = f"sess_{uuid.uuid4().hex[:6]}" # 2. Add User Message & 'Thinking' Placeholder (Dictionary format) # This matches the "messages" format required by newer Gradio versions history.append({"role": "user", "content": user_message}) history.append({"role": "assistant", "content": "⏳ **Agent is thinking...**"}) # Reset raw_sql state for new query raw_sql = [] # Yield initial state yield history, session_id, gr.update(visible=False, value=None), last_query, raw_sql try: # Construct history for backend (excluding the last "Thinking" message) backend_history = history[:-1] # Kickoff request kickoff_url = f"{BASE_URL}/kickoff" payload = { "inputs": { "session_id": session_id, "current_query": user_message, "conversation_history": backend_history } } post_resp = requests.post(kickoff_url, json=payload, headers=HEADERS, verify=False) post_resp.raise_for_status() kickoff_id = post_resp.json().get("kickoff_id") # Polling Loop status_url = f"{BASE_URL}/status/{kickoff_id}" result_data = None for i in range(200): get_resp = requests.get(status_url, headers=HEADERS, verify=False) status_data = get_resp.json() if status_data.get("state") == "SUCCESS": result_data = status_data.get("result") break # Update thinking status with progress # Access the last message dictionary and update its content history[-1]["content"] = f"⏳ **Agent is thinking...** ({i+1}s)" yield history, session_id, gr.update(visible=False), last_query, raw_sql time.sleep(1) # Handle Results if result_data: response_text = result_data.get("assistant_message", "No message returned.") sql_data = result_data.get("sql_result", []) # Replace 'Thinking' message with actual response history[-1]["content"] = response_text # If SQL data exists, update DataFrame & State if isinstance(sql_data, list) and len(sql_data) > 0: raw_sql = sql_data df = pd.DataFrame(sql_data) df = df.map(lambda x: x.strip("'") if isinstance(x, str) else x) yield history, session_id, gr.update(visible=True, value=df), last_query, raw_sql else: yield history, session_id, gr.update(visible=False, value=None), last_query, raw_sql else: history[-1]["content"] = "⚠️ Request timed out. Please try again." yield history, session_id, gr.update(visible=False), last_query, raw_sql except Exception as e: history[-1]["content"] = f"❌ Error: {str(e)}" yield history, session_id, gr.update(visible=False), last_query, raw_sql def generate_visualization(session_id, last_query, raw_sql): """ Handler for the Visualization button. """ if not raw_sql: return "

⚠️ No data available to visualize. Please run a query first.

" yield "

⏳ Generating Visualization...

" try: kickoff_url = f"{BASE_URL}/kickoff" payload = { "inputs": { "session_id": session_id, "current_query": last_query, "conversation_history": [], "router_flag": "viz", "sql_result": raw_sql } } post_resp = requests.post(kickoff_url, json=payload, headers=HEADERS, verify=False) post_resp.raise_for_status() kickoff_id = post_resp.json().get("kickoff_id") status_url = f"{BASE_URL}/status/{kickoff_id}" for i in range(200): get_resp = requests.get(status_url, headers=HEADERS, verify=False) status_data = get_resp.json() if status_data.get("state") == "SUCCESS": result = status_data.get("result", {}) raw_html = result.get("final_response", "") if not raw_html: yield "
No visualization content returned
" return try: html_b64 = base64.b64encode(raw_html.encode('utf-8')).decode('utf-8') iframe_html = f""" """ yield iframe_html except Exception as encode_error: yield f"
Error encoding visualization: {str(encode_error)}
" return time.sleep(1) yield "

⚠️ Visualization timed out.

" except Exception as e: yield f"

❌ Error generating visualization: {str(e)}

" # --- UI Setup --- with gr.Blocks() as demo: gr.Markdown("# 💰 ZBB GenAI Analysis") session_id = gr.State("") last_query_state = gr.State("") raw_sql_state = gr.State([]) # Pre-define dataframe (hidden initially) source_df = gr.DataFrame( label="Query Output", interactive=False, visible=False, wrap=True, render=False ) with gr.Row(): with gr.Column(scale=3): # FIXED: Removed type="messages". # Gradio 5+ defaults to messages format without the argument. # We must pass Dictionary data (done in process_zbb_query) chatbot = gr.Chatbot(label="ZBB Assistant", height=500) msg = gr.Textbox( placeholder="Type your query here...", label="Command" ) # --- SUGGESTIONS SECTION --- gr.Markdown("### 💡 Suggestions") examples = gr.Examples( examples=[["How is Overhead tracking vs Budget YTD for NAZ by Function?"], ["Show me NAZ overhead performance"], ["What is the travel budget for SAZ?"], ["Compare Marketing actuals vs budget for M1 2025"], ["Drill down into IT expenses for APAC"] ], inputs=[msg], label="Click to fill query" ) with gr.Column(scale=2): with gr.Accordion("📊 Data Sources ", open=False) as data_accordion: source_df.render() gr.Markdown("### Visualization") viz_btn = gr.Button("📊 Get Visualization", variant="primary") viz_output = gr.HTML(label="Chart", min_height=500) # 1. Chat Submission submit_event = msg.submit( process_zbb_query, inputs=[msg, chatbot, session_id, last_query_state, raw_sql_state], outputs=[chatbot, session_id, source_df, last_query_state, raw_sql_state] ).then(lambda: "", None, [msg]) # 2. Visualization Event viz_btn.click( generate_visualization, inputs=[session_id, last_query_state, raw_sql_state], outputs=[viz_output] ) if __name__ == "__main__": import os from dotenv import load_dotenv load_dotenv() demo.launch( auth=(os.getenv("ID"), os.getenv("PASS")), auth_message="Please enter your ABInBev credentials to access the ZBB GenAI Tool.", theme=gr.themes.Soft() )