Implement comprehensive enhancements: streaming improvements, conversation history, performance dashboard, settings UI, and loading states

app.py

@@ -3,6 +3,8 @@ import asyncio
 import threading
 import queue
 import os
+import time
+import json
 from datetime import datetime
 from modules.input_handler import validate_input
 from modules.retriever import perform_search
@@ -19,14 +21,13 @@ server_monitor = ServerMonitor()
 # Cat-themed greeting function
 def get_cat_greeting():
     """Generate a cat-themed greeting to test if the AI is operational"""
-    greeting_prompt = (
+    return (
         "Hello there! I'm a sophisticated AI research assistant, but right now I'm just a random cat preparing to make biscuits "
         "(that's cat slang for getting ready to do something awesome!). Today is " + datetime.now().strftime("%A, %B %d, %Y") + ". "
         "I'm purring with excitement to help you with your research questions! "
         "Meow... what delicious knowledge shall we hunt down today? "
         "Please ask me anything, and I'll pounce on the best information for you!"
     )
-    return greeting_prompt
 
 # Startup check function
 async def perform_startup_check():
@@ -121,23 +122,33 @@ def run_startup_check():
     wrapper = StartupCheckWrapper(coroutine)
     return wrapper
 
-# Main research assistant function
-async def research_assistant(query):
+# Enhanced streaming with markdown support
+async def research_assistant(query, history):
     log_request("Research started", query=query)
 
+    # Add typing indicator
+    history.append((query, "🔄 Searching for information..."))
+    yield history
+
     cached = get_cached_result(query)
     if cached:
         log_request("Cache hit", query=query)
-        yield cached
+        history[-1] = (query, cached)
+        yield history
         return
 
     try:
         validated_query = validate_input(query)
     except ValueError as e:
-        yield f"⚠️ Input Error: {str(e)}"
+        error_msg = f"⚠️ Input Error: {str(e)}"
+        history[-1] = (query, error_msg)
+        yield history
         return
 
     # Run context enhancement and search in parallel
+    history[-1] = (query, "🔍 Gathering context...")
+    yield history
+    
     weather_task = asyncio.create_task(add_weather_context())
     space_weather_task = asyncio.create_task(add_space_weather_context())
     search_task = asyncio.create_task(asyncio.to_thread(perform_search, validated_query))
@@ -148,7 +159,9 @@ async def research_assistant(query):
 
     # Handle search errors
     if isinstance(search_results, list) and len(search_results) > 0 and "error" in search_results[0]:
-        yield f"🔍 Search Error: {search_results[0]['error']}"
+        error_msg = f"🔍 Search Error: {search_results[0]['error']}"
+        history[-1] = (query, error_msg)
+        yield history
         return
 
     # Format search content for LLM
@@ -165,7 +178,7 @@ async def research_assistant(query):
     server_status = server_monitor.check_server_status()
     if not server_status["available"]:
         wait_time = server_status["estimated_wait"]
-        yield (
+        response = (
             f"⏳ **Server Initializing** ⏳\n\n"
             f"The AI model server is currently starting up. This happens automatically after periods of inactivity.\n\n"
             f"**Estimated wait time: {wait_time} minutes**\n\n"
@@ -175,15 +188,39 @@ async def research_assistant(query):
             f"- Check back shortly - the server will be ready soon!\n\n"
             f"*Technical Details: {server_status['message']}*"
         )
+        history[-1] = (query, response)
+        yield history
         return
 
     try:
+        history[-1] = (query, "🧠 Analyzing information...")
+        yield history
+        
         stream = analyze_with_model(enriched_input)
         full_response = ""
 
+        # Buffer for smoother streaming
+        buffer = ""
+        buffer_threshold = 20  # Characters before yielding
+
         for chunk in stream:
-            full_response += chunk
-            yield full_response
+            buffer += chunk
+            
+            # Yield when buffer is large enough or we have a complete line
+            if len(buffer) > buffer_threshold or '\n' in buffer:
+                full_response += buffer
+                history[-1] = (query, full_response)
+                yield history
+                buffer = ""
+            
+            # Small delay for smoother streaming
+            await asyncio.sleep(0.01)
+
+        # Flush remaining buffer
+        if buffer:
+            full_response += buffer
+            history[-1] = (query, full_response)
+            yield history
 
         citations = generate_citations(search_results)
         citation_text = format_citations(citations)
@@ -193,9 +230,14 @@ async def research_assistant(query):
         server_monitor.report_success()
         log_request("Research completed", result_length=len(full_output))
 
+        history[-1] = (query, full_output)
+        yield history
+
     except Exception as e:
         server_monitor.report_failure()
-        yield f"🤖 **Unexpected Error** 🤖\n\nAn unexpected error occurred:\n\n{str(e)}"
+        error_response = f"🤖 **Unexpected Error** 🤖\n\nAn unexpected error occurred:\n\n{str(e)}"
+        history[-1] = (query, error_response)
+        yield history
 
 # Thread-safe wrapper for async generator
 class AsyncGeneratorWrapper:
@@ -240,62 +282,95 @@ class AsyncGeneratorWrapper:
             raise StopIteration
         return item
 
-def research_assistant_wrapper(query):
-    async_gen = research_assistant(query)
+def research_assistant_wrapper(query, history):
+    async_gen = research_assistant(query, history)
     wrapper = AsyncGeneratorWrapper(async_gen)
     return wrapper
 
+# Performance dashboard data
+def get_performance_stats():
+    """Get performance statistics from Redis"""
+    try:
+        stats = server_monitor.get_system_stats()
+        if "error" in stats:
+            return {"status": "error", "message": stats["error"]}
+        
+        # Add more detailed stats
+        stats["current_time"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        stats["uptime"] = "Calculating..."
+        return stats
+    except Exception as e:
+        return {"status": "error", "message": str(e)}
+
 # Global variable to store startup check result
 startup_check_result = None
 
-# Gradio Interface for Streaming
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="amber", secondary_hue="orange"), title="AI Research Assistant") as demo:
+# Gradio Interface with all enhancements
+with gr.Blocks(
+    theme=gr.themes.Soft(primary_hue="amber", secondary_hue="orange"), 
+    title="AI Research Assistant"
+) as demo:
+    # State management
+    chat_history = gr.State([])
+    
     gr.Markdown("# 🧠 AI Research Assistant")
     gr.Markdown("This advanced AI assistant combines web search with contextual awareness to answer complex questions. "
                 "It incorporates current weather and space weather data for richer context.")
     
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("## System Status")
-            status_display = gr.Markdown("🔄 Checking system status...")
-            check_btn = gr.Button("🔁 Refresh Status")
-            
-            gr.Markdown("## How to Use")
-            gr.Markdown("""
-            1. Enter a research question in the input box
-            2. Click Submit or press Enter
-            3. Watch as the response streams in real-time
-            4. Review sources at the end of each response
-            
-            ## Features
-            - 🔍 Web search integration
-            - 🌤️ Weather context
-            - 🌌 Space weather context
-            - 📚 Real-time citations
-            - ⚡ Streaming output
-            """)
-            
-        with gr.Column(scale=2):
-            chatbot = gr.Chatbot(height=500, label="Research Conversation", latex_delimiters=[{"left": "$$", "right": "$$", "display": True}])
-            msg = gr.Textbox(
-                label="Research Question",
-                placeholder="Ask a complex research question...",
-                lines=3
-            )
-            submit_btn = gr.Button("Submit Research Query")
-            clear_btn = gr.Button("Clear Conversation")
-            
-            examples = gr.Examples(
-                examples=[
-                    "What are the latest developments in quantum computing?",
-                    "How does climate change affect ocean currents?",
-                    "Explain the significance of the James Webb Space Telescope findings",
-                    "What are the economic implications of renewable energy adoption?",
-                    "How do solar flares affect satellite communications?"
-                ],
-                inputs=msg,
-                label="Example Questions"
-            )
+    with gr.Tabs():
+        with gr.TabItem("💬 Chat"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    gr.Markdown("## System Status")
+                    status_display = gr.Markdown("🔄 Checking system status...")
+                    check_btn = gr.Button("🔁 Refresh Status")
+                    
+                    gr.Markdown("## How to Use")
+                    gr.Markdown("""
+                    1. Enter a research question in the input box
+                    2. Click Submit or press Enter
+                    3. Watch as the response streams in real-time
+                    4. Review sources at the end of each response
+                    
+                    ## Features
+                    - 🔍 Web search integration
+                    - 🌤️ Weather context
+                    - 🌌 Space weather context
+                    - 📚 Real-time citations
+                    - ⚡ Streaming output
+                    """)
+                    
+                with gr.Column(scale=2):
+                    chatbot = gr.Chatbot(
+                        height=500, 
+                        label="Research Conversation", 
+                        latex_delimiters=[{"left": "$$", "right": "$$", "display": True}],
+                        bubble_full_width=False
+                    )
+                    msg = gr.Textbox(
+                        label="Research Question",
+                        placeholder="Ask a complex research question...",
+                        lines=3
+                    )
+                    with gr.Row():
+                        submit_btn = gr.Button("Submit Research Query", variant="primary")
+                        clear_btn = gr.Button("Clear Conversation")
+                    
+                    examples = gr.Examples(
+                        examples=[
+                            "What are the latest developments in quantum computing?",
+                            "How does climate change affect ocean currents?",
+                            "Explain the significance of the James Webb Space Telescope findings",
+                            "What are the economic implications of renewable energy adoption?",
+                            "How do solar flares affect satellite communications?"
+                        ],
+                        inputs=msg,
+                        label="Example Questions"
+                    )
+        
+        with gr.TabItem("📊 Performance"):
+            perf_refresh_btn = gr.Button("🔄 Refresh Stats")
+            perf_display = gr.JSON(label="System Statistics")
     
     def update_status():
         """Update the system status display"""
@@ -309,7 +384,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="amber", secondary_hue="orange")
         if result["status"] == "operational":
             cat_greeting = get_cat_greeting()
             status_md = f"""
-{result["message"]}
+✅ **Server is operational and ready to assist!**
 
 🐾 **Cat Greeting:** 
 *{cat_greeting}*
@@ -318,17 +393,17 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="amber", secondary_hue="orange")
 """
         elif result["status"] == "initializing":
             status_md = f"""
-{result["message"]}
+⏳ **Server is currently initializing (503 error detected)**
 
 ⏳ **Estimated wait time:** 5 minutes
 
 While you wait, why not prepare some treats? I'll be ready to hunt for knowledge soon!
 """
         elif result["status"] == "checking":
-            status_md = result["message"]
+            status_md = "🔄 Performing startup checks..."
         else:
             status_md = f"""
-{result["message"]}
+❌ **Server check failed**
 
 📝 **Details:** {result["details"]}
 """
@@ -341,28 +416,39 @@ While you wait, why not prepare some treats? I'll be ready to hunt for knowledge
         startup_check_result = run_startup_check()
         return update_status()
     
-    def respond(message, chat_history):
-        # Add user message to chat history
-        chat_history.append((message, ""))
-        yield chat_history, update_status()
-        
+    def respond(message, history):
         # Get streaming response
-        full_response = ""
-        for partial_response in research_assistant_wrapper(message):
-            full_response = partial_response
-            # Update the last message in chat history with the partial response
-            chat_history[-1] = (message, full_response)
-            yield chat_history, update_status()
+        for updated_history in research_assistant_wrapper(message, history):
+            yield updated_history, update_status()
+    
+    def clear_conversation():
+        return [], []
+    
+    def update_performance_stats():
+        stats = get_performance_stats()
+        return stats
     
     # Set initial status on load
     demo.load(update_status, outputs=status_display)
+    demo.load(update_performance_stats, outputs=perf_display)
     
     # Button interactions
     check_btn.click(refresh_status, outputs=status_display)
-    submit_btn.click(respond, [msg, chatbot], [chatbot, status_display])
-    msg.submit(respond, [msg, chatbot], [chatbot, status_display])
+    submit_btn.click(
+        respond, 
+        [msg, chat_history], 
+        [chatbot, status_display]
+    )
+    msg.submit(
+        respond, 
+        [msg, chat_history], 
+        [chatbot, status_display]
+    )
+    
+    clear_btn.click(clear_conversation, outputs=[chat_history, chatbot])
     
-    clear_btn.click(lambda: (None, update_status()), outputs=[chatbot, status_display])
+    # Performance dashboard
+    perf_refresh_btn.click(update_performance_stats, outputs=perf_display)
 
 if __name__ == "__main__":
     demo.launch()