Update app.py

app.py

@@ -7,17 +7,17 @@ from sentence_transformers import SentenceTransformer, CrossEncoder
 from research_agent.config import AgentConfig
 from research_agent.agent import get_clarifying_questions, research_and_plan, write_report_stream
 
-# --- Improved CSS for a professional chatbot look ---
+# --- CSS for a professional chatbot look ---
 CSS = """
-body, .gradio-container { font-family: 'Inter', sans-serif; background-color: #F1F5F9; }
-.gradio-container { max-width: 800px !important; margin: auto !important; }
+body, .gradio-container { font-family: 'Inter', sans-serif; }
+.gradio-container { max-width: 800px !important; margin: auto !important; padding-top: 20px !important;}
 h1 { text-align: center; font-weight: 700; font-size: 2.5em; color: #1E293B; }
 .sub-header { text-align: center; color: #475569; margin-bottom: 20px; font-size: 1.1em; }
 .accordion { border: none !important; box-shadow: 0 1px 3px 0 rgba(0, 0, 0, 0.1), 0 1px 2px -1px rgba(0, 0, 0, 0.1) !important; }
-.chat-window { min-height: 450px; }
-.message { box-shadow: 0 1px 3px 0 rgba(0, 0, 0, 0.1), 0 1px 2px -1px rgba(0, 0, 0, 0.1) !important; }
-.message.user { background: #2563EB !important; color: white; border-bottom-right-radius: 2px !important; }
-.message.bot { background: #FFFFFF !important; color: #334155; border: 1px solid #E2E8F0; border-bottom-left-radius: 2px !important; }
+#chatbot { box-shadow: 0 1px 3px 0 rgba(0, 0, 0, 0.1), 0 1px 2px -1px rgba(0, 0, 0, 0.1) !important; }
+.message-bubble-container > .message-bubble { box-shadow: 0 1px 3px 0 rgba(0, 0, 0, 0.05) !important; }
+.message-bubble.user { background: #2563EB !important; color: white; }
+.message-bubble.bot { background: #FFFFFF !important; color: #334155; border: 1px solid #E2E8F0; }
 footer { display: none !important; }
 """
 
@@ -45,7 +45,7 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft(primary_hue="blue", secondary_hue="
     gr.Markdown("<h1>Mini DeepSearch Agent</h1>")
     gr.Markdown("<p class='sub-header'>Your AI partner for in-depth research and analysis.</p>")
 
-    with gr.Accordion("API & Settings", open=True, elem_classes="accordion") as settings:
+    with gr.Accordion("API & Settings", open=True, elem_classes="accordion") as settings_accordion:
         with gr.Row():
             google_api_key_input = gr.Textbox(label="Google API Key", type="password", placeholder="Enter Google AI API Key", scale=2)
             tavily_api_key_input = gr.Textbox(label="Tavily API Key", type="password", placeholder="Enter Tavily Search API Key", scale=2)
@@ -59,79 +59,88 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft(primary_hue="blue", secondary_hue="
         label="Research Agent",
         bubble_full_width=False,
         height=500,
-        avatar_images=(None, "https://www.gradio.app/images/logo.png"),
-        visible=False
+        visible=False,
+        render=False # We will render it manually later
     )
-    chat_input = gr.Textbox(placeholder="What topic would you like to research?", interactive=False, visible=False)
+    # Use the modern 'messages' type for the chatbot state
+    chat_history_state = gr.State([])
     
-    # State Management
+    chat_input = gr.Textbox(placeholder="Enter your research topic...", interactive=False, visible=False, render=False)
+    
+    # Agent state can be: "INITIAL", "CLARIFYING", "GENERATING"
     agent_state = gr.State("INITIAL") 
     initial_topic_state = gr.State("")
 
     def handle_initialization(google_key, tavily_key):
         init_status = initialize_models(google_key, tavily_key)
+        initial_message = {"role": "assistant", "content": "Agent initialized. Please enter your research topic to begin."}
         return {
             initialization_status: gr.update(value=f"**Status:** {init_status}", visible=True),
-            chatbot: gr.update(visible=True, value=[(None, "Agent initialized. Please enter your research topic to begin.")]),
+            chatbot: gr.update(visible=True, value=[initial_message]),
+            chat_history_state: [initial_message],
             chat_input: gr.update(interactive=True, visible=True),
-            settings: gr.update(open=False)
+            settings_accordion: gr.update(open=False)
         }
 
-    def chat_step(user_input, history):
-        current_state = agent_state.value
-        original_topic = initial_topic_state.value
-
-        if current_state == "INITIAL":
+    def chat_step(user_input, history, current_agent_state, topic_state):
+        history.append({"role": "user", "content": user_input})
+        
+        if current_agent_state == "INITIAL":
             # 1. User provides the initial topic
-            agent_state.update("CLARIFYING")
-            initial_topic_state.update(user_input)
-            history.append([user_input, None])
-            yield history, gr.update(interactive=False, placeholder="Thinking...")
+            new_agent_state = "CLARIFYING"
+            new_topic_state = user_input
+            history.append({"role": "assistant", "content": "Thinking..."})
+            yield history, history, new_agent_state, new_topic_state, gr.update(interactive=False)
             
             questions = get_clarifying_questions(planner_model, user_input)
-            history[-1][1] = "I can do that. To give you the best report, could you answer these questions for me?\n\n" + questions
-            yield history, gr.update(interactive=True, placeholder="Provide your answers to the questions above...")
+            history[-1] = {"role": "assistant", "content": "I can do that. To give you the best report, could you answer these questions for me?\n\n" + questions}
+            yield history, history, new_agent_state, new_topic_state, gr.update(interactive=True)
 
-        elif current_state == "CLARIFYING":
+        elif current_agent_state == "CLARIFYING":
             # 2. User provides answers to clarifying questions
-            agent_state.update("GENERATING")
-            history.append([user_input, None])
-            yield history, gr.update(interactive=False, placeholder="Generating full report...")
+            new_agent_state = "GENERATING"
+            new_topic_state = topic_state
+            history.append({"role": "assistant", "content": "Generating full report..."})
+            yield history, history, new_agent_state, new_topic_state, gr.update(interactive=False)
             
             try:
-                plan = research_and_plan(config, planner_model, tavily_client, original_topic, user_input)
+                plan = research_and_plan(config, planner_model, tavily_client, new_topic_state, user_input)
                 report_generator = write_report_stream(config, writer_model, tavily_client, embedding_model, reranker, plan)
                 
                 status_updates = ""
                 final_report_md = ""
                 for update in report_generator:
-                    if isinstance(update, str) and "Report Generation Complete" not in update:
-                         final_report_md = update # Keep track of the full report text
+                    final_report_md = update if "Report Generation Complete" not in update else final_report_md
                     status_updates += update
-                    history[-1][1] = status_updates # Stream thought process
-                    yield history, gr.update(interactive=False)
+                    history[-1] = {"role": "assistant", "content": status_updates}
+                    yield history, history, new_agent_state, new_topic_state, gr.update(interactive=False)
                 
-                history.append([None, final_report_md]) # Post the final report
-                agent_state.update("INITIAL") # Reset for next query
-                initial_topic_state.update("")
-                yield history, gr.update(interactive=True, placeholder="Research complete. What's the next topic?")
+                history.append({"role": "assistant", "content": final_report_md})
+                new_agent_state = "INITIAL"
+                new_topic_state = ""
+                yield history, history, new_agent_state, new_topic_state, gr.update(interactive=True)
             
             except Exception as e:
-                history.append([None, f"An error occurred: {str(e)}"])
-                agent_state.update("INITIAL")
-                initial_topic_state.update("")
-                yield history, gr.update(interactive=True, placeholder="Let's try again. What's the topic?")
+                error_message = f"An error occurred: {str(e)}"
+                history.append({"role": "assistant", "content": error_message})
+                new_agent_state = "INITIAL"
+                new_topic_state = ""
+                yield history, history, new_agent_state, new_topic_state, gr.update(interactive=True)
+
+    # Manually render components to control order
+    chatbot.render()
+    chat_input.render()
 
     init_button.click(
         fn=handle_initialization,
         inputs=[google_api_key_input, tavily_api_key_input],
-        outputs=[initialization_status, chatbot, chat_input, settings]
+        outputs=[initialization_status, chatbot, chat_history_state, chat_input, settings_accordion]
     )
     
     chat_input.submit(
         fn=chat_step,
-        inputs=[chat_input, chatbot],
-        outputs=[chatbot, chat_input]
+        inputs=[chat_input, chat_history_state, agent_state, initial_topic_state],
+        outputs=[chatbot, chat_history_state, agent_state, initial_topic_state, chat_input]
     ).then(
         lambda: "", None, chat_input, queue=False
     )