Update app.py

app.py

@@ -1,230 +1,265 @@
+# app.py
+
 import gradio as gr
 import os
-from llama_index.core import Settings
-from llama_index.llms.anthropic import Anthropic
-from llama_index.core.llms import ChatMessage, MessageRole
 import traceback
 import asyncio
-from dotenv import load_dotenv  # Make sure to use this
-import uuid
-from models.task_prompt import TaskPrompt  # Import from models directory
+from dotenv import load_dotenv
+from models.task_prompt import TaskPrompt
+import time
+from llama_index.core import Settings as LlamaSettings # Import at top level
+from llama_index.llms.anthropic import Anthropic # Import at top level
+from manager_agent import ManagerAgent # Ensure this path is correct
+import concurrent.futures # For running blocking code in a separate thread
 
 # Load environment variables from .env file
 load_dotenv()
 
-# Import your ManagerAgent and related classes
-from manager_agent import ManagerAgent  # Using the existing manager_agent.py file
-
 # --- Configuration ---
-LLM_MODEL = "claude-sonnet-4-20250514"  # Modèle Claude Opus 4
+LLM_MODEL = "claude-sonnet-4-20250514"
 
-# --- Global variables for update handling ---
-status_update_queue = asyncio.Queue()
+# --- Global variables ---
 current_status = "Ready"
+llm_global = None
+manager_agent_global = None
+# Settings_global is not strictly needed as a global if LlamaSettings is imported directly
+
+# Thread pool executor for running blocking agent tasks
+thread_pool_executor = concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count() or 1)
+
 
 # --- LlamaIndex LLM Initialization ---
-def initialize_llm_and_agent():
-    """Initialize LLM and Manager Agent with proper error handling"""
-    
-    # Check if API key is available
+def initialize_components():
+    global llm_global, manager_agent_global
+
     api_key = os.environ.get("ANTHROPIC_API_KEY")
     if not api_key:
         print("\n" + "="*60)
         print("⚠️ ERROR: ANTHROPIC_API_KEY not found in environment variables!")
-        print("Please set your API key using one of these methods:")
-        print("1. Create a .env file with: ANTHROPIC_API_KEY=your-api-key-here")
-        print("2. Set as environment variable: export ANTHROPIC_API_KEY='your-api-key-here'")
+        print("Please set your API key (e.g., in a .env file).")
         print("="*60 + "\n")
-        return None, None
-    
+        return
+
     try:
-        # Initialize LLM with lower temperature for more deterministic responses
-        llm = Anthropic(
+        llm_global = Anthropic(
             model=LLM_MODEL,
-            temperature=0.2,  # Lower temperature for more focused responses
+            temperature=0.2,
             max_tokens=4096
         )
-        Settings.llm = llm
+        LlamaSettings.llm = llm_global # Use the imported LlamaSettings directly
         print(f"Successfully initialized LlamaIndex with Anthropic model: {LLM_MODEL} (temperature=0.2)")
-        
-        # Initialize the ManagerAgent with update callback
-        agent = ManagerAgent(
-            llm, 
-            max_iterations=30,  # Increase max_iterations to avoid "Reached max iterations" error
-            update_callback=update_status  # Pass the update callback function
+
+        manager_agent_global = ManagerAgent(
+            llm_global,
+            max_iterations=30, # Keep this reasonable for testing
+            update_callback=update_status_callback
         )
         print("✅ ManagerAgent initialized successfully")
-        
-        return llm, agent
-        
+
     except Exception as e:
         print(f"Error initializing Anthropic LLM or ManagerAgent: {e}")
         traceback.print_exc()
-        return None, None
 
-# --- Update callback function ---
-def update_status(message):
-    """Callback function for the ManagerAgent to update status"""
+# --- Update callback function (called by ManagerAgent) ---
+def update_status_callback(message):
     global current_status
+    # This function is called from the ManagerAgent's thread (potentially)
+    # or the ReAct agent's execution context.
+    # It needs to update the global variable, which the Gradio polling thread will pick up.
     current_status = message
-    
-    # Add to queue for the status monitor
-    try:
-        status_update_queue.put_nowait(message)
-    except Exception as e:
-        print(f"Error adding status update to queue: {e}")
+    print(f"✅ UI_STATUS_UPDATE (via callback): {message}") # Differentiate console log
 
-# --- Status monitor function ---
-async def status_monitor():
-    """Monitor for status updates and update the Gradio component"""
+# --- Status retrieval function for Gradio polling ---
+def get_current_status_for_ui():
     global current_status
-    
-    try:
-        # Set initial status
-        yield current_status
-        
-        # Monitor for updates
-        while True:
-            try:
-                # Check for new updates with a timeout
-                message = await asyncio.wait_for(status_update_queue.get(), timeout=0.5)
-                yield message
-            except asyncio.TimeoutError:
-                # No updates, continue waiting
-                await asyncio.sleep(0.1)
-            except Exception as e:
-                print(f"Error in status monitor: {e}")
-                await asyncio.sleep(1)
-    except Exception as e:
-        print(f"Status monitor error: {e}")
-        yield f"Status monitor error: {str(e)}"
-
-# --- Gradio Chat Function ---
-async def respond(message, history):
-    """Fonction de réponse pour Gradio ChatInterface utilisant ManagerAgent"""
-    
-    # Vérifier si le ManagerAgent est initialisé
-    if manager_agent is None or Settings.llm is None:
-        yield "❌ ManagerAgent not initialized. Please check your ANTHROPIC_API_KEY environment variable and ensure all components are properly loaded."
-        return
-    
-    try:
-        print(f"\n🤖 ManagerAgent: Processing user message: '{message[:100]}{'...' if len(message) > 100 else ''}'")
-        
-        # Update status
-        update_status(f"Processing your request: '{message[:50]}{'...' if len(message) > 50 else ''}'")
-        
-        # Créer un TaskPrompt à partir du message utilisateur
-        # Le ManagerAgent va analyser le prompt et décider du workflow approprié
-        task_prompt = TaskPrompt(text=message)
-        
-        print("🔄 ManagerAgent: Starting task execution workflow...")
-        
-        # Get complete response
-        response = manager_agent.run_task(task_prompt)
-        
-        # Simuler un streaming progressif de la réponse pour une meilleure UX
-        words = response.split()
-        partial_response = ""
-        
-        for i, word in enumerate(words):
-            partial_response += word + " "
-            
-            # Yield la réponse progressive toutes les quelques mots pour un effet de streaming
-            if i % 3 == 0 or i == len(words) - 1:  # Toutes les 3 mots ou le dernier mot
-                yield partial_response.strip()
-                # Add a small delay for realistic streaming
-                await asyncio.sleep(0.01)
-        
-        print("✅ ManagerAgent: Task completed successfully")
-        update_status("Ready for your next request")
-            
-    except Exception as e:
-        error_message = f"❌ ManagerAgent Error: {str(e)}"
-        print(f"\n🚨 ManagerAgent Error: {e}")
-        traceback.print_exc()
-        update_status(f"Error: {str(e)}")
-        yield error_message
+    timestamp = time.time()
+    return f"{current_status}<span style='display:none;'>{timestamp}</span>"
 
 # --- Gradio Interface Setup ---
-def create_interface():
-    """Crée l'interface Gradio"""
-    
-    # Vérifier la clé API
+def create_gradio_interface():
     if "ANTHROPIC_API_KEY" not in os.environ:
-        print("\n" + "="*60)
-        print("⚠️ WARNING: ANTHROPIC_API_KEY not found in environment variables!")
-        print("Please set your API key:")
-        print("export ANTHROPIC_API_KEY='your-api-key-here'")
-        print("="*60 + "\n")
-    
-    # Create the chat interface
-    interface = gr.ChatInterface(
-        fn=respond,
-        chatbot=gr.Chatbot(
+        gr.Warning("ANTHROPIC_API_KEY not found in environment variables! ALITA may not function correctly.")
+
+    with gr.Blocks(theme="soft") as demo:
+        gr.Markdown("# ALITA")
+        gr.Markdown("ALITA is a self-learning AI agent that can search for information, analyze data, create tools, and orchestrate complex tasks.")
+
+        chatbot_component = gr.Chatbot(
+            label="Chat",
             height=500,
             show_label=False,
-            container=True,
-            type="messages"
-        ),
-        textbox=gr.Textbox(
-            placeholder="Tapez votre message ici...",
-            container=False,
-            scale=7
-        ),
-        title="ALITA",
-        description=f"ALITA is a self-learning AI agent that can search for information, analyze data, create tools, and orchestrate complex tasks.",
-        examples=[
-            "🔍 Recherche des informations sur l'intelligence artificielle",
-            "📊 Analyse les tendances du marché technologique",
-            "⚡ Crée un script pour automatiser une tâche répétitive",
-            "🌐 Trouve des ressources open source pour machine learning",
-        ],
-        theme="soft"
-    )
-    
-    # Add the status box
-    with interface:
+            # type='messages' # For Gradio 4.x+
+        )
+        gr.Markdown("Gradio version: " + gr.__version__ + " (Chatbot type defaults to 'tuples' for older versions. Consider `type='messages'` for newer Gradio if issues persist with chat display).")
+
+
         with gr.Row():
-            with gr.Column():
-                status_box = gr.Textbox(
-                    label="Agent Status",
-                    value="Ready",
-                    interactive=False
+            message_textbox = gr.Textbox(
+                placeholder="Type your message here...",
+                scale=7,
+                show_label=False,
+                container=False
+            )
+
+        gr.Examples(
+            examples=[
+            "🔍 Search for information on artificial intelligence",
+            "📊 Analyze technology market trends",
+            "⚡ Create a script to automate a repetitive task",
+            "🌐 Find open source resources for machine learning",
+            ],
+            inputs=message_textbox,
+        )
+
+        status_box_component = gr.Textbox(
+            label="Agent Status",
+            value=get_current_status_for_ui(),
+            interactive=False,
+            # elem_id="status_box_alita" # For potential direct JS manipulation if desperate (avoid)
+        )
+
+        def add_user_msg(user_input_text, chat_history_list):
+            if not user_input_text.strip():
+                return gr.update(), chat_history_list
+            # For older Gradio, history is list of [user_msg, bot_msg] tuples
+            chat_history_list.append((user_input_text, None))
+            return gr.update(value=""), chat_history_list
+
+        async def generate_bot_reply(chat_history_list):
+            if not chat_history_list or chat_history_list[-1][0] is None:
+                # This case should ideally not be reached if add_user_msg works correctly
+                yield chat_history_list
+                return
+
+            user_message = chat_history_list[-1][0]
+
+            if manager_agent_global is None or LlamaSettings.llm is None:
+                # This update_status_callback will set current_status
+                # The polling mechanism (continuous_status_updater) should pick it up.
+                update_status_callback("⚠️ Error: Agent or LLM not initialized. Check API key and logs.")
+                # For older Gradio, update the last tuple's second element
+                chat_history_list[-1] = (chat_history_list[-1][0], "❌ Critical Error: ALITA is not properly initialized. Please check server logs and API key.")
+                yield chat_history_list
+                return
+
+            try:
+                print(f"\n🤖 GRADIOLOG: Processing user message: '{user_message[:100]}{'...' if len(user_message) > 100 else ''}'")
+                update_status_callback(f"💬 Processing: '{user_message[:50]}{'...' if len(user_message) > 50 else ''}'")
+                await asyncio.sleep(0.01) # Allow UI to briefly update with "Processing..."
+
+                task_prompt = TaskPrompt(text=user_message)
+
+                update_status_callback("🔄 Analyzing request and determining optimal workflow...")
+                await asyncio.sleep(0.01) # Allow UI to briefly update
+
+                # Run the blocking manager_agent_global.run_task in a separate thread
+                loop = asyncio.get_event_loop()
+                response_text_from_agent = await loop.run_in_executor(
+                    thread_pool_executor,
+                    manager_agent_global.run_task, # The function to run
+                    task_prompt                     # Arguments to the function
                 )
+                # By this point, run_task has completed, and all its internal
+                # calls to update_status_callback (via send_update) should have occurred.
+                # The polling mechanism should have picked up these changes.
+
+                update_status_callback("✨ Generating final response stream...")
+                await asyncio.sleep(0.01)
+                final_bot_response = response_text_from_agent
+
+                words = final_bot_response.split()
+                accumulated_response_stream = ""
+                total_words = len(words)
                 
-                # Use the Interface for status monitoring
-                gr.Interface(
-                    fn=status_monitor,
-                    inputs=None,
-                    outputs=status_box,
-                    live=True,
-                    show_progress=False
-                )
-    
-    return interface
+                # Initialize bot's part of the message in history for older Gradio
+                current_user_message = chat_history_list[-1][0]
+                chat_history_list[-1] = (current_user_message, "")
+
+
+                if not words:
+                    chat_history_list[-1] = (current_user_message, final_bot_response.strip())
+                    yield chat_history_list
+                else:
+                    for i, word in enumerate(words):
+                        accumulated_response_stream += word + " "
+                        # These status updates are for the streaming part,
+                        # agent's internal updates should have already happened.
+                        if total_words > 0: # Avoid division by zero
+                            if i == total_words // 4: update_status_callback("🔄 Streaming response (25%)...")
+                            elif i == total_words // 2: update_status_callback("🔄 Streaming response (50%)...")
+                            elif i == (total_words * 3) // 4: update_status_callback("🔄 Streaming response (75%)...")
 
-# Initialize the components
-llm, manager_agent = initialize_llm_and_agent()
+                        if i % 3 == 0 or i == len(words) - 1:
+                            chat_history_list[-1] = (current_user_message, accumulated_response_stream.strip())
+                            yield chat_history_list
+                            await asyncio.sleep(0.01) # For streaming effect
+
+                # Ensure final complete response is set
+                if chat_history_list[-1][1] != final_bot_response.strip():
+                    chat_history_list[-1] = (current_user_message, final_bot_response.strip())
+                    yield chat_history_list
+                
+                print("✅ GRADIOLOG: Task processing and streaming completed.")
+                update_status_callback("✅ Ready for your next request")
+
+            except Exception as e:
+                error_message_for_ui = f"❌ Gradio/Agent Error: {str(e)}"
+                print(f"\n🚨 GRADIOLOG: Error in generate_bot_reply: {e}")
+                traceback.print_exc()
+                update_status_callback(f"❌ Error: {str(e)[:100]}...")
+                chat_history_list[-1] = (chat_history_list[-1][0], error_message_for_ui)
+                yield chat_history_list
+        
+        message_textbox.submit(
+            add_user_msg,
+            inputs=[message_textbox, chatbot_component],
+            outputs=[message_textbox, chatbot_component],
+            show_progress="hidden", # Gradio 3.x might not have this, can be ignored
+        ).then(
+            generate_bot_reply,
+            inputs=[chatbot_component],
+            outputs=[chatbot_component],
+            api_name=False, # Good practice
+            # show_progress="hidden", # Gradio 3.x might not have this
+        )
+
+        async def continuous_status_updater(update_interval_seconds=0.3): # Slightly faster poll
+            """Continuously yields status updates for the status_box_component."""
+            print("GRADIOLOG: Starting continuous_status_updater loop.")
+            while True:
+                # print(f"POLL: Fetching status: {current_status}") # DEBUG: very verbose
+                yield get_current_status_for_ui()
+                await asyncio.sleep(update_interval_seconds)
+
+        demo.load(continuous_status_updater, inputs=None, outputs=status_box_component)
+        print("GRADIOLOG: Continuous status updater loaded.")
+    return demo
+
+# Initialize LLM and Agent components
+initialize_components()
 
 # --- Launch the Application ---
 if __name__ == "__main__":
+    print(f"Gradio version: {gr.__version__}")
+
     print("🚀 Starting Gradio ALITA Chat Application...")
-    
-    # Créer et lancer l'interface
-    app = create_interface()
-    
+    alita_interface = create_gradio_interface()
+
     try:
-        app.launch(
+        alita_interface.launch(
             share=False,
             server_name="127.0.0.1",
-            server_port=7825,
-            show_error=True
+            server_port=6126,
+            show_error=True,
+            # debug=True # Can be helpful
         )
     except KeyboardInterrupt:
         print("\n👋 Application stopped by user")
     except Exception as e:
         print(f"\n❌ Error launching application: {e}")
         traceback.print_exc()
-    
+    finally:
+        print("Shutting down thread pool executor...")
+        thread_pool_executor.shutdown(wait=True) # Clean up threads
+
     print("✅ Gradio application stopped.")