Update app.py

app.py

@@ -7,20 +7,18 @@ 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
 
-# --- CSS for styling the Gradio app ---
+# --- Improved CSS for a professional chatbot look ---
 CSS = """
-body, .gradio-container { font-family: 'Inter', sans-serif; }
-.gradio-container { max-width: 960px !important; margin: auto !important; }
-h1 { text-align: center; font-size: 2.2em; color: #334155; }
-.sub-header { text-align: center; color: #64748B; margin-bottom: 20px; }
-.gr-button { background-color: #2563EB; color: white; }
-.gr-button:hover { background-color: #1E4ED8; }
-.accordion { border: 1px solid #E5E7EB !important; border-radius: 8px !important; }
-.chat-bubble-container { display: flex; flex-direction: column; gap: 5px; }
-.chat-bubble-message { padding: 10px; border-radius: 10px; font-size: 0.95rem; }
-.user-message { background-color: #EFF6FF; color: #1E3A8A; align-self: flex-end; border-bottom-right-radius: 2px; }
-.bot-message { background-color: #F8FAFC; color: #334155; align-self: flex-start; border: 1px solid #E2E8F0; border-bottom-left-radius: 2px; }
-.thinking { color: #64748B; font-style: italic; text-align: center; padding: 10px; }
+body, .gradio-container { font-family: 'Inter', sans-serif; background-color: #F1F5F9; }
+.gradio-container { max-width: 800px !important; margin: auto !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; }
+footer { display: none !important; }
 """
 
 # --- Model Initialization ---
@@ -38,97 +36,104 @@ def initialize_models(google_key, tavily_key):
         planner_model = genai.GenerativeModel(config.WRITER_MODEL)
         embedding_model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
         reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', device='cpu')
-        print("Models initialized successfully!")
+        return "Models initialized successfully!"
     except Exception as e:
         raise gr.Error(f"Failed to initialize models. Error: {str(e)}")
 
 # --- Gradio Application Logic ---
-with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as app:
-    gr.Markdown("# Mini DeepSearch Agent")
-    gr.Markdown("This agent performs in-depth research using a multi-step AI process.", elem_classes="sub-header")
+with gr.Blocks(css=CSS, theme=gr.themes.Soft(primary_hue="blue", secondary_hue="slate")) as app:
+    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.Row():
-        with gr.Column(scale=1):
-            with gr.Accordion("API & Settings", open=True, elem_classes="accordion") as settings:
-                google_api_key_input = gr.Textbox(label="Google API Key", type="password", placeholder="Enter Google AI API Key")
-                tavily_api_key_input = gr.Textbox(label="Tavily API Key", type="password", placeholder="Enter Tavily Search API Key")
-                init_button = gr.Button("Initialize Agent")
-        
-        with gr.Column(scale=3):
-            chatbot = gr.Chatbot(
-                [],
-                label="Research Agent Conversation",
-                bubble_full_width=False,
-                avatar_images=(None, "https://www.gradio.app/images/logo.png"),
-                elem_id="chatbot"
-            )
-            chat_input = gr.Textbox(placeholder="What would you like to research?", interactive=False)
+    with gr.Accordion("API & Settings", open=True, elem_classes="accordion") as settings:
+        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)
+            init_button = gr.Button("Initialize Agent", scale=1, variant="primary")
+    
+    initialization_status = gr.Markdown(visible=False)
 
-    # --- State Management ---
-    # agent_state can be: "INITIAL", "CLARIFYING", "GENERATING", "DONE"
+    chatbot = gr.Chatbot(
+        [],
+        elem_id="chatbot",
+        label="Research Agent",
+        bubble_full_width=False,
+        height=500,
+        avatar_images=(None, "https://www.gradio.app/images/logo.png"),
+        visible=False
+    )
+    chat_input = gr.Textbox(placeholder="What topic would you like to research?", interactive=False, visible=False)
+    
+    # State Management
     agent_state = gr.State("INITIAL") 
     initial_topic_state = gr.State("")
 
-    # --- Event Handlers ---
     def handle_initialization(google_key, tavily_key):
-        initialize_models(google_key, tavily_key)
+        init_status = initialize_models(google_key, tavily_key)
         return {
-            chat_input: gr.update(interactive=True, placeholder="Models Initialized! What would you like to research?"),
-            init_button: gr.update(value="Agent Ready!", interactive=False)
+            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.")]),
+            chat_input: gr.update(interactive=True, visible=True),
+            settings: gr.update(open=False)
         }
 
-    def chat_step(user_input, history, state, original_topic):
-        history = history or []
-        
-        if state == "INITIAL":
+    def chat_step(user_input, history):
+        current_state = agent_state.value
+        original_topic = initial_topic_state.value
+
+        if current_state == "INITIAL":
+            # 1. User provides the initial topic
+            agent_state.update("CLARIFYING")
+            initial_topic_state.update(user_input)
             history.append([user_input, None])
-            # Show a thinking message
-            yield {chatbot: history, agent_state: "CLARIFYING", initial_topic_state: user_input, chat_input: gr.update(interactive=False, placeholder="Thinking...")}
+            yield history, gr.update(interactive=False, placeholder="Thinking...")
             
-            # Get clarifying questions
             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 {chatbot: history, chat_input: gr.update(interactive=True, placeholder="Provide your answers to the questions above...")}
+            yield history, gr.update(interactive=True, placeholder="Provide your answers to the questions above...")
 
-        elif state == "CLARIFYING":
+        elif current_state == "CLARIFYING":
+            # 2. User provides answers to clarifying questions
+            agent_state.update("GENERATING")
             history.append([user_input, None])
-            yield {chatbot: history, agent_state: "GENERATING", chat_input: gr.update(interactive=False, placeholder="Generating full report...")}
-            
-            # Plan the research
-            status_updates = "### Agent Status\n"
-            plan = research_and_plan(config, planner_model, tavily_client, original_topic, user_input)
-            
-            status_updates += f"**Research Plan:**\n- **Topic:** {plan['detailed_topic']}\n- **Sections:** {[s.title for s in plan['sections']]}\n\n---\n"
-            history[-1][1] = status_updates
-            yield {chatbot: history}
+            yield history, gr.update(interactive=False, placeholder="Generating full report...")
             
-            # Generate the report, streaming updates
-            report_generator = write_report_stream(config, writer_model, tavily_client, embedding_model, reranker, plan)
-            final_report_md = ""
-            for update in report_generator:
-                if isinstance(update, str):
-                    final_report_md = update
-                status_updates += update
-                history[-1][1] = status_updates
-                yield {chatbot: history}
+            try:
+                plan = research_and_plan(config, planner_model, tavily_client, original_topic, 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
+                    status_updates += update
+                    history[-1][1] = status_updates # Stream thought process
+                    yield history, 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?")
             
-            # Append the final report as a new message
-            history.append([None, final_report_md])
-            yield {chatbot: history, agent_state: "DONE", chat_input: gr.update(interactive=True, placeholder="Research complete. What's next?")}
+            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?")
 
     init_button.click(
         fn=handle_initialization,
         inputs=[google_api_key_input, tavily_api_key_input],
-        outputs=[chat_input, init_button]
+        outputs=[initialization_status, chatbot, chat_input, settings]
     )
     
     chat_input.submit(
         fn=chat_step,
-        inputs=[chat_input, chatbot, agent_state, initial_topic_state],
-        outputs=[chatbot, agent_state, initial_topic_state, chat_input]
+        inputs=[chat_input, chatbot],
+        outputs=[chatbot, chat_input]
     ).then(
-        fn=lambda: "",
-        outputs=[chat_input]
+        lambda: "", None, chat_input, queue=False
     )
 
 app.launch(debug=True)