Update app.py

app.py

@@ -1,609 +1,608 @@
-# app.py - MasterLLM v2.0 with Bedrock Fallback System
-"""
-MasterLLM Pipeline Orchestrator v2.0
-- Bedrock (priority) + Gemini (fallback) for pipeline generation
-- Bedrock LangChain (priority) + CrewAI (fallback) for execution
-- MongoDB session management
-- Complete REST API
-- Gradio UI with fancy displays
-"""
-import os
-import json
-import uuid
-from datetime import datetime
-from typing import List, Optional
-
-import gradio as gr
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-from contextlib import asynccontextmanager
-import asyncio
-
-# Import our new services
-from services.pipeline_generator import generate_pipeline, format_pipeline_for_display
-from services.pipeline_executor import execute_pipeline_streaming
-from services.session_manager import session_manager
-from api_routes import router as api_router
-
-
-# ========================
-# BACKGROUND CLEANUP TASK
-# ========================
-
-async def periodic_cleanup():
-    """Cleanup old sessions every hour"""
-    while True:
-        await asyncio.sleep(3600)  # Run every hour
-        try:
-            removed = session_manager.cleanup_old_sessions(max_age_hours=24)
-            if removed > 0:
-                print(f"🧹 Cleaned up {removed} inactive sessions")
-        except Exception as e:
-            print(f"⚠️  Cleanup error: {e}")
-
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    """Manage application lifecycle"""
-    # Startup
-    print("🚀 Starting MasterLLM v2.0...")
-    task = asyncio.create_task(periodic_cleanup())
-    yield
-    # Shutdown
-    task.cancel()
-    session_manager.close()
-    print("🛑 MasterLLM shut down gracefully")
-
-
-# ========================
-# FASTAPI APP
-# ========================
-
-app = FastAPI(
-    title="MasterLLM v2.0 - AI Pipeline Orchestrator",
-    description="Bedrock + Gemini fallback system with MongoDB sessions",
-    version="2.0.0",
-    lifespan=lifespan
-)
-
-# CORS Configuration
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=[os.getenv("FRONTEND_ORIGIN", "http://localhost:3000")],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Mount API routes
-app.include_router(api_router)
-
-
-# ========================
-# CONVERSATION STATE
-# ========================
-
-class ConversationState:
-    INITIAL = "initial"
-    PIPELINE_PROPOSED = "pipeline_proposed"
-    PIPELINE_APPROVED = "pipeline_approved"
-    EXECUTING = "executing"
-    COMPLETED = "completed"
-    ERROR = "error"
-
-
-# ========================
-# GRADIO UI HANDLERS
-# ========================
-
-def create_new_session():
-    """Create a new session"""
-    return session_manager.create_session()
-
-
-def handle_file_upload(file_path, session_id):
-    """Handle file upload"""
-    if not file_path:
-        return None, json.dumps({
-            "status": "error",
-            "message": "No file uploaded"
-        }, indent=2), session_id
-    
-    if not session_id:
-        session_id = create_new_session()
-    
-    file_name = os.path.basename(file_path)
-    
-    # Update session
-    session_manager.update_session(session_id, {
-        "current_file": file_path,
-        "state": ConversationState.INITIAL
-    })
-    
-    # Add system message
-    session_manager.add_message(
-        session_id,
-        "system",
-        f"File uploaded: {file_name}"
-    )
-    
-    status = {
-        "status": "success",
-        "message": f"File '{file_name}' uploaded successfully",
-        "file_info": {
-            "name": file_name,
-            "path": file_path,
-            "size_bytes": os.path.getsize(file_path) if os.path.exists(file_path) else 0
-        },
-        "next_action": "💬 Now tell me what you'd like to do with this document"
-    }
-    
-    return file_path, json.dumps(status, indent=2), session_id
-
-
-def chatbot_response_streaming(message: str, history: List, session_id: str, file_path: str = None):
-    """
-    Handle chat messages with streaming updates
-    Uses Bedrock (priority) → Gemini (fallback) for both generation and execution
-    """
-    # Get or create session
-    session = session_manager.get_session(session_id)
-    if not session:
-        session_id = create_new_session()
-        session = session_manager.get_session(session_id)
-    
-    # Update file path if provided
-    if file_path:
-        session_manager.update_session(session_id, {"current_file": file_path})
-        session = session_manager.get_session(session_id)
-    
-    # Add user message to session
-    session_manager.add_message(session_id, "user", message)
-    
-    current_state = session.get("state", ConversationState.INITIAL)
-    
-    # ========================
-    # STATE: INITIAL - Generate Pipeline
-    # ========================
-    if current_state == ConversationState.INITIAL:
-        # Check if file is uploaded
-        if not session.get("current_file"):
-            response = {
-                "status": "error",
-                "message": "Please upload a document first",
-                "action": "📁 Click 'Upload Document' to begin"
-            }
-            response_text = f"```json\n{json.dumps(response, indent=2)}\n```"
-            session_manager.add_message(session_id, "assistant", response_text)
-            yield history + [[message, response_text]]
-            return
-        
-        try:
-            # Generate pipeline using Bedrock → Gemini fallback
-            yield history + [[message, "🤖 Generating pipeline with AI...\n⏳ Trying Bedrock first..."]]
-            
-            pipeline = generate_pipeline(
-                user_input=message,
-                file_path=session.get("current_file"),
-                prefer_bedrock=True
-            )
-            
-            # Save proposed pipeline to session
-            session_manager.update_session(session_id, {
-                "proposed_pipeline": pipeline,
-                "state": ConversationState.PIPELINE_PROPOSED
-            })
-            
-            # Format for display
-            formatted_display = format_pipeline_for_display(pipeline)
-            
-            # Create response with both fancy display and JSON
-            response_text = formatted_display + f"\n\n```json\n{json.dumps(pipeline, indent=2)}\n```"
-            
-            session_manager.add_message(session_id, "assistant", response_text)
-            yield history + [[message, response_text]]
-            return
-            
-        except Exception as e:
-            error_response = {
-                "status": "error",
-                "message": "Failed to generate pipeline",
-                "error": str(e),
-                "action": "Please try rephrasing your request"
-            }
-            response_text = f"```json\n{json.dumps(error_response, indent=2)}\n```"
-            session_manager.add_message(session_id, "assistant", response_text)
-            yield history + [[message, response_text]]
-            return
-    
-    # ========================
-    # STATE: PIPELINE_PROPOSED - Handle Approval/Rejection
-    # ========================
-    elif current_state == ConversationState.PIPELINE_PROPOSED:
-        user_input = message.lower().strip()
-        
-        # APPROVE - Execute the pipeline
-        if "approve" in user_input or "yes" in user_input:
-            session_manager.update_session(session_id, {"state": ConversationState.EXECUTING})
-            
-            plan = session.get("proposed_pipeline", {})
-            
-            # Initial status
-            initial_status = {
-                "status": "executing",
-                "message": "🚀 Starting pipeline execution...",
-                "pipeline": plan.get("pipeline_name", "unknown"),
-                "executor": "Attempting Bedrock LangChain first",
-                "steps": []
-            }
-            accumulated_response = f"```json\n{json.dumps(initial_status, indent=2)}\n```"
-            yield history + [[message, accumulated_response]]
-            
-            steps_completed = []
-            final_payload = None
-            executor_used = "unknown"
-            
-            try:
-                # Execute pipeline with Bedrock → CrewAI fallback
-                for event in execute_pipeline_streaming(
-                    pipeline=plan,
-                    file_path=session.get("current_file"),
-                    session_id=session_id,
-                    prefer_bedrock=True
-                ):
-                    event_type = event.get("type")
-                    
-                    # Info events (fallback notifications, etc.)
-                    if event_type == "info":
-                        info_status = {
-                            "status": "info",
-                            "message": event.get("message"),
-                            "executor": event.get("executor", "unknown")
-                        }
-                        accumulated_response = f"```json\n{json.dumps(info_status, indent=2)}\n```"
-                        yield history + [[message, accumulated_response]]
-                    
-                    # Step updates
-                    elif event_type == "step":
-                        step_info = {
-                            "step": event.get("step", 0),
-                            "tool": event.get("tool", "processing"),
-                            "status": event.get("status", "running"),
-                            "executor": event.get("executor", "unknown")
-                        }
-                        steps_completed.append(step_info)
-                        executor_used = event.get("executor", executor_used)
-                        
-                        progress_status = {
-                            "status": "executing",
-                            "message": f"📍 Step {event.get('step', 0)}: {event.get('tool', 'processing')}...",
-                            "pipeline": plan.get("pipeline_name", ""),
-                            "executor": executor_used,
-                            "steps_completed": steps_completed
-                        }
-                        accumulated_response = f"```json\n{json.dumps(progress_status, indent=2)}\n```"
-                        yield history + [[message, accumulated_response]]
-                    
-                    # Final result
-                    elif event_type == "final":
-                        final_payload = event.get("data")
-                        executor_used = event.get("executor", executor_used)
-                    
-                    # Error
-                    elif event_type == "error":
-                        error_result = {
-                            "status": "failed",
-                            "error": event.get("error"),
-                            "steps_completed": steps_completed,
-                            "executor": event.get("executor", "unknown")
-                        }
-                        final_response = f"```json\n{json.dumps(error_result, indent=2)}\n```"
-                        session_manager.update_session(session_id, {"state": ConversationState.INITIAL})
-                        session_manager.add_message(session_id, "assistant", final_response)
-                        yield history + [[message, final_response]]
-                        return
-                
-                # Process final result
-                if final_payload:
-                    session_manager.update_session(session_id, {
-                        "pipeline_result": final_payload,
-                        "state": ConversationState.INITIAL
-                    })
-                    
-                    # Save execution to MongoDB
-                    session_manager.save_pipeline_execution(
-                        session_id=session_id,
-                        pipeline=plan,
-                        result=final_payload,
-                        file_path=session.get("current_file"),
-                        executor=executor_used
-                    )
-                    
-                    # Format final response
-                    final_display = {
-                        "status": "completed",
-                        "executor": executor_used,
-                        "pipeline": plan.get("pipeline_name"),
-                        "result": final_payload,
-                        "summary": {
-                            "total_steps": len(steps_completed),
-                            "completed_successfully": len([s for s in steps_completed if s.get("status") == "completed"])
-                        }
-                    }
-                    final_response = f"```json\n{json.dumps(final_display, indent=2)}\n```"
-                else:
-                    final_response = f"```json\n{json.dumps({'status': 'completed', 'steps': steps_completed, 'executor': executor_used}, indent=2)}\n```"
-                    session_manager.update_session(session_id, {"state": ConversationState.INITIAL})
-                
-                session_manager.add_message(session_id, "assistant", final_response)
-                yield history + [[message, final_response]]
-                return
-                
-            except Exception as e:
-                error_result = {
-                    "error": str(e),
-                    "status": "failed",
-                    "message": "Pipeline execution failed",
-                    "steps_completed": steps_completed
-                }
-                final_response = f"```json\n{json.dumps(error_result, indent=2)}\n```"
-                session_manager.update_session(session_id, {"state": ConversationState.INITIAL})
-                session_manager.add_message(session_id, "assistant", final_response)
-                yield history + [[message, final_response]]
-                return
-        
-        # REJECT - Cancel the pipeline
-        elif "reject" in user_input or "no" in user_input:
-            session_manager.update_session(session_id, {
-                "state": ConversationState.INITIAL,
-                "proposed_pipeline": None
-            })
-            response_data = {
-                "status": "rejected",
-                "message": "Pipeline rejected by user",
-                "action": "💬 Please provide a new instruction"
-            }
-            response = f"```json\n{json.dumps(response_data, indent=2)}\n```"
-            session_manager.add_message(session_id, "assistant", response)
-            yield history + [[message, response]]
-            return
-        
-        # EDIT - Request modifications
-        elif "edit" in user_input or "modify" in user_input:
-            current_pipeline = session.get("proposed_pipeline", {})
-            edit_help = {
-                "status": "edit_mode",
-                "message": "To modify the plan, describe your changes",
-                "current_plan": current_pipeline,
-                "examples": [
-                    "Add summarization at the end",
-                    "Remove table extraction",
-                    "Only process pages 1-3",
-                    "Translate to French instead of Spanish"
-                ],
-                "action": "Describe your changes, or say 'approve' to run as-is"
-            }
-            response = f"```json\n{json.dumps(edit_help, indent=2)}\n```"
-            session_manager.add_message(session_id, "assistant", response)
-            yield history + [[message, response]]
-            return
-        
-        # Try to modify pipeline based on user input
-        else:
-            if len(message.strip()) > 5:
-                try:
-                    original_plan = session.get("proposed_pipeline", {})
-                    edit_context = f"Original: {original_plan.get('pipeline_name')}. User wants: {message}"
-                    
-                    # Generate new pipeline with modification
-                    new_pipeline = generate_pipeline(
-                        user_input=edit_context,
-                        file_path=session.get("current_file"),
-                        prefer_bedrock=True
-                    )
-                    
-                    session_manager.update_session(session_id, {
-                        "proposed_pipeline": new_pipeline,
-                        "state": ConversationState.PIPELINE_PROPOSED
-                    })
-                    
-                    formatted = format_pipeline_for_display(new_pipeline)
-                    response = formatted + f"\n\n```json\n{json.dumps(new_pipeline, indent=2)}\n```"
-                    session_manager.add_message(session_id, "assistant", response)
-                    yield history + [[message, response]]
-                    return
-                    
-                except Exception as e:
-                    error_response = {
-                        "status": "edit_failed",
-                        "error": str(e),
-                        "message": "Could not modify the plan",
-                        "action": "Try 'approve' to run as-is, or 'reject' to start over"
-                    }
-                    response = f"```json\n{json.dumps(error_response, indent=2)}\n```"
-                    session_manager.add_message(session_id, "assistant", response)
-                    yield history + [[message, response]]
-                    return
-            
-            # Default waiting message
-            response_data = {
-                "status": "waiting_for_confirmation",
-                "message": "Please type 'approve', 'reject', or describe changes",
-                "hint": "You can also say 'edit' for modification hints"
-            }
-            response = f"```json\n{json.dumps(response_data, indent=2)}\n```"
-            session_manager.add_message(session_id, "assistant", response)
-            yield history + [[message, response]]
-            return
-    
-    # Default fallback
-    response = json.dumps({"status": "ready", "message": "Ready for your next instruction"}, indent=2)
-    session_manager.add_message(session_id, "assistant", response)
-    yield history + [[message, response]]
-
-
-# ========================
-# GRADIO UI
-# ========================
-
-with gr.Blocks(
-    title="MasterLLM v2.0 - AI Pipeline Orchestrator",
-    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
-    css="""
-    .gradio-container {
-        max-width: 1400px !important;
-    }
-    """
-) as demo:
-    gr.Markdown("""
-    # 🤖 MasterLLM v2.0 - AI Pipeline Orchestrator
-    
-    **🏆 Bedrock Priority** with Gemini Fallback | **💾 MongoDB Sessions** | **📡 Complete REST API**
-    
-    Upload a document, describe what you want, and watch AI orchestrate the perfect pipeline!
-    """)
-    
-    # State management
-    session_id_state = gr.State(value=create_new_session())
-    file_state = gr.State(value=None)
-    
-    with gr.Row():
-        with gr.Column(scale=3):
-            # Chat interface
-            chatbot = gr.Chatbot(
-                height=650,
-                show_label=False,
-                bubble_full_width=False,
-                show_copy_button=True,
-                avatar_images=(None, "🤖"),
-            )
-            
-            # Text input
-            msg = gr.Textbox(
-                placeholder="💬 Type your instruction... (e.g., 'extract text from pages 1-5 and summarize')",
-                show_label=False,
-                lines=2,
-                max_lines=4,
-                container=False,
-            )
-            
-            with gr.Row():
-                submit_btn = gr.Button("🚀 Send", variant="primary", scale=2)
-                clear_btn = gr.Button("🗑️ Clear Chat", scale=1)
-        
-        with gr.Column(scale=1):
-            # File upload section
-            gr.Markdown("### 📁 Upload Document")
-            file_upload = gr.File(
-                label="PDF or Image",
-                file_types=[".pdf", ".png", ".jpg", ".jpeg", ".gif", ".bmp"],
-                type="filepath",
-            )
-            
-            upload_status = gr.Textbox(
-                label="📊 Upload Status",
-                interactive=False,
-                lines=10,
-                max_lines=15,
-            )
-            
-            # Session info
-            gr.Markdown("### 🔗 Session Info")
-            session_display = gr.Textbox(
-                label="Session ID",
-                interactive=False,
-                value=lambda: session_id_state.value[:8] + "...",
-            )
-            
-            # Examples
-            gr.Markdown("### 💡 Example Pipelines")
-            gr.Examples(
-                examples=[
-                    "extract text from pages 1-5",
-                    "extract text and summarize",
-                    "extract text, tables, and translate to Spanish",
-                    "get tables from pages 2-4 and summarize",
-                    "text-classify-ner from entire document",
-                    "describe images and summarize findings",
-                    "extract text, detect signatures and stamps",
-                ],
-                inputs=msg,
-            )
-            
-            # System info
-            gr.Markdown("""
-            ### ℹ️ System Features
-            - ✅ **Bedrock** (Claude 3.5 Sonnet) priority
-            - ✅ **Gemini** (gemini-2.0-flash) fallback
-            - ✅ **MongoDB** session persistence
-            - ✅ **Streaming** real-time updates
-            - ✅ **Component-level** JSON output
-            - ✅ **REST API** for integration
-            
-            ### 📊 Pipeline Flow:
-            1. **Upload** your document
-            2. **Describe** what you want
-            3. **Review** AI-generated pipeline
-            4. **Approve** to execute
-            5. **Watch** streaming updates
-            6. **Get** complete JSON results
-            """)
-    
-    # Event handlers
-    file_upload.upload(
-        fn=handle_file_upload,
-        inputs=[file_upload, session_id_state],
-        outputs=[file_state, upload_status, session_id_state],
-    )
-    
-    msg.submit(
-        fn=chatbot_response_streaming,
-        inputs=[msg, chatbot, session_id_state, file_state],
-        outputs=[chatbot],
-    ).then(
-        lambda: "",
-        outputs=msg,
-    )
-    
-    submit_btn.click(
-        fn=chatbot_response_streaming,
-        inputs=[msg, chatbot, session_id_state, file_state],
-        outputs=[chatbot],
-    ).then(
-        lambda: "",
-        outputs=msg,
-    )
-    
-    clear_btn.click(
-        fn=lambda: ([], create_new_session(), None, None, "", ""),
-        outputs=[chatbot, session_id_state, file_state, file_upload, msg, upload_status],
-    )
-
-# Mount Gradio on FastAPI
-app = gr.mount_gradio_app(app, demo, path="/")
-
-
-# ========================
-# LAUNCH
-# ========================
-
-if __name__ == "__main__":
-    import uvicorn
-    port = int(os.getenv("PORT", 7860))
-    print(f"""
-    ╔════════════════════════════════════════════════════════════╗
-    ║                                                            ║
-    ║           🚀 MasterLLM v2.0 Starting...                   ║
-    ║                                                            ║
-    ║  🌐 Gradio UI:  http://localhost:{port}                      ║
-    ║  📡 REST API:   http://localhost:{port}/api/v1                ║
-    ║  📚 API Docs:   http://localhost:{port}/docs                  ║
-    ║                                                            ║
-    ║  🏆 Bedrock:    Priority (Claude 3.5 Sonnet)             ║
-    ║  🔄 Gemini:     Fallback (gemini-2.0-flash)              ║
-    ║  💾 MongoDB:    Session management                         ║
-    ║                                                            ║
-    ╚════════════════════════════════════════════════════════════╝
-    """)
-    
+# app.py - MasterLLM v2.0 with Bedrock Fallback System
+"""
+MasterLLM Pipeline Orchestrator v2.0
+- Bedrock (priority) + Gemini (fallback) for pipeline generation
+- Bedrock LangChain (priority) + CrewAI (fallback) for execution
+- MongoDB session management
+- Complete REST API
+- Gradio UI with fancy displays
+"""
+import os
+import json
+import uuid
+from datetime import datetime
+from typing import List, Optional
+
+import gradio as gr
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from contextlib import asynccontextmanager
+import asyncio
+
+# Import our new services
+from services.pipeline_generator import generate_pipeline, format_pipeline_for_display
+from services.pipeline_executor import execute_pipeline_streaming
+from services.session_manager import session_manager
+from api_routes import router as api_router
+
+
+# ========================
+# BACKGROUND CLEANUP TASK
+# ========================
+
+async def periodic_cleanup():
+    """Cleanup old sessions every hour"""
+    while True:
+        await asyncio.sleep(3600)  # Run every hour
+        try:
+            removed = session_manager.cleanup_old_sessions(max_age_hours=24)
+            if removed > 0:
+                print(f"🧹 Cleaned up {removed} inactive sessions")
+        except Exception as e:
+            print(f"⚠️  Cleanup error: {e}")
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Manage application lifecycle"""
+    # Startup
+    print("🚀 Starting MasterLLM v2.0...")
+    task = asyncio.create_task(periodic_cleanup())
+    yield
+    # Shutdown
+    task.cancel()
+    session_manager.close()
+    print("🛑 MasterLLM shut down gracefully")
+
+
+# ========================
+# FASTAPI APP
+# ========================
+
+app = FastAPI(
+    title="MasterLLM v2.0 - AI Pipeline Orchestrator",
+    description="Bedrock + Gemini fallback system with MongoDB sessions",
+    version="2.0.0",
+    lifespan=lifespan
+)
+
+# CORS Configuration
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=[os.getenv("FRONTEND_ORIGIN", "http://localhost:3000")],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Mount API routes
+app.include_router(api_router)
+
+
+# ========================
+# CONVERSATION STATE
+# ========================
+
+class ConversationState:
+    INITIAL = "initial"
+    PIPELINE_PROPOSED = "pipeline_proposed"
+    PIPELINE_APPROVED = "pipeline_approved"
+    EXECUTING = "executing"
+    COMPLETED = "completed"
+    ERROR = "error"
+
+
+# ========================
+# GRADIO UI HANDLERS
+# ========================
+
+def create_new_session():
+    """Create a new session"""
+    return session_manager.create_session()
+
+
+def handle_file_upload(file_path, session_id):
+    """Handle file upload"""
+    if not file_path:
+        return None, json.dumps({
+            "status": "error",
+            "message": "No file uploaded"
+        }, indent=2), session_id
+    
+    if not session_id:
+        session_id = create_new_session()
+    
+    file_name = os.path.basename(file_path)
+    
+    # Update session
+    session_manager.update_session(session_id, {
+        "current_file": file_path,
+        "state": ConversationState.INITIAL
+    })
+    
+    # Add system message
+    session_manager.add_message(
+        session_id,
+        "system",
+        f"File uploaded: {file_name}"
+    )
+    
+    status = {
+        "status": "success",
+        "message": f"File '{file_name}' uploaded successfully",
+        "file_info": {
+            "name": file_name,
+            "path": file_path,
+            "size_bytes": os.path.getsize(file_path) if os.path.exists(file_path) else 0
+        },
+        "next_action": "💬 Now tell me what you'd like to do with this document"
+    }
+    
+    return file_path, json.dumps(status, indent=2), session_id
+
+
+def chatbot_response_streaming(message: str, history: List, session_id: str, file_path: str = None):
+    """
+    Handle chat messages with streaming updates
+    Uses Bedrock (priority) → Gemini (fallback) for both generation and execution
+    """
+    # Get or create session
+    session = session_manager.get_session(session_id)
+    if not session:
+        session_id = create_new_session()
+        session = session_manager.get_session(session_id)
+    
+    # Update file path if provided
+    if file_path:
+        session_manager.update_session(session_id, {"current_file": file_path})
+        session = session_manager.get_session(session_id)
+    
+    # Add user message to session
+    session_manager.add_message(session_id, "user", message)
+    
+    current_state = session.get("state", ConversationState.INITIAL)
+    
+    # ========================
+    # STATE: INITIAL - Generate Pipeline
+    # ========================
+    if current_state == ConversationState.INITIAL:
+        # Check if file is uploaded
+        if not session.get("current_file"):
+            response = {
+                "status": "error",
+                "message": "Please upload a document first",
+                "action": "📁 Click 'Upload Document' to begin"
+            }
+            response_text = f"```json\n{json.dumps(response, indent=2)}\n```"
+            session_manager.add_message(session_id, "assistant", response_text)
+            yield history + [[message, response_text]]
+            return
+        
+        try:
+            # Generate pipeline using Bedrock → Gemini fallback
+            yield history + [[message, "🤖 Generating pipeline with AI...\n⏳ Trying Bedrock first..."]]
+            
+            pipeline = generate_pipeline(
+                user_input=message,
+                file_path=session.get("current_file"),
+                prefer_bedrock=True
+            )
+            
+            # Save proposed pipeline to session
+            session_manager.update_session(session_id, {
+                "proposed_pipeline": pipeline,
+                "state": ConversationState.PIPELINE_PROPOSED
+            })
+            
+            # Format for display
+            formatted_display = format_pipeline_for_display(pipeline)
+            
+            # Create response with both fancy display and JSON
+            response_text = formatted_display + f"\n\n```json\n{json.dumps(pipeline, indent=2)}\n```"
+            
+            session_manager.add_message(session_id, "assistant", response_text)
+            yield history + [[message, response_text]]
+            return
+            
+        except Exception as e:
+            error_response = {
+                "status": "error",
+                "message": "Failed to generate pipeline",
+                "error": str(e),
+                "action": "Please try rephrasing your request"
+            }
+            response_text = f"```json\n{json.dumps(error_response, indent=2)}\n```"
+            session_manager.add_message(session_id, "assistant", response_text)
+            yield history + [[message, response_text]]
+            return
+    
+    # ========================
+    # STATE: PIPELINE_PROPOSED - Handle Approval/Rejection
+    # ========================
+    elif current_state == ConversationState.PIPELINE_PROPOSED:
+        user_input = message.lower().strip()
+        
+        # APPROVE - Execute the pipeline
+        if "approve" in user_input or "yes" in user_input:
+            session_manager.update_session(session_id, {"state": ConversationState.EXECUTING})
+            
+            plan = session.get("proposed_pipeline", {})
+            
+            # Initial status
+            initial_status = {
+                "status": "executing",
+                "message": "🚀 Starting pipeline execution...",
+                "pipeline": plan.get("pipeline_name", "unknown"),
+                "executor": "Attempting Bedrock LangChain first",
+                "steps": []
+            }
+            accumulated_response = f"```json\n{json.dumps(initial_status, indent=2)}\n```"
+            yield history + [[message, accumulated_response]]
+            
+            steps_completed = []
+            final_payload = None
+            executor_used = "unknown"
+            
+            try:
+                # Execute pipeline with Bedrock → CrewAI fallback
+                for event in execute_pipeline_streaming(
+                    pipeline=plan,
+                    file_path=session.get("current_file"),
+                    session_id=session_id,
+                    prefer_bedrock=True
+                ):
+                    event_type = event.get("type")
+                    
+                    # Info events (fallback notifications, etc.)
+                    if event_type == "info":
+                        info_status = {
+                            "status": "info",
+                            "message": event.get("message"),
+                            "executor": event.get("executor", "unknown")
+                        }
+                        accumulated_response = f"```json\n{json.dumps(info_status, indent=2)}\n```"
+                        yield history + [[message, accumulated_response]]
+                    
+                    # Step updates
+                    elif event_type == "step":
+                        step_info = {
+                            "step": event.get("step", 0),
+                            "tool": event.get("tool", "processing"),
+                            "status": event.get("status", "running"),
+                            "executor": event.get("executor", "unknown")
+                        }
+                        steps_completed.append(step_info)
+                        executor_used = event.get("executor", executor_used)
+                        
+                        progress_status = {
+                            "status": "executing",
+                            "message": f"📍 Step {event.get('step', 0)}: {event.get('tool', 'processing')}...",
+                            "pipeline": plan.get("pipeline_name", ""),
+                            "executor": executor_used,
+                            "steps_completed": steps_completed
+                        }
+                        accumulated_response = f"```json\n{json.dumps(progress_status, indent=2)}\n```"
+                        yield history + [[message, accumulated_response]]
+                    
+                    # Final result
+                    elif event_type == "final":
+                        final_payload = event.get("data")
+                        executor_used = event.get("executor", executor_used)
+                    
+                    # Error
+                    elif event_type == "error":
+                        error_result = {
+                            "status": "failed",
+                            "error": event.get("error"),
+                            "steps_completed": steps_completed,
+                            "executor": event.get("executor", "unknown")
+                        }
+                        final_response = f"```json\n{json.dumps(error_result, indent=2)}\n```"
+                        session_manager.update_session(session_id, {"state": ConversationState.INITIAL})
+                        session_manager.add_message(session_id, "assistant", final_response)
+                        yield history + [[message, final_response]]
+                        return
+                
+                # Process final result
+                if final_payload:
+                    session_manager.update_session(session_id, {
+                        "pipeline_result": final_payload,
+                        "state": ConversationState.INITIAL
+                    })
+                    
+                    # Save execution to MongoDB
+                    session_manager.save_pipeline_execution(
+                        session_id=session_id,
+                        pipeline=plan,
+                        result=final_payload,
+                        file_path=session.get("current_file"),
+                        executor=executor_used
+                    )
+                    
+                    # Format final response
+                    final_display = {
+                        "status": "completed",
+                        "executor": executor_used,
+                        "pipeline": plan.get("pipeline_name"),
+                        "result": final_payload,
+                        "summary": {
+                            "total_steps": len(steps_completed),
+                            "completed_successfully": len([s for s in steps_completed if s.get("status") == "completed"])
+                        }
+                    }
+                    final_response = f"```json\n{json.dumps(final_display, indent=2)}\n```"
+                else:
+                    final_response = f"```json\n{json.dumps({'status': 'completed', 'steps': steps_completed, 'executor': executor_used}, indent=2)}\n```"
+                    session_manager.update_session(session_id, {"state": ConversationState.INITIAL})
+                
+                session_manager.add_message(session_id, "assistant", final_response)
+                yield history + [[message, final_response]]
+                return
+                
+            except Exception as e:
+                error_result = {
+                    "error": str(e),
+                    "status": "failed",
+                    "message": "Pipeline execution failed",
+                    "steps_completed": steps_completed
+                }
+                final_response = f"```json\n{json.dumps(error_result, indent=2)}\n```"
+                session_manager.update_session(session_id, {"state": ConversationState.INITIAL})
+                session_manager.add_message(session_id, "assistant", final_response)
+                yield history + [[message, final_response]]
+                return
+        
+        # REJECT - Cancel the pipeline
+        elif "reject" in user_input or "no" in user_input:
+            session_manager.update_session(session_id, {
+                "state": ConversationState.INITIAL,
+                "proposed_pipeline": None
+            })
+            response_data = {
+                "status": "rejected",
+                "message": "Pipeline rejected by user",
+                "action": "💬 Please provide a new instruction"
+            }
+            response = f"```json\n{json.dumps(response_data, indent=2)}\n```"
+            session_manager.add_message(session_id, "assistant", response)
+            yield history + [[message, response]]
+            return
+        
+        # EDIT - Request modifications
+        elif "edit" in user_input or "modify" in user_input:
+            current_pipeline = session.get("proposed_pipeline", {})
+            edit_help = {
+                "status": "edit_mode",
+                "message": "To modify the plan, describe your changes",
+                "current_plan": current_pipeline,
+                "examples": [
+                    "Add summarization at the end",
+                    "Remove table extraction",
+                    "Only process pages 1-3",
+                    "Translate to French instead of Spanish"
+                ],
+                "action": "Describe your changes, or say 'approve' to run as-is"
+            }
+            response = f"```json\n{json.dumps(edit_help, indent=2)}\n```"
+            session_manager.add_message(session_id, "assistant", response)
+            yield history + [[message, response]]
+            return
+        
+        # Try to modify pipeline based on user input
+        else:
+            if len(message.strip()) > 5:
+                try:
+                    original_plan = session.get("proposed_pipeline", {})
+                    edit_context = f"Original: {original_plan.get('pipeline_name')}. User wants: {message}"
+                    
+                    # Generate new pipeline with modification
+                    new_pipeline = generate_pipeline(
+                        user_input=edit_context,
+                        file_path=session.get("current_file"),
+                        prefer_bedrock=True
+                    )
+                    
+                    session_manager.update_session(session_id, {
+                        "proposed_pipeline": new_pipeline,
+                        "state": ConversationState.PIPELINE_PROPOSED
+                    })
+                    
+                    formatted = format_pipeline_for_display(new_pipeline)
+                    response = formatted + f"\n\n```json\n{json.dumps(new_pipeline, indent=2)}\n```"
+                    session_manager.add_message(session_id, "assistant", response)
+                    yield history + [[message, response]]
+                    return
+                    
+                except Exception as e:
+                    error_response = {
+                        "status": "edit_failed",
+                        "error": str(e),
+                        "message": "Could not modify the plan",
+                        "action": "Try 'approve' to run as-is, or 'reject' to start over"
+                    }
+                    response = f"```json\n{json.dumps(error_response, indent=2)}\n```"
+                    session_manager.add_message(session_id, "assistant", response)
+                    yield history + [[message, response]]
+                    return
+            
+            # Default waiting message
+            response_data = {
+                "status": "waiting_for_confirmation",
+                "message": "Please type 'approve', 'reject', or describe changes",
+                "hint": "You can also say 'edit' for modification hints"
+            }
+            response = f"```json\n{json.dumps(response_data, indent=2)}\n```"
+            session_manager.add_message(session_id, "assistant", response)
+            yield history + [[message, response]]
+            return
+    
+    # Default fallback
+    response = json.dumps({"status": "ready", "message": "Ready for your next instruction"}, indent=2)
+    session_manager.add_message(session_id, "assistant", response)
+    yield history + [[message, response]]
+
+
+# ========================
+# GRADIO UI
+# ========================
+
+with gr.Blocks(
+    title="MasterLLM v2.0 - AI Pipeline Orchestrator",
+    css="""
+    .gradio-container {
+        max-width: 1400px !important;
+    }
+    """
+) as demo:
+    gr.Markdown("""
+    # 🤖 MasterLLM v2.0 - AI Pipeline Orchestrator
+    
+    **🏆 Bedrock Priority** with Gemini Fallback | **💾 MongoDB Sessions** | **📡 Complete REST API**
+    
+    Upload a document, describe what you want, and watch AI orchestrate the perfect pipeline!
+    """)
+    
+    # State management
+    session_id_state = gr.State(value=create_new_session())
+    file_state = gr.State(value=None)
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            # Chat interface
+            chatbot = gr.Chatbot(
+                height=650,
+                show_label=False,
+                bubble_full_width=False,
+                show_copy_button=True,
+                avatar_images=(None, "🤖"),
+            )
+            
+            # Text input
+            msg = gr.Textbox(
+                placeholder="💬 Type your instruction... (e.g., 'extract text from pages 1-5 and summarize')",
+                show_label=False,
+                lines=2,
+                max_lines=4,
+                container=False,
+            )
+            
+            with gr.Row():
+                submit_btn = gr.Button("🚀 Send", variant="primary", scale=2)
+                clear_btn = gr.Button("🗑️ Clear Chat", scale=1)
+        
+        with gr.Column(scale=1):
+            # File upload section
+            gr.Markdown("### 📁 Upload Document")
+            file_upload = gr.File(
+                label="PDF or Image",
+                file_types=[".pdf", ".png", ".jpg", ".jpeg", ".gif", ".bmp"],
+                type="filepath",
+            )
+            
+            upload_status = gr.Textbox(
+                label="📊 Upload Status",
+                interactive=False,
+                lines=10,
+                max_lines=15,
+            )
+            
+            # Session info
+            gr.Markdown("### 🔗 Session Info")
+            session_display = gr.Textbox(
+                label="Session ID",
+                interactive=False,
+                value=lambda: session_id_state.value[:8] + "...",
+            )
+            
+            # Examples
+            gr.Markdown("### 💡 Example Pipelines")
+            gr.Examples(
+                examples=[
+                    "extract text from pages 1-5",
+                    "extract text and summarize",
+                    "extract text, tables, and translate to Spanish",
+                    "get tables from pages 2-4 and summarize",
+                    "text-classify-ner from entire document",
+                    "describe images and summarize findings",
+                    "extract text, detect signatures and stamps",
+                ],
+                inputs=msg,
+            )
+            
+            # System info
+            gr.Markdown("""
+            ### ℹ️ System Features
+            - ✅ **Bedrock** (Claude 3.5 Sonnet) priority
+            - ✅ **Gemini** (gemini-2.0-flash) fallback
+            - ✅ **MongoDB** session persistence
+            - ✅ **Streaming** real-time updates
+            - ✅ **Component-level** JSON output
+            - ✅ **REST API** for integration
+            
+            ### 📊 Pipeline Flow:
+            1. **Upload** your document
+            2. **Describe** what you want
+            3. **Review** AI-generated pipeline
+            4. **Approve** to execute
+            5. **Watch** streaming updates
+            6. **Get** complete JSON results
+            """)
+    
+    # Event handlers
+    file_upload.upload(
+        fn=handle_file_upload,
+        inputs=[file_upload, session_id_state],
+        outputs=[file_state, upload_status, session_id_state],
+    )
+    
+    msg.submit(
+        fn=chatbot_response_streaming,
+        inputs=[msg, chatbot, session_id_state, file_state],
+        outputs=[chatbot],
+    ).then(
+        lambda: "",
+        outputs=msg,
+    )
+    
+    submit_btn.click(
+        fn=chatbot_response_streaming,
+        inputs=[msg, chatbot, session_id_state, file_state],
+        outputs=[chatbot],
+    ).then(
+        lambda: "",
+        outputs=msg,
+    )
+    
+    clear_btn.click(
+        fn=lambda: ([], create_new_session(), None, None, "", ""),
+        outputs=[chatbot, session_id_state, file_state, file_upload, msg, upload_status],
+    )
+
+# Mount Gradio on FastAPI
+app = gr.mount_gradio_app(app, demo, path="/")
+
+
+# ========================
+# LAUNCH
+# ========================
+
+if __name__ == "__main__":
+    import uvicorn
+    port = int(os.getenv("PORT", 7860))
+    print(f"""
+    ╔════════════════════════════════════════════════════════════╗
+    ║                                                            ║
+    ║           🚀 MasterLLM v2.0 Starting...                   ║
+    ║                                                            ║
+    ║  🌐 Gradio UI:  http://localhost:{port}                      ║
+    ║  📡 REST API:   http://localhost:{port}/api/v1                ║
+    ║  📚 API Docs:   http://localhost:{port}/docs                  ║
+    ║                                                            ║
+    ║  🏆 Bedrock:    Priority (Claude 3.5 Sonnet)             ║
+    ║  🔄 Gemini:     Fallback (gemini-2.0-flash)              ║
+    ║  💾 MongoDB:    Session management                         ║
+    ║                                                            ║
+    ╚════════════════════════════════════════════════════════════╝
+    """)
+    
     uvicorn.run(app, host="0.0.0.0", port=port)