app.py

# app.py (Final Version: Visuals + Financials + Robust Streaming + Multi-Provider + Model Transparency)
import gradio as gr
import httpx
import os
import json
# Import the visualization logic from our new module
from visuals import create_progress_chart, create_calibration_table, create_cost_summary
# Import config to display model defaults
import config 

# --- Config ---
BLAXEL_BASE_URL = os.getenv("BLAXEL_BACKEND_URL")
BLAXEL_API_KEY = os.getenv("BLAXEL_API_KEY") 

# --- Backend Client ---
def call_blaxel_backend(user_problem, google_key, anthropic_key, sambanova_key, openai_key, nebius_key):
    if not BLAXEL_BASE_URL or not BLAXEL_API_KEY:
        yield {status_output: "Configuration Error: Secrets not set."}
        return
    
    if not google_key:
        yield {status_output: "Input Error: Google API Key is required."}
        return

    full_endpoint_url = f"{BLAXEL_BASE_URL.rstrip('/')}/solve_problem"
    
    headers = {
        "Authorization": f"Bearer {BLAXEL_API_KEY}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "problem": user_problem,
        "keys": {
            "google": google_key,
            "anthropic": anthropic_key or None,
            "sambanova": sambanova_key or None,
            "openai": openai_key or None,
            "nebius": nebius_key or None,
        }
    }

    yield {status_output: f"Connecting to MudabbirAI..."}
    
    try:
        # High timeout for complex agentic workflows
        with httpx.stream("POST", full_endpoint_url, json=payload, headers=headers, timeout=600.0) as response:
            if response.status_code != 200:
                yield {status_output: f"HTTP Error: {response.status_code}"}
                return 

            final_json = None
            full_log = ""

            # Robust SSE Line Parsing
            for line in response.iter_lines():
                if not line: continue 
                if hasattr(line, "decode"): line = line.decode("utf-8")
                if line.startswith(":"): continue # Skip keep-alive comments

                if line.startswith("data: "):
                    content = line.replace("data: ", "", 1).strip()
                    
                    # Check for FINAL payload
                    if content.startswith("FINAL:"):
                        try:
                            json_str = content.replace("FINAL:", "", 1)
                            final_json = json.loads(json_str)
                        except json.JSONDecodeError as e:
                            full_log += f"\n[Error Parsing Final JSON]: {e}"
                    else:
                        # Streaming Log Update
                        full_log += content + "\n"
                        yield {status_output: full_log}

            # Final Rendering
            if final_json:
                log_data = final_json.get("log")
                
                # Create Visuals & Cost Report
                chart = create_progress_chart(log_data)
                calib_table = create_calibration_table(log_data)
                cost_md = create_cost_summary(log_data)

                yield {
                    status_output: full_log + "\nDone!",
                    final_text_output: final_json.get("text"),
                    final_audio_output: final_json.get("audio"),
                    final_json_log: log_data,
                    progress_plot: chart,
                    calibration_data: calib_table,
                    cost_display: cost_md
                }
            else:
                yield {status_output: full_log + "\n[Stream ended without final payload]"}

    except Exception as e:
        yield {status_output: f"Connection Error: {str(e)}"}

# --- Helper to Format Model Info ---
def get_model_info():
    info = "### 🤖 System Configuration (Default Models)\n"
    info += f"* **Judge & Orchestrator:** {config.MODELS['Gemini']['judge']}\n"
    info += f"* **Classifier:** {config.MODELS['Gemini']['classifier']}\n"
    info += f"* **Homogeneous Team:** {config.MODELS['Gemini']['default']} (x3)\n"
    info += "\n*Note: In 'Cognitive Labyrinth' mode, the system will audition models from all provided keys to find the best fit for each role.*"
    return info

# --- Gradio UI ---
with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald", secondary_hue="slate")) as demo:
    gr.Markdown(
        """
        # 🧠 MudabbirAI: The Strategic Selector
        ### A Multi-Agent System that *Auditions* Models to Build the Perfect Team.
        """
    )

    with gr.Row():
        with gr.Column(scale=2):
            problem_input = gr.Textbox(label="1. Enter Your 'Business Problem'", lines=4, placeholder="e.g., Design a sustainable urban farming initiative for Tokyo...")
            
            with gr.Accordion("2. Sponsor API Keys (Expand)", open=False):
                # --- UPDATED LABELS WITH MODEL NAMES ---
                google_key_input = gr.Textbox(
                    label=f"Google API Key ({config.MODELS['Gemini']['default']}) (Required)", 
                    type="password"
                )
                with gr.Row():
                    anthropic_key_input = gr.Textbox(
                        label=f"Anthropic API Key ({config.MODELS['Anthropic']['default']}) (Optional)", 
                        type="password"
                    )
                    sambanova_key_input = gr.Textbox(
                        label=f"SambaNova API Key ({config.MODELS['SambaNova']['default']}) (Optional)", 
                        type="password"
                    )
                with gr.Row():
                    openai_key_input = gr.Textbox(
                        label=f"OpenAI API Key ({config.MODELS['OpenAI']['default']}) (Optional)", 
                        type="password"
                    )
                    nebius_key_input = gr.Textbox(
                        label=f"Nebius API Key ({config.MODELS['Nebius']['default']}) (Optional)", 
                        type="password"
                    )
        
        with gr.Column(scale=1):
            # Display Default Model Info here for transparency
            gr.Markdown(get_model_info())
            submit_button = gr.Button("🚀 Deploy MudabbirAI", variant="primary", size="lg")

    with gr.Tabs():
        with gr.TabItem("📊 Analytics Dashboard", id="viz_tab"):
            with gr.Row():
                progress_plot = gr.Plot(label="Improvement Trajectory (Draft 1 vs. Final)")
            with gr.Row():
                calibration_data = gr.Dataframe(label="Team Calibration Results (Model Scores)")
            with gr.Row():
                cost_display = gr.Markdown(label="Financial Intelligence")
                
        with gr.TabItem("📝 Final Briefing", id="result_tab"):
            final_text_output = gr.Markdown(label="Strategic Report")
            final_audio_output = gr.Audio(label="Audio Briefing")
            
        with gr.TabItem("⚙️ Internal Logs", id="log_tab"):
            status_output = gr.Textbox(label="Live System Logs", lines=15, interactive=False, autoscroll=True)
            final_json_log = gr.JSON(label="Full Execution Trace (JSON)")
        
    submit_button.click(
        fn=call_blaxel_backend,
        inputs=[problem_input, google_key_input, anthropic_key_input, sambanova_key_input, openai_key_input, nebius_key_input],
        outputs=[status_output, final_text_output, final_audio_output, final_json_log, progress_plot, calibration_data, cost_display]
    )

demo.launch()