Update app.py

app.py

@@ -9,158 +9,126 @@ from research_agent.agent import get_clarifying_questions, research_and_plan, wr
 
 # --- CSS for styling the Gradio app ---
 CSS = """
-body { font-family: 'Inter', sans-serif; background-color: #F0F2F6; }
+body, .gradio-container { font-family: 'Inter', sans-serif; }
 .gradio-container { max-width: 960px !important; margin: auto !important; }
-h1 { text-align: center; font-size: 2.5em; color: #1E3A8A; }
+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; }
-.status_box {
-    background-color: #FFFFFF;
-    border-radius: 8px;
-    padding: 15px;
-    border: 1px solid #E5E7EB;
-    box-shadow: 0 2px 4px rgba(0,0,0,0.05);
-}
-.report_output {
-    background-color: #FFFFFF;
-    border-radius: 8px;
-    padding: 20px;
-    border: 1px solid #E5E7EB;
-    box-shadow: 0 4px 8px rgba(0,0,0,0.05);
-}
+.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; }
 """
 
-# --- Global variables for models (to avoid reloading) ---
-writer_model = None
-planner_model = None
-embedding_model = None
-reranker = None
-tavily_client = None
+# --- Model Initialization ---
 config = AgentConfig()
+writer_model, planner_model, embedding_model, reranker, tavily_client = None, None, None, None, None
 
-def initialize_models(google_api_key, tavily_api_key):
-    """Initializes all the necessary models and API clients."""
+def initialize_models(google_key, tavily_key):
     global writer_model, planner_model, embedding_model, reranker, tavily_client
-    
-    if not google_api_key or not tavily_api_key:
+    if not google_key or not tavily_key:
         raise gr.Error("API keys are required. Please provide both Google and Tavily API keys.")
-    
-    try:
-        genai.configure(api_key=google_api_key)
-        tavily_client = TavilyClient(api_key=tavily_api_key)
-
-        if writer_model is None:
-            writer_model = genai.GenerativeModel(config.WRITER_MODEL)
-        if planner_model is None:
-            planner_model = genai.GenerativeModel(config.WRITER_MODEL)
-        if embedding_model is None:
-            embedding_model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
-        if reranker is None:
-            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. Please check your API keys. Error: {str(e)}")
-
-def start_research_phase(topic, google_key, tavily_key):
-    """Phase 1: Get user topic and return clarifying questions."""
-    initialize_models(google_key, tavily_key)
-    
-    if not topic:
-        raise gr.Error("Research topic cannot be empty.")
-    
-    questions = get_clarifying_questions(planner_model, topic)
-    
-    # Show the next stage of the UI
-    return {
-        clarification_ui: gr.update(visible=True),
-        clarification_questions_display: gr.update(value=questions),
-        initial_ui: gr.update(visible=False)
-    }
-
-def generate_report_phase(topic, answers, google_key, tavily_key):
-    """Phase 2: Take answers and generate the full report, streaming progress."""
-    initialize_models(google_key, tavily_key)
-    
-    status_updates = "### Agent Status\n"
-    yield {
-        status_box: gr.update(value=status_updates + "-> Planning research...\n"),
-        final_report: gr.update(value=None)
-    }
-    
     try:
-        plan = research_and_plan(config, planner_model, tavily_client, topic, answers)
+        genai.configure(api_key=google_key)
+        tavily_client = TavilyClient(api_key=tavily_key)
+        writer_model = genai.GenerativeModel(config.WRITER_MODEL)
+        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!")
     except Exception as e:
-        raise gr.Error(f"Failed during planning phase: {e}")
-        
-    status_updates += f"**Research Plan:**\n- **Topic:** {plan['detailed_topic']}\n- **Sections:** {[s.title for s in plan['sections']]}\n\n---\n"
-    yield { status_box: gr.update(value=status_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
-        yield { status_box: gr.update(value=status_updates) }
-        
-    yield { final_report: gr.update(value=final_report_md) }
+        raise gr.Error(f"Failed to initialize models. Error: {str(e)}")
 
-# --- Build the Gradio Interface ---
+# --- 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 on a given topic, using AI to plan, search, and write a comprehensive report.")
-    
-    # State to hold the original topic
-    topic_state = gr.State()
+    gr.Markdown("This agent performs in-depth research using a multi-step AI process.", elem_classes="sub-header")
 
-    # --- UI Stage 1: Initial Query ---
-    with gr.Column(visible=True) as initial_ui:
-        with gr.Row():
-            google_api_key_input = gr.Textbox(label="Google API Key", type="password", placeholder="Enter your Google AI API Key")
-            tavily_api_key_input = gr.Textbox(label="Tavily API Key", type="password", placeholder="Enter your Tavily Search API Key")
+    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")
         
-        topic_input = gr.Textbox(label="Research Topic", placeholder="e.g., The future of renewable energy")
-        start_button = gr.Button("Start Research", variant="primary")
-
-    # --- UI Stage 2: Clarification ---
-    with gr.Column(visible=False) as clarification_ui:
-        gr.Markdown("### To give you the most relevant report, could you please clarify:")
-        clarification_questions_display = gr.Markdown(elem_classes="status_box")
-        clarification_answers_input = gr.Textbox(label="Your Answers", placeholder="Provide your answers to the questions above to tailor the research...")
-        generate_report_button = gr.Button("Generate Full Report", variant="primary")
+        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)
 
-    # --- UI Stage 3: Output ---
-    with gr.Column():
-        status_box = gr.Markdown(elem_classes="status_box", label="Agent Thought Process", visible=False)
-        final_report = gr.Markdown(elem_classes="report_output", label="Final Research Report", visible=False)
+    # --- State Management ---
+    # agent_state can be: "INITIAL", "CLARIFYING", "GENERATING", "DONE"
+    agent_state = gr.State("INITIAL") 
+    initial_topic_state = gr.State("")
 
     # --- Event Handlers ---
-    def show_outputs():
+    def handle_initialization(google_key, tavily_key):
+        initialize_models(google_key, tavily_key)
         return {
-            status_box: gr.update(visible=True),
-            final_report: gr.update(visible=True)
+            chat_input: gr.update(interactive=True, placeholder="Models Initialized! What would you like to research?"),
+            init_button: gr.update(value="Agent Ready!", interactive=False)
         }
 
-    start_button.click(
-        fn=start_research_phase,
-        inputs=[topic_input, google_api_key_input, tavily_api_key_input],
-        outputs=[initial_ui, clarification_ui, clarification_questions_display]
-    ).then(
-        fn=lambda topic: topic,
-        inputs=[topic_input],
-        outputs=[topic_state] # Save the topic for the next step
+    def chat_step(user_input, history, state, original_topic):
+        history = history or []
+        
+        if state == "INITIAL":
+            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...")}
+            
+            # 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...")}
+
+        elif state == "CLARIFYING":
+            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}
+            
+            # 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}
+            
+            # 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?")}
+
+    init_button.click(
+        fn=handle_initialization,
+        inputs=[google_api_key_input, tavily_api_key_input],
+        outputs=[chat_input, init_button]
     )
     
-    generate_report_button.click(
-        fn=show_outputs,
-        outputs=[status_box, final_report]
+    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(
-        fn=generate_report_phase,
-        inputs=[topic_state, clarification_answers_input, google_api_key_input, tavily_api_key_input],
-        outputs=[status_box, final_report]
+        fn=lambda: "",
+        outputs=[chat_input]
     )
 
 app.launch(debug=True)