Update app.py

app.py

@@ -3,39 +3,320 @@ import gradio as gr
 import google.generativeai as genai
 from tavily import TavilyClient
 from sentence_transformers import SentenceTransformer, CrossEncoder
+import markdown
+from weasyprint import HTML, CSS as WeasyCSS
+from datetime import datetime
+import tempfile
+import re
 
 from research_agent.config import AgentConfig
 from research_agent.agent import get_clarifying_questions, research_and_plan, write_report_stream
 
-
 google_key = os.getenv("GOOGLE_API_KEY")
 tavily_key = os.getenv("TAVILY_API_KEY")
 
 if not google_key or not tavily_key:
     raise ValueError("API keys not found.")
 
+# Enhanced CSS for a professional research interface
+CSS = """
+@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
 
+:root {
+    --primary-color: #2563eb;
+    --primary-hover: #1d4ed8;
+    --bg-primary: #0f172a;
+    --bg-secondary: #1e293b;
+    --bg-tertiary: #334155;
+    --text-primary: #f1f5f9;
+    --text-secondary: #cbd5e1;
+    --text-muted: #94a3b8;
+    --border-color: #334155;
+    --success-color: #10b981;
+    --warning-color: #f59e0b;
+    --error-color: #ef4444;
+}
 
-CSS = """
-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; }
-#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; }
+body, .gradio-container {
+    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
+    background-color: var(--bg-primary) !important;
+    color: var(--text-primary) !important;
+}
+
+.gradio-container {
+    max-width: 1200px !important;
+    margin: 0 auto !important;
+    padding: 2rem !important;
+}
+
+/* Header Styling */
+.header-container {
+    text-align: center;
+    margin-bottom: 3rem;
+    padding: 2rem;
+    background: linear-gradient(135deg, var(--bg-secondary) 0%, var(--bg-tertiary) 100%);
+    border-radius: 16px;
+    border: 1px solid var(--border-color);
+}
+
+h1 {
+    font-size: 3rem;
+    font-weight: 700;
+    background: linear-gradient(135deg, #60a5fa 0%, #a78bfa 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    margin-bottom: 0.5rem;
+}
+
+.subtitle {
+    color: var(--text-secondary);
+    font-size: 1.25rem;
+    font-weight: 400;
+}
+
+/* Status Bar */
+.status-bar {
+    background: var(--bg-secondary);
+    border: 1px solid var(--border-color);
+    border-radius: 12px;
+    padding: 1rem 1.5rem;
+    margin-bottom: 2rem;
+    display: flex;
+    align-items: center;
+    justify-content: space-between;
+}
+
+.status-indicator {
+    display: flex;
+    align-items: center;
+    gap: 0.5rem;
+}
+
+.status-dot {
+    width: 8px;
+    height: 8px;
+    border-radius: 50%;
+    background: var(--success-color);
+    animation: pulse 2s infinite;
+}
+
+@keyframes pulse {
+    0% { opacity: 1; }
+    50% { opacity: 0.5; }
+    100% { opacity: 1; }
+}
+
+/* Chat Interface */
+#chatbot {
+    background: var(--bg-secondary) !important;
+    border: 1px solid var(--border-color) !important;
+    border-radius: 16px !important;
+    overflow: hidden !important;
+}
+
+#chatbot .message {
+    border: none !important;
+    padding: 1.5rem !important;
+}
+
+#chatbot .user {
+    background: var(--bg-tertiary) !important;
+    border-left: 4px solid var(--primary-color) !important;
+}
+
+#chatbot .bot {
+    background: var(--bg-secondary) !important;
+}
+
+/* Progress Indicators */
+.progress-container {
+    background: var(--bg-tertiary);
+    border-radius: 8px;
+    padding: 1rem;
+    margin: 1rem 0;
+}
+
+.progress-bar {
+    height: 4px;
+    background: var(--border-color);
+    border-radius: 2px;
+    overflow: hidden;
+    margin-top: 0.5rem;
+}
+
+.progress-fill {
+    height: 100%;
+    background: linear-gradient(90deg, var(--primary-color) 0%, #60a5fa 100%);
+    transition: width 0.3s ease;
+    animation: shimmer 2s infinite;
+}
+
+@keyframes shimmer {
+    0% { opacity: 0.8; }
+    50% { opacity: 1; }
+    100% { opacity: 0.8; }
+}
+
+/* Input Area */
+.input-container {
+    background: var(--bg-secondary);
+    border: 1px solid var(--border-color);
+    border-radius: 12px;
+    padding: 1.5rem;
+    margin-top: 2rem;
+}
+
+#chat-input textarea {
+    background: var(--bg-tertiary) !important;
+    color: var(--text-primary) !important;
+    border: 1px solid var(--border-color) !important;
+    border-radius: 8px !important;
+    padding: 1rem !important;
+    font-size: 1rem !important;
+    transition: all 0.2s ease !important;
+}
+
+#chat-input textarea:focus {
+    border-color: var(--primary-color) !important;
+    box-shadow: 0 0 0 3px rgba(37, 99, 235, 0.1) !important;
+}
+
+/* Buttons */
+.gr-button {
+    background: var(--primary-color) !important;
+    color: white !important;
+    border: none !important;
+    border-radius: 8px !important;
+    padding: 0.75rem 1.5rem !important;
+    font-weight: 600 !important;
+    transition: all 0.2s ease !important;
+    cursor: pointer !important;
+}
+
+.gr-button:hover {
+    background: var(--primary-hover) !important;
+    transform: translateY(-1px);
+    box-shadow: 0 4px 12px rgba(37, 99, 235, 0.3) !important;
+}
+
+.gr-button.secondary {
+    background: var(--bg-tertiary) !important;
+    color: var(--text-primary) !important;
+}
+
+.gr-button.secondary:hover {
+    background: #475569 !important;
+}
+
+/* Report Display */
+.report-section {
+    background: var(--bg-secondary);
+    border: 1px solid var(--border-color);
+    border-radius: 12px;
+    padding: 2rem;
+    margin: 1rem 0;
+}
+
+.report-section h2 {
+    color: var(--text-primary);
+    font-size: 1.75rem;
+    font-weight: 600;
+    margin-bottom: 1rem;
+    padding-bottom: 0.75rem;
+    border-bottom: 2px solid var(--border-color);
+}
+
+.report-section h3 {
+    color: var(--text-secondary);
+    font-size: 1.25rem;
+    font-weight: 500;
+    margin: 1.5rem 0 0.75rem 0;
+}
+
+.source-list {
+    background: var(--bg-tertiary);
+    border-radius: 8px;
+    padding: 1rem;
+    margin-top: 1rem;
+}
+
+.source-item {
+    display: flex;
+    align-items: center;
+    gap: 0.5rem;
+    padding: 0.5rem 0;
+    color: var(--text-secondary);
+    text-decoration: none;
+    transition: color 0.2s ease;
+}
+
+.source-item:hover {
+    color: var(--primary-color);
+}
+
+/* Loading States */
+.thinking-indicator {
+    display: flex;
+    align-items: center;
+    gap: 0.75rem;
+    color: var(--text-secondary);
+    font-style: italic;
+}
+
+.thinking-dots {
+    display: flex;
+    gap: 0.25rem;
+}
+
+.thinking-dots span {
+    width: 6px;
+    height: 6px;
+    background: var(--text-muted);
+    border-radius: 50%;
+    animation: bounce 1.4s infinite ease-in-out both;
+}
+
+.thinking-dots span:nth-child(1) { animation-delay: -0.32s; }
+.thinking-dots span:nth-child(2) { animation-delay: -0.16s; }
+
+@keyframes bounce {
+    0%, 80%, 100% { transform: scale(0); }
+    40% { transform: scale(1); }
+}
+
+/* Export Options */
+.export-container {
+    background: var(--bg-secondary);
+    border: 1px solid var(--border-color);
+    border-radius: 12px;
+    padding: 1.5rem;
+    margin-top: 2rem;
+}
+
+.export-buttons {
+    display: flex;
+    gap: 1rem;
+    margin-top: 1rem;
+}
+
+/* Responsive Design */
+@media (max-width: 768px) {
+    .gradio-container {
+        padding: 1rem !important;
+    }
+    
+    h1 {
+        font-size: 2rem;
+    }
+    
+    .export-buttons {
+        flex-direction: column;
+    }
+}
 """
 
-# --- Model Initialization ---
+# Initialize models
 config = AgentConfig()
-
 writer_model, planner_model, embedding_model, reranker, tavily_client = None, None, None, None, None
 IS_PROCESSING = False
 
@@ -43,7 +324,6 @@ def initialize_models():
     """Initializes all the models and clients using keys from environment variables."""
     global writer_model, planner_model, embedding_model, reranker, tavily_client, IS_PROCESSING
     try:
-
         genai.configure(api_key=google_key)
         tavily_client = TavilyClient(api_key=tavily_key)
         writer_model = genai.GenerativeModel(config.WRITER_MODEL)
@@ -51,119 +331,252 @@ def initialize_models():
         embedding_model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
         reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', device='cpu')
     except Exception as e:
-        # This error will be displayed in the Hugging Face logs
         print(f"FATAL: Failed to initialize models. Error: {str(e)}")
-        # Raise an exception to stop the app from running incorrectly
         raise gr.Error(f"Failed to initialize models. Please check the logs. Error: {str(e)}")
     IS_PROCESSING = False
     print("Models and clients initialized successfully.")
 
-# --- 3. Initialize models on application startup ---
+# Initialize models on startup
 initialize_models()
 
-# --- Gradio Application Logic ---
-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>")
+# Helper functions for better UI
+def format_progress_message(message):
+    """Formats progress messages with visual indicators"""
+    if "Step" in message:
+        return f"🔄 **{message}**"
+    elif "Searching" in message:
+        return f"🔍 {message}"
+    elif "Found" in message:
+        return f"✅ {message}"
+    elif "Processing" in message:
+        return f"⚙️ {message}"
+    elif "Writing" in message or "Synthesizing" in message:
+        return f"✍️ {message}"
+    elif "Fact-checking" in message:
+        return f"🔎 {message}"
+    else:
+        return message
 
-    agent_state = gr.State("INITIAL")
-    initial_topic_state = gr.State("")
+def export_to_pdf(report_content, filename="research_report.pdf"):
+    """Exports the report to PDF with proper formatting"""
+    try:
+        # Convert markdown to HTML
+        html_content = markdown.markdown(report_content, extensions=['extra', 'codehilite'])
+        
+        # Add CSS for PDF
+        pdf_css = """
+        @page { size: A4; margin: 2cm; }
+        body { font-family: Arial, sans-serif; line-height: 1.6; color: #333; }
+        h1 { color: #2563eb; border-bottom: 2px solid #2563eb; padding-bottom: 10px; }
+        h2 { color: #1e40af; margin-top: 30px; }
+        h3 { color: #3730a3; }
+        pre { background: #f3f4f6; padding: 10px; border-radius: 4px; }
+        code { background: #e5e7eb; padding: 2px 4px; border-radius: 2px; }
+        blockquote { border-left: 4px solid #2563eb; padding-left: 16px; color: #6b7280; }
+        """
+        
+        # Create PDF
+        with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as tmp_file:
+            HTML(string=f"<html><body>{html_content}</body></html>").write_pdf(
+                tmp_file.name,
+                stylesheets=[WeasyCSS(string=pdf_css)]
+            )
+            return tmp_file.name
+    except Exception as e:
+        print(f"Error creating PDF: {e}")
+        return None
 
+def chat_step_wrapper(user_input, history, current_agent_state, topic_state, progress_state):
+    """Enhanced wrapper with progress tracking"""
+    global IS_PROCESSING
+    if IS_PROCESSING:
+        print("Ignoring duplicate request while processing.")
+        if False: yield
+        return
 
+    IS_PROCESSING = True
+    try:
+        for update in chat_step(user_input, history, current_agent_state, topic_state, progress_state):
+            yield update
+    except Exception as e:
+        error_message = f"❌ **Error**: {str(e)}"
+        history.append((None, error_message))
+        yield history, "INITIAL", "", {}, gr.update(interactive=True, placeholder="Let's try again. What would you like to research?"), None, gr.update(visible=False)
+    finally:
+        IS_PROCESSING = False
+        print("Processing finished. Lock released.")
+
+def chat_step(user_input, history, current_agent_state, topic_state, progress_state):
+    """Enhanced chat step with visual progress tracking"""
+    history = history or []
+    history.append((user_input, None))
+    
+    if current_agent_state == "INITIAL":
+        yield history, "CLARIFYING", user_input, progress_state, gr.update(interactive=False, placeholder="Analyzing your topic..."), None, gr.update(visible=False)
+        
+        # Show thinking animation
+        thinking_msg = """<div class="thinking-indicator">
+            <span>Analyzing your research topic</span>
+            <div class="thinking-dots">
+                <span></span><span></span><span></span>
+            </div>
+        </div>"""
+        history[-1] = (user_input, thinking_msg)
+        yield history, "CLARIFYING", user_input, progress_state, gr.update(interactive=False), None, gr.update(visible=False)
+        
+        questions = get_clarifying_questions(planner_model, user_input)
+        formatted_questions = f"""
+### 🎯 Let's refine your research
+
+To create the most comprehensive report on **{user_input}**, I'd like to understand your specific interests:
+
+{questions}
+
+Please provide your answers below to help me tailor the research to your needs.
+"""
+        history[-1] = (user_input, formatted_questions)
+        yield history, "CLARIFYING", user_input, progress_state, gr.update(interactive=True, placeholder="Type your answers here..."), None, gr.update(visible=False)
+
+    elif current_agent_state == "CLARIFYING":
+        # Show initial processing message
+        history[-1] = (user_input, "📋 **Perfect! I have all the information I need.**\n\nStarting deep research process...")
+        yield history, "GENERATING", topic_state, {"current_step": 1, "total_steps": 5}, gr.update(interactive=False, placeholder="Generating report..."), None, gr.update(visible=False)
+        
+        try:
+            # Research and planning phase
+            plan = research_and_plan(config, planner_model, tavily_client, topic_state, user_input)
+            
+            # Show research plan
+            sections_preview = "\n".join([f"  {i+1}. {s.title}" for i, s in enumerate(plan['sections'])])
+            planning_update = f"""
+### 📊 Research Plan Created
+
+**Topic**: {plan['detailed_topic']}
+
+**Report Structure**:
+{sections_preview}
+
+Now conducting deep research and writing each section...
+"""
+            history[-1] = (user_input, planning_update)
+            yield history, "GENERATING", topic_state, {"current_step": 2, "total_steps": 5}, gr.update(interactive=False), None, gr.update(visible=False)
+            
+            # Stream report generation
+            report_generator = write_report_stream(config, writer_model, tavily_client, embedding_model, reranker, plan)
+            
+            full_report = ""
+            for update in report_generator:
+                # Format the update for better display
+                if update.startswith("#"):
+                    full_report = update
+                    # Add progress indicators to the report display
+                    display_report = full_report
+                else:
+                    # Show progress updates
+                    progress_msg = format_progress_message(update)
+                    display_report = f"{planning_update}\n\n---\n\n**Current Progress**: {progress_msg}\n\n---\n\n### 📄 Report Preview:\n\n{full_report}"
+                
+                history[-1] = (user_input, display_report)
+                yield history, "GENERATING", topic_state, progress_state, gr.update(interactive=False), None, gr.update(visible=False)
+            
+            # Final report display
+            completion_message = f"""
+### ✅ Research Complete!
+
+Your comprehensive research report is ready. You can:
+- 📥 Download as PDF using the button below
+- 📋 Copy the text directly from the report
+- 🔄 Start a new research topic
+
+---
+
+{full_report}
+"""
+            history[-1] = (user_input, completion_message)
+            
+            # Enable PDF download
+            pdf_path = export_to_pdf(full_report)
+            
+            yield history, "INITIAL", "", {}, gr.update(interactive=True, placeholder="What would you like to research next?"), pdf_path, gr.update(visible=True)
+            
+        except Exception as e:
+            error_msg = f"❌ **Error during research**: {str(e)}\n\nPlease try again with a different topic or check your API keys."
+            history.append((None, error_msg))
+            yield history, "INITIAL", "", {}, gr.update(interactive=True, placeholder="Let's try again. What would you like to research?"), None, gr.update(visible=False)
+
+# Build the Gradio interface
+with gr.Blocks(css=CSS, theme=gr.themes.Base()) as app:
+    # Header
+    gr.HTML("""
+        <div class="header-container">
+            <h1>DeepSearch Research Agent</h1>
+            <p class="subtitle">AI-powered comprehensive research and analysis</p>
+        </div>
+    """)
+    
+    # Status bar
+    gr.HTML("""
+        <div class="status-bar">
+            <div class="status-indicator">
+                <span class="status-dot"></span>
+                <span>System Online</span>
+            </div>
+            <div>
+                <span style="color: var(--text-muted);">Powered by Gemini & Tavily</span>
+            </div>
+        </div>
+    """)
+    
+    # State management
+    agent_state = gr.State("INITIAL")
+    initial_topic_state = gr.State("")
+    progress_state = gr.State({})
+    
+    # Chat interface
     chatbot = gr.Chatbot(
         elem_id="chatbot",
         bubble_full_width=False,
-        height=500,
-        # --- 5. Make the chatbot visible from the start ---
+        height=600,
         visible=True,
-        value=[(None, "Agent is ready. What would you like to research?")]
+        value=[(None, "👋 **Welcome to DeepSearch!**\n\nI'm your AI research assistant. I can help you create comprehensive, well-researched reports on any topic.\n\n**How it works:**\n1. Tell me what you'd like to research\n2. I'll ask a few clarifying questions\n3. I'll conduct deep research and write a detailed report\n4. You'll get a downloadable PDF with all sources\n\n**What would you like to research today?**")],
+        avatar_images=(None, "🔬")
     )
-
-    with gr.Row(elem_id="chat-input-container"):
-        chat_input = gr.Textbox(
-            placeholder="What would you like to research?",
-            # --- 6. Make the input box interactive and visible from the start ---
-            interactive=True,
-            visible=True,
-            show_label=False,
-            scale=8
-        )
-        submit_button = gr.Button("Submit", elem_id="submit-button", visible=True, scale=1)
-
-    # The handle_initialization function is no longer needed.
-
-    def chat_step_wrapper(user_input, history, current_agent_state, topic_state):
-        """A wrapper to manage the processing lock."""
-        global IS_PROCESSING
-        if IS_PROCESSING:
-            print("Ignoring duplicate request while processing.")
-            # This is a generator, so we must yield something. An empty yield is fine.
-            if False: yield
-            return
-
-        IS_PROCESSING = True
-        try:
-            # Yield all updates from the actual agent logic
-            for update in chat_step(user_input, history, current_agent_state, topic_state):
-                yield update
-        except Exception as e:
-            error_message = f"An error occurred: {str(e)}"
-            history.append((None, error_message))
-            yield history, "INITIAL", "", gr.update(interactive=True, placeholder="Let's try again. What's the topic?")
-        finally:
-            # Release the lock once the generator is exhausted or an error occurs
-            IS_PROCESSING = False
-            print("Processing finished. Lock released.")
-
-    def chat_step(user_input, history, current_agent_state, topic_state):
-        history = history or []
-        history.append((user_input, None))
-
-        if current_agent_state == "INITIAL":
-            yield history, "CLARIFYING", user_input, gr.update(interactive=False, placeholder="Thinking...")
-            questions = get_clarifying_questions(planner_model, user_input)
-            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":
-            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)
-
-                stream_content = ""
-                for update in report_generator:
-                    stream_content = update
-                    history[-1] = (user_input, stream_content)
-                    yield history, "GENERATING", topic_state, gr.update(interactive=False)
-
-                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((None, error_message))
-                yield history, "INITIAL", "", gr.update(interactive=True, placeholder="Let's try again. What's the topic?")
-
-    # --- Event Listeners ---
-
+    
+    # Input area
+    with gr.Group(elem_classes="input-container"):
+        with gr.Row():
+            chat_input = gr.Textbox(
+                placeholder="Enter your research topic (e.g., 'Impact of AI on healthcare', 'Climate change solutions', 'History of quantum computing')",
+                interactive=True,
+                visible=True,
+                show_label=False,
+                scale=8,
+                elem_id="chat-input"
+            )
+            submit_button = gr.Button("🚀 Start Research", scale=2, variant="primary")
+    
+    # Export section
+    with gr.Group(elem_classes="export-container", visible=False) as export_group:
+        gr.Markdown("### 📥 Export Options")
+        with gr.Row(elem_classes="export-buttons"):
+            pdf_download = gr.File(label="Download PDF Report", visible=False)
+            
+    # Event handlers
     submit_event = submit_button.click(
         fn=chat_step_wrapper,
-        inputs=[chat_input, chatbot, agent_state, initial_topic_state],
-        outputs=[chatbot, agent_state, initial_topic_state, chat_input],
+        inputs=[chat_input, chatbot, agent_state, initial_topic_state, progress_state],
+        outputs=[chatbot, agent_state, initial_topic_state, progress_state, chat_input, pdf_download, export_group],
     ).then(
         fn=lambda: gr.update(value=""),
         inputs=None,
         outputs=[chat_input],
         queue=False
     )
-
+    
     chat_input.submit(
         fn=chat_step_wrapper,
-        inputs=[chat_input, chatbot, agent_state, initial_topic_state],
-        outputs=[chatbot, agent_state, initial_topic_state, chat_input],
+        inputs=[chat_input, chatbot, agent_state, initial_topic_state, progress_state],
+        outputs=[chatbot, agent_state, initial_topic_state, progress_state, chat_input, pdf_download, export_group],
     ).then(
         fn=lambda: gr.update(value=""),
         inputs=None,
@@ -171,4 +584,7 @@ with gr.Blocks(css=CSS, theme=gr.themes.Base()) as app:
         queue=False
     )
 
-app.launch(debug=True)
+# Launch the app
+if __name__ == "__main__":
+    app.queue()
+    app.launch(debug=True, share=False)