Update app.py

app.py

@@ -19,20 +19,16 @@ import sys
 import asyncio
 import json
 import zipfile
+import shutil
 from pathlib import Path
-from typing import Dict, Any, Optional, Tuple, AsyncGenerator
+from typing import Dict, Any, List, Optional, Tuple, AsyncGenerator
+
 from datetime import datetime
 
 import gradio as gr
 import plotly.graph_objects as go
 import networkx as nx
 
-# Optional: upload generated ZIPs back into the Space repo
-try:
-    from huggingface_hub import HfApi
-except Exception:
-    HfApi = None  # Safe fallback if not installed
-
 # Add project root to path
 sys.path.insert(0, str(Path(__file__).parent))
 
@@ -44,29 +40,12 @@ from ui.voice_interface import voice
 
 # Load environment variables
 from dotenv import load_dotenv
-load_dotenv()
 
+load_dotenv()
 
-# =============================================================================
-# Helpers
-# =============================================================================
-
-def to_jsonable(obj: Any) -> Any:
-    """
-    Recursively convert objects to JSON-serializable equivalents.
-    - Path -> str
-    - datetime -> isoformat
-    - dict/list/tuple/set -> recurse
-    """
-    if isinstance(obj, Path):
-        return str(obj)
-    if isinstance(obj, datetime):
-        return obj.isoformat()
-    if isinstance(obj, dict):
-        return {k: to_jsonable(v) for k, v in obj.items()}
-    if isinstance(obj, (list, tuple, set)):
-        return [to_jsonable(v) for v in obj]
-    return obj
+# ============================================================================
+# Download + Hub Upload Functionality
+# ============================================================================
 
 
 def create_download_zip(server_metadata: Dict[str, Any]) -> Optional[str]:
@@ -74,81 +53,80 @@ def create_download_zip(server_metadata: Dict[str, Any]) -> Optional[str]:
     Create a ZIP file of the generated MCP server for download.
 
     Returns:
-        Path to the ZIP file as a string, or None if creation fails.
+        Path to the ZIP file, or None if creation fails
     """
     try:
         server_dir = Path(server_metadata["directory"])
         server_id = server_metadata["server_id"]
         zip_path = server_dir.parent / f"{server_id}.zip"
 
+        # Make sure parent exists
+        zip_path.parent.mkdir(parents=True, exist_ok=True)
+
         with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
             for file_path in server_dir.rglob("*"):
                 if file_path.is_file():
                     arcname = file_path.relative_to(server_dir.parent)
                     zipf.write(file_path, arcname)
 
-        print(f"[OK] Created ZIP at {zip_path}")
+        print(f"[ZIP] Created MCP archive at {zip_path}")
         return str(zip_path)
     except Exception as e:
         print(f"[ERROR] Failed to create ZIP: {e}")
         return None
 
 
-def push_zip_to_space_repo(zip_path: str) -> Optional[str]:
+def push_zip_to_space_repo(zip_path: Path) -> Optional[str]:
     """
-    Upload the generated ZIP into this Space's git repo under generated_mcps/.
+    Upload the generated ZIP into the same Space repository using HF API.
 
     Uses:
-      - SPACE_ID      (provided automatically by HF for Spaces)
-      - HF_WRITE_TOKEN (your write token, set as a secret)
-        or HF_TOKEN as a fallback.
+      - SPACE_ID       (your space id, e.g. MCP-1st-Birthday/OmniMind-Orchestrator)
+      - HF_WRITE_TOKEN (or HF_TOKEN) with write permission
 
     Returns:
-        Public Hub URL, or None if upload failed or skipped.
+        A Hub URL (string) pointing to the uploaded file, or None on failure.
     """
-    if not zip_path:
-        return None
-
-    space_id = os.getenv("SPACE_ID")  # auto-injected in Spaces
-    token = os.getenv("HF_WRITE_TOKEN") or os.getenv("HF_TOKEN")
-
-    if HfApi is None:
-        print("[HF] huggingface_hub not installed - skipping Hub upload.")
+    try:
+        from huggingface_hub import HfApi
+    except Exception as e:
+        print(f"[HF] huggingface_hub not installed - skipping Space upload ({e})")
         return None
 
+    space_id = os.getenv("SPACE_ID")
     if not space_id:
-        print("[HF] SPACE_ID env var not found – skipping Hub upload.")
+        print("[HF] SPACE_ID env var not set - skipping Space upload")
         return None
 
+    token = os.getenv("HF_WRITE_TOKEN") or os.getenv("HF_TOKEN")
     if not token:
-        print("[HF] HF_WRITE_TOKEN / HF_TOKEN not set – skipping Hub upload.")
+        print("[HF] HF_WRITE_TOKEN / HF_TOKEN not set - skipping Space upload")
         return None
 
-    try:
-        api = HfApi(token=token)
+    api = HfApi(token=token)
 
-        zip_path_obj = Path(zip_path)
-        path_in_repo = f"generated_mcps/{zip_path_obj.name}"
+    # We store files under generated_mcps/ in the repo
+    path_in_repo = f"generated_mcps/{zip_path.name}"
 
-        api.upload_file(
-            path_or_fileobj=str(zip_path_obj),
+    try:
+        hub_url = api.upload_file(
+            path_or_fileobj=str(zip_path),
             path_in_repo=path_in_repo,
             repo_id=space_id,
             repo_type="space",
         )
-
-        hub_url = f"https://huggingface.co/spaces/{space_id}/blob/main/{path_in_repo}"
-        print(f"[HF] Uploaded ZIP to Hub: {hub_url}")
-        return hub_url
-
+        # `hub_url` is already the correct URL to visualize the uploaded file
+        print(f"[HF] Uploaded ZIP to Hub at: {hub_url}")
+        return str(hub_url)
     except Exception as e:
-        print(f"[HF][ERROR] Failed to upload ZIP to Hub: {e}")
+        print(f"[HF] Failed to upload ZIP to Hub: {e}")
         return None
 
 
-# =============================================================================
+# ============================================================================
 # Agent Visualization (Blaxel Integration)
-# =============================================================================
+# ============================================================================
+
 
 def create_agent_graph(agent_state: Dict[str, Any]) -> go.Figure:
     """
@@ -195,8 +173,8 @@ def create_agent_graph(agent_state: Dict[str, Any]) -> go.Figure:
     node_colors = []
 
     color_map = {
-        "planning": "#3B82F6",   # Blue
-        "generating": "#10B981", # Green
+        "planning": "#3B82F6",  # Blue
+        "generating": "#10B981",  # Green
         "deploying": "#F59E0B",  # Orange
         "executing": "#8B5CF6",  # Purple
         "completed": "#6B7280",  # Gray
@@ -217,11 +195,7 @@ def create_agent_graph(agent_state: Dict[str, Any]) -> go.Figure:
         hoverinfo="text",
         text=node_text,
         textposition="top center",
-        marker=dict(
-            size=30,
-            color=node_colors,
-            line=dict(width=2, color="white"),
-        ),
+        marker=dict(size=30, color=node_colors, line=dict(width=2, color="white")),
     )
 
     # Create figure
@@ -242,9 +216,10 @@ def create_agent_graph(agent_state: Dict[str, Any]) -> go.Figure:
     return fig
 
 
-# =============================================================================
+# ============================================================================
 # Core Agent Orchestration
-# =============================================================================
+# ============================================================================
+
 
 async def orchestrate_task(
     user_request: str,
@@ -255,7 +230,7 @@ async def orchestrate_task(
     Main orchestration function - the brain of OmniMind.
 
     Yields:
-        (status_text, agent_graph, metadata, download_zip_path)
+        (status_text, agent_graph, metadata, zip_path_for_download)
     """
     output = "# 🤖 OmniMind Orchestrator\n\n"
     output += f"**Request:** {user_request}\n\n"
@@ -266,7 +241,6 @@ async def orchestrate_task(
         "edges": [],
     }
 
-    # Initial state
     yield (output, create_agent_graph(agent_state), {}, None)
 
     # Step 1: Analyze request with multi-model router
@@ -294,7 +268,9 @@ Respond in JSON:
 """
 
     analysis = await router.generate(
-        analysis_prompt, task_type=TaskType.PLANNING, temperature=0.3
+        analysis_prompt,
+        task_type=TaskType.PLANNING,
+        temperature=0.3,
     )
 
     try:
@@ -317,12 +293,7 @@ Respond in JSON:
     )
     agent_state["edges"].append({"from": "start", "to": "analyze"})
 
-    yield (
-        output,
-        create_agent_graph(agent_state),
-        to_jsonable(analysis_data),
-        None,
-    )
+    yield (output, create_agent_graph(agent_state), analysis_data, None)
 
     # Step 2: Get knowledge context (if enabled)
     context = None
@@ -341,15 +312,10 @@ Respond in JSON:
             output += "**No relevant context found**\n\n"
 
         agent_state["nodes"][-1]["type"] = "completed"
-        yield (
-            output,
-            create_agent_graph(agent_state),
-            {"has_context": bool(context)},
-            None,
-        )
+        yield (output, create_agent_graph(agent_state), {}, None)
 
     # Step 3: Generate MCP (if needed)
-    server_metadata: Optional[Dict[str, Any]] = None
+    server_metadata = None
     zip_path: Optional[str] = None
 
     if analysis_data.get("needs_custom_mcp", False):
@@ -370,8 +336,7 @@ Respond in JSON:
 
         output += f"✅ **Generated:** {server_metadata['server_name']}\n"
         output += (
-            f"**Tools:** "
-            f"{', '.join([t['name'] for t in server_metadata['tools']])}\n"
+            f"**Tools:** {', '.join([t['name'] for t in server_metadata['tools']])}\n"
         )
         output += f"**Location:** `{server_metadata['directory']}`\n\n"
 
@@ -379,9 +344,10 @@ Respond in JSON:
         output += "### 📄 Generated Code Preview\n\n"
         output += "```python\n"
         try:
-            app_file = server_metadata["files"]["app"]
-            with open(app_file, "r", encoding="utf-8") as f:
-                code_lines = f.readlines()[:30]  # Show first 30 lines
+            with open(
+                server_metadata["files"]["app"], "r", encoding="utf-8"
+            ) as f_code:
+                code_lines = f_code.readlines()[:30]  # Show first 30 lines
                 output += "".join(code_lines)
                 if len(code_lines) >= 30:
                     output += "\n... (truncated - full code saved locally)\n"
@@ -399,25 +365,22 @@ Respond in JSON:
             server_metadata["zip_path"] = zip_path
             output += "📦 **Download button updated below!**\n\n"
 
-            # Try to push to Hub as well (so it appears in Files tab)
-            hub_url = push_zip_to_space_repo(zip_path)
+            # Upload ZIP into this same Space repo (if configured)
+            hub_url = push_zip_to_space_repo(Path(zip_path))
             if hub_url:
                 server_metadata["hub_url"] = hub_url
-                output += f"📡 **Saved to Hub:** {hub_url}\n\n"
+                output += f"🔗 **Saved to Hub:** {hub_url}\n\n"
             else:
                 output += (
-                    "⚠️ Could not push ZIP to the Hub repo. "
-                    "You can still use the download button and upload manually.\n\n"
+                    "⚠️ **Could not save ZIP to Hub (check Space logs / env vars).**\n\n"
                 )
-        else:
-            output += "⚠️ Failed to create ZIP archive for download.\n\n"
 
         agent_state["nodes"][-1]["type"] = "completed"
         yield (
             output,
             create_agent_graph(agent_state),
-            to_jsonable(server_metadata),
-            zip_path,
+            server_metadata,
+            zip_path if zip_path else None,
         )
 
         # Step 4: Deploy to Modal
@@ -426,29 +389,25 @@ Respond in JSON:
             {"id": "deploy", "label": "Deploy", "type": "deploying"}
         )
         agent_state["edges"].append({"from": "generate", "to": "deploy"})
-        yield (output, create_agent_graph(agent_state), {}, zip_path)
+        yield (output, create_agent_graph(agent_state), {}, None)
 
         deployment = await deployer.deploy_mcp_server(server_metadata)
 
         if deployment.get("simulated"):
-            output += "⚠️  **Simulated deployment** (configure MODAL_TOKEN for real deployment)\n"
+            output += (
+                "⚠️  **Simulated deployment** (configure MODAL_TOKEN for real deployment)\n"
+            )
 
         if deployment.get("status") == "failed":
             output += (
-                f"⚠️  **Deployment skipped:** "
-                f"{deployment.get('error', 'Unknown error')}\n\n"
+                f"⚠️  **Deployment skipped:** {deployment.get('error', 'Unknown error')}\n\n"
             )
         else:
             output += f"**URL:** {deployment.get('modal_url', 'N/A')}\n"
             output += f"**Status:** {deployment.get('status', 'unknown')}\n\n"
 
         agent_state["nodes"][-1]["type"] = "completed"
-        yield (
-            output,
-            create_agent_graph(agent_state),
-            to_jsonable(deployment),
-            zip_path,
-        )
+        yield (output, create_agent_graph(agent_state), deployment, zip_path)
 
     # Step 5: Final response generation
     output += "## ✨ Step 5: Generating Response\n\n"
@@ -474,7 +433,9 @@ Provide a helpful response explaining what was accomplished and how the user can
 """
 
     final_response = await router.generate(
-        response_prompt, task_type=TaskType.REASONING, temperature=0.7
+        response_prompt,
+        task_type=TaskType.REASONING,
+        temperature=0.7,
     )
 
     output += final_response["response"] + "\n\n"
@@ -486,27 +447,29 @@ Provide a helpful response explaining what was accomplished and how the user can
     if use_voice and voice.client:
         output += "\n🔊 **Generating voice response...**\n"
         yield (output, create_agent_graph(agent_state), {}, zip_path)
-        # Voice generation would happen here (skipped in this demo)
+        # (Voice generation omitted in this demo)
 
     output += "\n---\n\n"
     output += "**Model Usage:**\n"
     stats = router.get_usage_stats()
     output += f"- Total Requests: {stats['total_requests']}\n"
     output += f"- Total Cost: ${stats['total_cost']}\n"
-    output += f"- Claude: {stats['by_model']['claude']['requests']}\n"
-    output += f"- Gemini: {stats['by_model']['gemini']['requests']}\n"
-    output += f"- GPT-4: {stats['by_model']['gpt4']['requests']}\n"
+    output += f"- Claude: {stats['by_model']['claude']['requests']} requests\n"
+    output += f"- Gemini: {stats['by_model']['gemini']['requests']} requests\n"
+    output += f"- GPT-4: {stats['by_model']['gpt4']['requests']} requests\n"
 
-    yield (output, create_agent_graph(agent_state), to_jsonable(stats), zip_path)
+    yield (output, create_agent_graph(agent_state), stats, zip_path)
 
 
-# =============================================================================
+# ============================================================================
 # Gradio UI
-# =============================================================================
+# ============================================================================
+
 
 def build_ui() -> gr.Blocks:
-    """Build the Gradio interface"""
+    """Build the Gradio 6 interface"""
 
+    # Custom CSS for professional look
     custom_css = """
     .gradio-container {
         font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
@@ -528,7 +491,7 @@ def build_ui() -> gr.Blocks:
     """
 
     with gr.Blocks(title="OmniMind Orchestrator - MCP Hackathon") as app:
-        # Inject CSS via HTML instead of css= argument (for older Gradio versions)
+        # Inject CSS (Blocks in this Gradio version doesn't accept css=...)
         gr.HTML(f"<style>{custom_css}</style>")
 
         gr.HTML(
@@ -536,7 +499,7 @@ def build_ui() -> gr.Blocks:
         <div class="main-header">
             <h1>🧠 OmniMind Orchestrator</h1>
             <p>The World's First Self-Evolving Multi-Agent MCP Ecosystem</p>
-            <p style="font-size: 0.9em; opacity: 0.9%;">
+            <p style="font-size: 0.9em; opacity: 0.9;">
                 Track 2 Submission - MCP's 1st Birthday Hackathon
             </p>
         </div>
@@ -563,10 +526,7 @@ def build_ui() -> gr.Blocks:
 
                 user_input = gr.Textbox(
                     label="What do you need?",
-                    placeholder=(
-                        "Example: Create a tool that monitors my competitor's "
-                        "pricing every hour"
-                    ),
+                    placeholder="Example: Create a tool that monitors my competitor's pricing every hour",
                     lines=3,
                 )
 
@@ -575,8 +535,7 @@ def build_ui() -> gr.Blocks:
                     use_kb = gr.Checkbox(label="📚 Use Knowledge Base", value=False)
 
                 submit_btn = gr.Button(
-                    "🚀 Let OmniMind Handle It",
-                    variant="primary",
+                    "🚀 Let OmniMind Handle It", variant="primary", size="lg"
                 )
 
                 gr.Markdown(
@@ -592,14 +551,12 @@ def build_ui() -> gr.Blocks:
 
             with gr.Column(scale=2):
                 output_md = gr.Markdown(
-                    value="**Results will appear here**",
-                    label="Agent Output",
+                    value="**Results will appear here**", label="Agent Output"
                 )
 
                 agent_graph = gr.Plot(label="🧠 Agent Brain (Real-Time)")
 
-                # Use DownloadButton for file download
-                download_btn = gr.DownloadButton(
+                download_file = gr.File(
                     label="📦 Download Generated MCP Server",
                     visible=False,
                 )
@@ -667,12 +624,12 @@ def build_ui() -> gr.Blocks:
             1. **Never-Before-Done**: First agent that creates agents
             2. **All Sponsors**: Uses every sponsor technology meaningfully
             3. **Real Impact**: Saves enterprises weeks of custom development
-            4. **Beautiful UX**: Gradio + voice + visualizations
+            4. **Beautiful UX**: Gradio 6 + voice + visualizations
             5. **Production-Ready**: Clean code, error handling, scalability
 
             ### Built With
 
-            - Gradio
+            - Gradio 6.0
             - LangGraph
             - Claude Sonnet 4
             - Gemini 2.0 Flash
@@ -694,39 +651,37 @@ def build_ui() -> gr.Blocks:
         async def handle_submit(request, voice_enabled, kb_enabled):
             """Handle user submissions"""
             async for (
-                output,
+                out_text,
                 graph,
                 metadata,
                 download_zip,
             ) in orchestrate_task(request, voice_enabled, kb_enabled):
-
                 if download_zip:
-                    # Provide a file path to DownloadButton
                     yield (
-                        output,
+                        out_text,
                         graph,
                         metadata,
                         gr.update(value=download_zip, visible=True),
                     )
                 else:
                     yield (
-                        output,
+                        out_text,
                         graph,
                         metadata,
-                        gr.update(visible=False),
+                        gr.update(value=None, visible=False),
                     )
 
         submit_btn.click(
             fn=handle_submit,
             inputs=[user_input, use_voice, use_kb],
-            outputs=[output_md, agent_graph, metadata_json, download_btn],
+            outputs=[output_md, agent_graph, metadata_json, download_file],
         )
 
         gr.Markdown(
             """
         ---
         <div style="text-align: center; padding: 1rem; color: #666;">
-            🎉 Built for MCP's 1st Birthday Hackathon | Hosted by Anthropic &amp; Gradio
+            🎉 Built for MCP's 1st Birthday Hackathon | Hosted by Anthropic & Gradio
         </div>
         """
         )
@@ -734,9 +689,10 @@ def build_ui() -> gr.Blocks:
     return app
 
 
-# =============================================================================
+# ============================================================================
 # Main Execution
-# =============================================================================
+# ============================================================================
+
 
 if __name__ == "__main__":
     print("=" * 60)