Update app.py

app.py

@@ -1,166 +1,166 @@
-import os
-import gradio as gr
-import google.generativeai as genai
-from tavily import TavilyClient
-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 ---
-CSS = """
-body { font-family: 'Inter', sans-serif; background-color: #F0F2F6; }
-.gradio-container { max-width: 960px !important; margin: auto !important; }
-h1 { text-align: center; font-size: 2.5em; color: #1E3A8A; }
-.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);
-}
-"""
-
-# --- Global variables for models (to avoid reloading) ---
-writer_model = None
-planner_model = None
-embedding_model = None
-reranker = None
-tavily_client = None
-config = AgentConfig()
-
-def initialize_models(google_api_key, tavily_api_key):
-    """Initializes all the necessary models and API clients."""
-    global writer_model, planner_model, embedding_model, reranker, tavily_client
-    
-    if not google_api_key or not tavily_api_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)
-    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) }
-
-# --- Build the Gradio Interface ---
-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()
-
-    # --- UI Stage 1: Initial Query ---
-    with gr.Box(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")
-        
-        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.Box(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")
-
-    # --- 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)
-
-    # --- Event Handlers ---
-    def show_outputs():
-        return {
-            status_box: gr.update(visible=True),
-            final_report: gr.update(visible=True)
-        }
-
-    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
-    )
-    
-    generate_report_button.click(
-        fn=show_outputs,
-        outputs=[status_box, final_report]
-    ).then(
-        fn=generate_report_phase,
-        inputs=[topic_state, clarification_answers_input, google_api_key_input, tavily_api_key_input],
-        outputs=[status_box, final_report]
-    )
-
+import os
+import gradio as gr
+import google.generativeai as genai
+from tavily import TavilyClient
+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 ---
+CSS = """
+body { font-family: 'Inter', sans-serif; background-color: #F0F2F6; }
+.gradio-container { max-width: 960px !important; margin: auto !important; }
+h1 { text-align: center; font-size: 2.5em; color: #1E3A8A; }
+.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);
+}
+"""
+
+# --- Global variables for models (to avoid reloading) ---
+writer_model = None
+planner_model = None
+embedding_model = None
+reranker = None
+tavily_client = None
+config = AgentConfig()
+
+def initialize_models(google_api_key, tavily_api_key):
+    """Initializes all the necessary models and API clients."""
+    global writer_model, planner_model, embedding_model, reranker, tavily_client
+    
+    if not google_api_key or not tavily_api_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)
+    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) }
+
+# --- Build the Gradio Interface ---
+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()
+
+    # --- 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")
+        
+        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")
+
+    # --- 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)
+
+    # --- Event Handlers ---
+    def show_outputs():
+        return {
+            status_box: gr.update(visible=True),
+            final_report: gr.update(visible=True)
+        }
+
+    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
+    )
+    
+    generate_report_button.click(
+        fn=show_outputs,
+        outputs=[status_box, final_report]
+    ).then(
+        fn=generate_report_phase,
+        inputs=[topic_state, clarification_answers_input, google_api_key_input, tavily_api_key_input],
+        outputs=[status_box, final_report]
+    )
+
 app.launch(debug=True)