Update app.py

app.py

@@ -7,17 +7,24 @@ 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 a professional chatbot look ---
+# --- CSS for a professional, ChatGPT-inspired look ---
 CSS = """
-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; min-height: 450px; }
-.message-bubble.user { background: #2563EB !important; color: white; }
-.message-bubble.bot { background: #FFFFFF !important; color: #334155; border: 1px solid #E2E8F0; }
+body, .gradio-container { font-family: 'Inter', sans-serif; background-color: #343541; color: #ECECEC; }
+.gradio-container { max-width: 800px !important; margin: auto !important; padding-top: 2rem !important;}
+h1 { text-align: center; font-weight: 700; font-size: 2.5em; color: white; }
+.sub-header { text-align: center; color: #C5C5D2; margin-bottom: 2rem; font-size: 1.1em; }
+.accordion { background-color: #40414F; border: 1px solid #565869 !important; border-radius: 8px !important; }
+.accordion .gr-button { background-color: #4B4C5A; color: white; }
+#chatbot { box-shadow: none !important; border: none !important; background-color: transparent !important; }
+.message-bubble { background: #40414F !important; border: 1px solid #565869 !important; color: #ECECEC !important;}
+.message-bubble.user { background: #343541 !important; border: none !important; }
 footer { display: none !important; }
+.gr-box.gradio-container { padding: 0 !important; }
+.gr-form { background-color: transparent !important; border: none !important; box-shadow: none !important; }
+.gradio-container .gr-form .gr-button { display: none; } /* Hide the default submit button */
+#chat-input-container { position: relative; }
+#chat-input-container textarea { background-color: #40414F; color: white; border: 1px solid #565869 !important; }
+#submit-button { position: absolute; right: 10px; top: 50%; transform: translateY(-50%); background: #2563EB; color: white; border-radius: 4px; padding: 4px 8px; }
 """
 
 # --- Model Initialization ---
@@ -27,7 +34,7 @@ writer_model, planner_model, embedding_model, reranker, tavily_client = None, No
 def initialize_models(google_key, tavily_key):
     global writer_model, planner_model, embedding_model, reranker, tavily_client
     if not google_key or not tavily_key:
-        raise gr.Error("API keys are required. Please provide both Google and Tavily API keys.")
+        raise gr.Error("API keys are required.")
     try:
         genai.configure(api_key=google_key)
         tavily_client = TavilyClient(api_key=tavily_key)
@@ -35,93 +42,96 @@ 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')
-        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(primary_hue="blue", secondary_hue="slate")) as app:
+with gr.Blocks(css=CSS, theme=gr.themes.Base()) 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>")
 
+    agent_state = gr.State("INITIAL") 
+    initial_topic_state = gr.State("")
+
     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)
-            init_button = gr.Button("Initialize Agent", scale=1, variant="primary")
+            init_button = gr.Button("Initialize Agent", scale=1)
     
-    initialization_status = gr.Markdown(visible=False)
-
     chatbot = gr.Chatbot(
         elem_id="chatbot",
         bubble_full_width=False,
         height=500,
-        visible=False,
-        type="messages"  # Use the modern 'messages' format
+        visible=False
     )
-    chat_input = gr.Textbox(placeholder="Enter your research topic...", interactive=False, visible=False)
     
-    # State variables
-    agent_state = gr.State("INITIAL") 
-    initial_topic_state = gr.State("")
+    with gr.Row(elem_id="chat-input-container"):
+        chat_input = gr.Textbox(placeholder="What would you like to research?", interactive=False, visible=False, show_label=False, scale=8)
+        submit_button = gr.Button("Submit", elem_id="submit-button", visible=False, scale=1)
 
     def handle_initialization(google_key, tavily_key):
-        init_status = initialize_models(google_key, tavily_key)
-        initial_messages = [{"role": "assistant", "content": "Agent initialized. Please enter your research topic to begin."}]
+        initialize_models(google_key, tavily_key)
         return {
-            initialization_status: gr.update(value=f"**Status:** {init_status}", visible=True),
-            chatbot: gr.update(visible=True, value=initial_messages),
+            chatbot: gr.update(visible=True, value=[(None, "Agent initialized. Please enter your research topic to begin.")]),
             chat_input: gr.update(interactive=True, visible=True),
+            submit_button: gr.update(visible=True),
             settings_accordion: gr.update(open=False)
         }
 
     def chat_step(user_input, history, current_agent_state, topic_state):
-        history.append({"role": "user", "content": user_input})
+        history = history or []
+        history.append((user_input, None))
         
         if current_agent_state == "INITIAL":
-            history.append({"role": "assistant", "content": "Thinking..."})
-            yield history, "CLARIFYING", user_input, gr.update(interactive=False)
-            
+            yield history, "CLARIFYING", user_input, gr.update(interactive=False, placeholder="Thinking...")
             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, "CLARIFYING", user_input, gr.update(interactive=True)
+            history[-1] = (user_input, "To give you the best report, could you answer these questions for me?\n\n" + questions)
+            yield history, "CLARIFYING", user_input, gr.update(interactive=True, placeholder="Provide your answers to the questions above...")
 
         elif current_agent_state == "CLARIFYING":
-            history.append({"role": "assistant", "content": "Generating full report..."})
-            yield history, "GENERATING", topic_state, gr.update(interactive=False)
+            thinking_message = "Got it. Generating your full research report. This will take a moment..."
+            history[-1] = (user_input, thinking_message)
+            yield history, "GENERATING", topic_state, gr.update(interactive=False, placeholder="Generating...")
             
             try:
                 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 = ""
-                final_report_md = ""
                 for update in report_generator:
-                    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}
+                    history[-1] = (user_input, status_updates)
                     yield history, "GENERATING", topic_state, gr.update(interactive=False)
                 
-                history.append({"role": "assistant", "content": final_report_md})
-                yield history, "INITIAL", "", gr.update(interactive=True)
+                full_report = status_updates # The final state of the stream is the report
+                history.append((None, full_report))
+                yield history, "INITIAL", "", gr.update(interactive=True, placeholder="Research complete. What's the next topic?")
             
             except Exception as e:
                 error_message = f"An error occurred: {str(e)}"
-                history.append({"role": "assistant", "content": error_message})
-                yield history, "INITIAL", "", gr.update(interactive=True)
+                history.append((None, error_message))
+                yield history, "INITIAL", "", 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=[initialization_status, chatbot, chat_input, settings_accordion]
+        outputs=[chatbot, chat_input, submit_button, settings_accordion]
     )
     
+    submit_action = submit_button.click(
+        fn=chat_step,
+        inputs=[chat_input, chatbot, agent_state, initial_topic_state],
+        outputs=[chatbot, agent_state, initial_topic_state, chat_input]
+    )
     chat_input.submit(
         fn=chat_step,
         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
     )
 
+    # Clear input after submission
+    submit_action.then(lambda: "", None, chat_input)
+    chat_input.submit(lambda: "", None, chat_input)
+
 app.launch(debug=True)