collapsible side bar

app.py

@@ -47,33 +47,87 @@ def initialize_agent():
         return False
 
 
-def chat(message, history, session_id):
+def chat_with_streaming(message, history, session_id):
     """
-    Handle chat interaction
+    Stream response for better UX
     
     Args:
         message: User's message
-        history: Chat history (list of [user_msg, bot_msg] pairs)
+        history: Chat history (not used in streaming)
         session_id: Unique session identifier for conversation memory
+    
+    Yields:
+        Partial responses as they become available
     """
     if initialization_error:
-        return initialization_error
+        yield initialization_error
+        return
     
     if not agent:
-        return "⚠️ Agent not initialized. Please refresh the page."
+        yield "⚠️ Agent not initialized. Please refresh the page."
+        return
     
     if not message.strip():
-        return ""
+        return
     
     try:
-        # Use session_id as thread_id for conversation memory
-        answer = agent.query(message, thread_id=session_id)
-        return answer
+        # Show that we're processing
+        yield "🤔 Analyzing your question..."
+        
+        # Debug: Check if stream_query exists
+        if not hasattr(agent, 'stream_query'):
+            print("⚠️ WARNING: agent.stream_query() not found, falling back to regular query")
+            yield agent.query(message, thread_id=session_id)
+            return
+        
+        # Stream events from agent
+        has_answer = False
+        event_count = 0
+        
+        for event in agent.stream_query(message, thread_id=session_id):
+            event_count += 1
+            print(f"📡 Stream event {event_count}: {list(event.keys()) if event else 'None'}")
+            
+            if not event:
+                continue
+            
+            # Get current node
+            node = list(event.keys())[0]
+            print(f"   Processing node: {node}")
+            
+            # Show progress based on which node is executing
+            if node == "generate_research_brief":
+                yield "📋 Planning research strategy..."
+            elif node == "execute_stage":
+                yield "📊 Retrieving data from surveys..."
+            elif node == "extract_stage_context":
+                yield "🔗 Processing retrieved data..."
+            elif node == "verify_results":
+                yield "✓ Verifying data completeness..."
+            elif node == "synthesize_response":
+                yield "✍️ Synthesizing answer..."
+                # Get final answer
+                state = event[node]
+                final_answer = state.get("final_answer")
+                if final_answer:
+                    print(f"   Got final answer ({len(final_answer)} chars)")
+                    yield final_answer
+                    has_answer = True
+                    return
+        
+        print(f"📡 Stream complete. Total events: {event_count}, Has answer: {has_answer}")
+        
+        # Fallback if streaming didn't provide answer
+        if not has_answer:
+            print("⚠️ No answer from streaming, using regular query")
+            yield agent.query(message, thread_id=session_id)
         
     except Exception as e:
         error_msg = f"❌ Error processing query: {str(e)}"
-        print(f"Error details: {e}")  # Log to console
-        return error_msg
+        print(f"Error details: {e}")
+        import traceback
+        traceback.print_exc()
+        yield error_msg
 
 
 def create_new_session():
@@ -179,44 +233,44 @@ with gr.Blocks(title="Survey Analysis Agent", theme=gr.themes.Soft()) as demo:
                 label="💡 Example Questions"
             )
         
-        # Sidebar with info
+        # Collapsible sidebar with info
         with gr.Column(scale=1):
-            gr.Markdown("## 📋 Available Data")
-            survey_info = gr.Markdown(
-                value=get_available_surveys() if init_success else "Agent not initialized",
-            )
-            
-            refresh_info = gr.Button("🔄 Refresh", size="sm")
-            
-            gr.Markdown("""
-            ## 🎯 What I Can Do
-            
-            **📝 Questionnaires**
-            - Question text & options
-            - Topics and themes
-            - Skip logic & sampling
-            - Question sequencing
-            
-            **📊 Response Data**
-            - Overall percentages
-            - Demographic breakdowns
-            - Cross-tabulations
-            - Time comparisons
-            
-            ## 💡 Tips
+            with gr.Accordion("📋 Available Data", open=False):
+                survey_info = gr.Markdown(
+                    value=get_available_surveys() if init_success else "Agent not initialized",
+                )
+                refresh_info = gr.Button("🔄 Refresh", size="sm")
             
-            - Specify time periods when relevant
-            - Ask follow-up questions for more detail
-            - I maintain conversation context
-            - Request comparisons across time periods
+            with gr.Accordion("🎯 What I Can Do", open=False):
+                gr.Markdown("""
+                **📝 Questionnaires**
+                - Question text & options
+                - Topics and themes
+                - Skip logic & sampling
+                - Question sequencing
+                
+                **📊 Response Data**
+                - Overall percentages
+                - Demographic breakdowns
+                - Cross-tabulations
+                - Time comparisons
+                """)
             
-            ## 🔧 Current Status
+            with gr.Accordion("💡 Tips", open=False):
+                gr.Markdown("""
+                - Specify time periods when relevant
+                - Ask follow-up questions for more detail
+                - I maintain conversation context
+                - Request comparisons across time periods
+                """)
             
-            ✅ Questionnaire data  
-            ✅ Toplines (response %)  
-            ✅ Crosstabs (demographics)  
-            ⏳ SQL queries (coming soon)
-            """)
+            with gr.Accordion("🔧 Current Status", open=False):
+                gr.Markdown("""
+                ✅ Questionnaire data  
+                ✅ Toplines (response %)  
+                ✅ Crosstabs (demographics)  
+                ⏳ SQL queries (coming soon)
+                """)
     
     # Footer
     gr.Markdown("""
@@ -227,20 +281,23 @@ with gr.Blocks(title="Survey Analysis Agent", theme=gr.themes.Soft()) as demo:
     
     # Event handlers
     def respond(message, chat_history, session_id):
-        """Handle message and update chat history"""
+        """Handle message with streaming updates"""
         if not message.strip():
             return chat_history, ""
         
         # Add user message
         chat_history.append({"role": "user", "content": message})
         
-        # Get bot response
-        bot_message = chat(message, chat_history, session_id)
+        # Add placeholder for assistant response
+        chat_history.append({"role": "assistant", "content": ""})
         
-        # Add bot message
-        chat_history.append({"role": "assistant", "content": bot_message})
+        # Stream updates
+        for partial_response in chat_with_streaming(message, chat_history, session_id):
+            chat_history[-1]["content"] = partial_response
+            yield chat_history, ""
         
-        return chat_history, ""
+        # Final return
+        yield chat_history, ""
     
     def clear_chat():
         """Clear chat and create new session"""