Update app.py

app.py

@@ -1,76 +1,146 @@
-# app.py (Smoke Test - Now with Authentication)
+# app.py (New Secure & Robust Version)
 import gradio as gr
 import httpx
 import os
+import json
 
 # --- Config ---
-# Read the BASE URL from Hugging Face secrets
+# Secrets for our *own* backend
 BLAXEL_BASE_URL = os.getenv("BLAXEL_BACKEND_URL")
-# Read the API KEY from Hugging Face secrets
 BLAXEL_API_KEY = os.getenv("BLAXEL_API_KEY") 
 
 # --- Backend Client ---
-def call_blaxel_backend(user_problem):
+def call_blaxel_backend(
+    user_problem: str, 
+    google_key: str, 
+    anthropic_key: str, 
+    sambanova_key: str
+):
     
-    # 1. Check if the URL secret was loaded
-    if not BLAXEL_BASE_URL:
-        yield "Configuration Error: BLAXEL_BACKEND_URL secret is not set in Hugging Face."
-        return 
-
-    # 2. Check if the API KEY secret was loaded
-    if not BLAXEL_API_KEY:
-        yield "Configuration Error: BLAXEL_API_KEY secret is not set in Hugging Face."
-        yield "Please add your Blaxel API Key to the 'Secrets' section."
+    # --- 1. Validate Keys ---
+    if not BLAXEL_BASE_URL or not BLAXEL_API_KEY:
+        yield "Configuration Error: The app's own BLAXEL secrets are not set."
+        return
+    
+    # We require at least the Google key
+    if not google_key:
+        yield {status_output: "Input Error: Please enter your Google API Key to continue. It is required for the 'Judge' agent."}
         return
 
-    # 3. Construct the correct, full URL
+    # --- 2. Set up connection details ---
     full_endpoint_url = f"{BLAXEL_BASE_URL.rstrip('/')}/solve_problem"
     
-    # 4. Create the authentication headers
     headers = {
         "Authorization": f"Bearer {BLAXEL_API_KEY}",
         "Content-Type": "application/json"
     }
+    
+    # --- 3. Securely package the user's keys ---
+    payload = {
+        "problem": user_problem,
+        "keys": {
+            "google": google_key,
+            "anthropic": anthropic_key or None,  # Send None if empty
+            "sambanova": sambanova_key or None, # Send None if empty
+        }
+    }
 
-    yield f"Attempting to connect to MudabbirAI (at {full_endpoint_url})..."
+    yield f"Connecting to MudabbirAI (at {full_endpoint_url})..."
     
     try:
-        # 5. Pass the headers to the request
-        with httpx.stream("POST", full_endpoint_url, json={"problem": user_problem}, headers=headers, timeout=60) as response:
+        # Set a long timeout for all the LLM calls
+        with httpx.stream("POST", full_endpoint_url, json=payload, headers=headers, timeout=300.0) as response:
             
             if response.status_code != 200:
-                # Manually read the error response text
                 error_details = response.read().decode()
-                if response.status_code == 401 or response.status_code == 403:
-                    yield f"Authentication Error (HTTP {response.status_code}): Incorrect or invalid BLAXEL_API_KEY secret."
-                else:
-                    yield f"HTTP Error: Received status code {response.status_code} from Blaxel."
-                yield f"Details: {error_details}"
+                yield {status_output: f"HTTP Error: Received status code {response.status_code} from Blaxel.\nDetails: {error_details}"}
                 return 
 
-            # If we get here, the status is 200!
+            final_json = None
+            full_log = ""
             for chunk in response.iter_text():
-                # Each chunk is a status update
-                yield chunk
+                if chunk.startswith("FINAL:"):
+                    # This is our final package, don't yield it
+                    final_json = json.loads(chunk.replace("FINAL:", ""))
+                else:
+                    # This is a live log update
+                    full_log += chunk + "\n"
+                    yield {status_output: full_log} # Update log
+
+            # --- 4. Return final results to all UI components ---
+            if final_json:
+                yield {
+                    status_output: full_log + "\nDone!",
+                    final_text_output: final_json.get("text"),
+                    final_audio_output: final_json.get("audio")
+                }
 
     except httpx.ConnectError as e:
-        yield f"Connection Error: Could not connect to Blaxel. Is the URL correct and the backend running?\nDetails: {str(e)}"
+        yield {status_output: f"Connection Error: Could not connect to Blaxel.\nDetails: {str(e)}"}
     except httpx.ReadTimeout:
-        yield "Error: The connection timed out."
+        yield {status_output: "Error: The connection timed out. (This is a complex agent and may take several minutes.)"}
     except Exception as e:
-        yield f"An unknown error occurred: {str(e)}"
+        yield {status_output: f"An unknown error occurred: {str(e)}"}
 
-# --- Gradio UI ---
+# --- Gradio UI (Now with Optional Keys) ---
 with gr.Blocks() as demo:
-    gr.Markdown("# MudabbirAI - Connection Test")
-    problem_input = gr.Textbox(label="Enter any text to test")
-    submit_button = gr.Button("Test Connection")
-    status_output = gr.Textbox(label="Live Log from Blaxel", interactive=False)
+    gr.Markdown("# MudabbirAI: The Self-Calibrating Strategic Selector")
+    gr.Markdown("Based on the research by Youssef Hariri (Papers 1 & 2)")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            problem_input = gr.Textbox(
+                label="Enter Your 'Wicked Problem'",
+                placeholder="e.g., 'Develop a new market entry strategy for Southeast Asia...'",
+                lines=5
+            )
+            
+            gr.Markdown("### Sponsor API Keys")
+            gr.Markdown("Your keys are sent securely, used only for this request, and never stored.")
+            
+            google_key_input = gr.Textbox(
+                label="Google API Key (Required)", type="password",
+                placeholder="Required for 'Judge' and 'default' agent"
+            )
+            
+            with gr.Row():
+                anthropic_key_input = gr.Textbox(
+                    label="Anthropic API Key (Optional)", type="password",
+                    placeholder="Enter to enable Anthropic models"
+                )
+                sambanova_key_input = gr.Textbox(
+                    label="SambaNova API Key (Optional)", type="password",
+                    placeholder="Enter to enable SambaNova models"
+                )
+            
+            submit_button = gr.Button("Deploy Agent", variant="primary")
+
+        with gr.Column(scale=1):
+            status_output = gr.Textbox(
+                label="Agent's Internal Monologue (Live Log)",
+                lines=20, # Made log even taller
+                interactive=False
+            )
+            final_text_output = gr.Markdown(label="Final Briefing")
+            final_audio_output = gr.Audio(label="Audio Presentation")
+            
+    all_inputs = [
+        problem_input,
+        google_key_input,
+        anthropic_key_input,
+        sambanova_key_input
+    ]
     
+    all_outputs = {
+        status_output,
+        final_text_output,
+        final_audio_output
+    }
+
     submit_button.click(
         fn=call_blaxel_backend,
-        inputs=problem_input,
-        outputs=status_output
+        inputs=all_inputs,
+        outputs=list(all_outputs) # Convert set to list
     )
 
 demo.launch()