Update app.py

app.py

@@ -1,55 +1,95 @@
 import gradio as gr
+import asyncio
+from functools import lru_cache
 from search import search_google
 from llm import generate_answer
 from memory import ConversationMemory
+from utils import async_timeout
 
 # Initialize conversation memory
 memory = ConversationMemory()
 
-def ask_agent(question):
+@async_timeout(30)  # Timeout after 30 seconds
+async def ask_agent(question, progress=gr.Progress()):
+    # Track progress
+    progress(0.1, desc="📡 Searching the web...")
+    
     # Retrieve conversation context
     context = memory.get_context()
     
-    # Search for information
-    search_results = search_google(question, num_results=5)
+    # Search for information (with caching)
+    search_results = await cached_search_async(question, num_results=5)
     
     if not search_results:
         return "I couldn't find any relevant information about that. Could you try rephrasing your question?"
     
+    # Track progress
+    progress(0.4, desc="🧠 Processing information...")
+    
     # Generate human-like response
-    answer = generate_answer(
+    answer = await generate_answer(
         question=question,
         context=context,
         search_results=search_results
     )
     
     # Update conversation history
-    memory.add_exchange(question, answer)
+    memory.add_exchange(question, answer['response'])
     
     # Format response with sources
     formatted_response = f"""
-🤖 **Assistant**: {answer['response']}
-    
-🔍 **Sources I used**:
+🤖 **Assistant**: {answer['response']}\n
+🔍 **Sources I used**:\n
 """
     for source in answer['sources']:
         formatted_response += f"- [{source['title']}]({source['url']})\n"
     
+    # Track progress
+    progress(1.0, desc="✅ Response ready")
     return formatted_response
 
-# Gradio chat interface
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🧠 **AI Research Assistant**")
-    chatbot = gr.Chatbot(height=500)
-    msg = gr.Textbox(label="Your Question")
+# Cached async search
+@lru_cache(maxsize=100)
+def cached_search_async(query, num_results=5):
+    return search_google(query, num_results)
+
+# Gradio chat interface with progress tracking
+with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 800px; margin: auto;}") as demo:
+    gr.Markdown("""
+    <div style="text-align: center;">
+        <h1>🧠 AI Research Assistant</h1>
+        <p>I can help you find information on any topic!</p>
+    </div>
+    """)
+    
+    chatbot = gr.Chatbot(height=400, bubble_full_width=False)
+    msg = gr.Textbox(label="Your Question", placeholder="Ask me anything...")
     clear = gr.Button("Clear History")
+    status = gr.Textbox("", label="Status", interactive=False)
     
-    def respond(message, chat_history):
-        bot_message = ask_agent(message)
+    async def respond(message, chat_history):
+        # Create progress tracker
+        tracker = []
+        
+        # Wrap in try/except for better error handling
+        try:
+            bot_message = await ask_agent(
+                message, 
+                progress=lambda p, d, t=tracker: tracker.append((p, d))
+            )
+            
+            # Update status
+            if tracker:
+                status.value = tracker[-1][1]
+        except Exception as e:
+            bot_message = f"⚠️ Sorry, I encountered an error: {str(e)[:100]}"
+            status.value = "Error occurred"
+        
         chat_history.append((message, bot_message))
         return "", chat_history
     
     msg.submit(respond, [msg, chatbot], [msg, chatbot])
-    clear.click(lambda: None, None, chatbot, queue=False)
+    clear.click(lambda: (memory.clear(), None), None, chatbot, queue=False)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.queue(concurrency_count=4).launch()