app.py

import gradio as gr
import requests
import uuid
import time
import urllib3
import os
from dotenv import load_dotenv

load_dotenv()

# Disable SSL warnings for corporate networks/self-signed certs
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):
    if not session_id:
        session_id = f"sess_{uuid.uuid4().hex[:6]}"

    # Gradio 6.0 Message format
    history.append({"role": "user", "content": [{"type": "text", "text": user_message}]})
    history.append({"role": "assistant", "content": [{"type": "text", "text": "⏳ Thinking your request..."}]})
    yield history, session_id

    try:
        # POST to Kickoff
        kickoff_url = f"{BASE_URL}/kickoff"
        payload = {
            "inputs": {
                "session_id": session_id, 
                "current_query": user_message, 
                "conversation_history": history[:-1] 
            }
        }
        
        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(60):
            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
            
            history[-1]["content"][0]["text"] = f"⏳ Processing... ({i+1}s)"
            yield history, session_id
            time.sleep(1)

        if result_data:
            response_text = result_data.get("assistant_message", "No message returned.")
            history[-1]["content"][0]["text"] = response_text
        else:
            history[-1]["content"][0]["text"] = "⚠️ Request timed out. Please try again."
        
        yield history, session_id

    except Exception as e:
        history[-1]["content"][0]["text"] = f"❌ Error: {str(e)}"
        yield history, session_id

# --- UI Setup ---
with gr.Blocks() as demo:
    gr.Markdown("# 💰 ZBB GenAI Analysis")
    session_id = gr.State("")
    chatbot = gr.Chatbot(label="ZBB Assistant")
    msg = gr.Textbox(placeholder="Enter query...")
    
    msg.submit(
        process_zbb_query, 
        inputs=[msg, chatbot, session_id], 
        outputs=[chatbot, session_id]
    ).then(lambda: "", None, [msg])

if __name__ == "__main__":
    # --- AUTHENTICATION GATE ---
    # username: ABINBEV | password: Venkat
    demo.launch(
        auth=("ABINBEV", "Venkat"),
        auth_message="Please enter your ABInBev credentials to access the ZBB GenAI Tool.",
        theme=gr.themes.Soft()
    )