Update app.py

app.py

@@ -9,13 +9,12 @@ from research_agent.agent import get_clarifying_questions, research_and_plan, wr
 
 # --- CSS for a professional chatbot look ---
 CSS = """
-body, .gradio-container { font-family: 'Inter', sans-serif; }
+body, .gradio-container { font-family: 'Inter', sans-serif; background-color: #F1F5F9; }
 .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; }
-#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; }
+#chatbot { box-shadow: 0 1px 3px 0 rgba(0, 0, 0, 0.1), 0 1px 2px -1px rgba(0, 0, 0, 0.1) !important; min-height: 450px; }
 .message-bubble.user { background: #2563EB !important; color: white; }
 .message-bubble.bot { background: #FFFFFF !important; color: #334155; border: 1px solid #E2E8F0; }
 footer { display: none !important; }
@@ -54,30 +53,24 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft(primary_hue="blue", secondary_hue="
     initialization_status = gr.Markdown(visible=False)
 
     chatbot = gr.Chatbot(
-        [],
         elem_id="chatbot",
-        label="Research Agent",
         bubble_full_width=False,
         height=500,
         visible=False,
-        render=False # We will render it manually later
+        type="messages"  # Use the modern 'messages' format
     )
-    # Use the modern 'messages' type for the chatbot state
-    chat_history_state = gr.State([])
+    chat_input = gr.Textbox(placeholder="Enter your research topic...", interactive=False, visible=False)
     
-    chat_input = gr.Textbox(placeholder="Enter your research topic...", interactive=False, visible=False, render=False)
-    
-    # Agent state can be: "INITIAL", "CLARIFYING", "GENERATING"
+    # State variables
     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."}
+        initial_messages = [{"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=[initial_message]),
-            chat_history_state: [initial_message],
+            chatbot: gr.update(visible=True, value=initial_messages),
             chat_input: gr.update(interactive=True, visible=True),
             settings_accordion: gr.update(open=False)
         }
@@ -86,25 +79,19 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft(primary_hue="blue", secondary_hue="
         history.append({"role": "user", "content": user_input})
         
         if current_agent_state == "INITIAL":
-            # 1. User provides the initial topic
-            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)
+            yield history, "CLARIFYING", user_input, gr.update(interactive=False)
             
             questions = get_clarifying_questions(planner_model, user_input)
             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)
+            yield history, "CLARIFYING", user_input, gr.update(interactive=True)
 
         elif current_agent_state == "CLARIFYING":
-            # 2. User provides answers to clarifying questions
-            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)
+            yield history, "GENERATING", topic_state, gr.update(interactive=False)
             
             try:
-                plan = research_and_plan(config, planner_model, tavily_client, new_topic_state, user_input)
+                plan = research_and_plan(config, planner_model, tavily_client, topic_state, user_input)
                 report_generator = write_report_stream(config, writer_model, tavily_client, embedding_model, reranker, plan)
                 
                 status_updates = ""
@@ -113,34 +100,26 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft(primary_hue="blue", secondary_hue="
                     final_report_md = update if "Report Generation Complete" not in update else final_report_md
                     status_updates += update
                     history[-1] = {"role": "assistant", "content": status_updates}
-                    yield history, history, new_agent_state, new_topic_state, gr.update(interactive=False)
+                    yield history, "GENERATING", topic_state, gr.update(interactive=False)
                 
                 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)
+                yield history, "INITIAL", "", gr.update(interactive=True)
             
             except Exception as e:
                 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()
+                yield history, "INITIAL", "", gr.update(interactive=True)
 
     init_button.click(
         fn=handle_initialization,
         inputs=[google_api_key_input, tavily_api_key_input],
-        outputs=[initialization_status, chatbot, chat_history_state, chat_input, settings_accordion]
+        outputs=[initialization_status, chatbot, chat_input, settings_accordion]
     )
     
     chat_input.submit(
         fn=chat_step,
-        inputs=[chat_input, chat_history_state, agent_state, initial_topic_state],
-        outputs=[chatbot, chat_history_state, agent_state, initial_topic_state, chat_input]
+        inputs=[chat_input, chatbot, agent_state, initial_topic_state],
+        outputs=[chatbot, agent_state, initial_topic_state, chat_input]
     ).then(
         lambda: "", None, chat_input, queue=False
     )