Update research_agent/agent.py

research_agent/agent.py

@@ -1,152 +1,723 @@
 import os
+import gradio as gr
+import google.generativeai as genai
+from tavily import TavilyClient
+from sentence_transformers import SentenceTransformer, CrossEncoder
+from datetime import datetime
 import json
-from pydantic import BaseModel, Field
-from IPython.display import display, Markdown
-
-# Local module imports
-from .config import AgentConfig
-from . import prompts
-from .llm_utils import run_gemini_json_completion, run_gemini_text_completion
-from .search import gather_research
-from .rag_pipeline import RAGPipeline
-
-# For running async code in notebook
-import nest_asyncio
-nest_asyncio.apply()
-
-class Section(BaseModel):
-    title: str = Field(description="The title of the report section.")
-    description: str = Field(description="A detailed description of what the section will cover, including key sub-topics.")
-
-def run_verification_step(writer_model, section_text: str, research_context: str):
-    """A new step to verify claims and check for hallucinations."""
-    verification_prompt = prompts.verification_prompt_template.format(
-        section_text=section_text,
-        research_context=research_context
-    )
+import time
+
+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.")
+
+# Professional CSS with dark theme
+CSS = """
+@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&family=JetBrains+Mono&display=swap');
+
+* {
+    transition: all 0.3s ease;
+}
+
+body, .gradio-container {
+    font-family: 'Inter', -apple-system, system-ui, sans-serif !important;
+    background: #0a0a0a !important;
+    color: #e4e4e7 !important;
+}
+
+.gradio-container {
+    max-width: 1200px !important;
+    margin: 0 auto !important;
+}
+
+/* Header */
+.header {
+    background: linear-gradient(135deg, #1e1b4b 0%, #312e81 100%);
+    padding: 3rem 2rem;
+    border-radius: 24px;
+    margin-bottom: 2rem;
+    text-align: center;
+    box-shadow: 0 20px 40px rgba(0,0,0,0.5);
+}
+
+.header h1 {
+    font-size: 3.5rem;
+    font-weight: 800;
+    background: linear-gradient(135deg, #818cf8 0%, #c084fc 50%, #f472b6 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    margin: 0;
+    letter-spacing: -1px;
+}
+
+.header p {
+    color: #a5b4fc;
+    font-size: 1.25rem;
+    margin-top: 0.5rem;
+    font-weight: 300;
+}
+
+/* Status Bar */
+.status-bar {
+    background: #18181b;
+    border: 1px solid #27272a;
+    border-radius: 16px;
+    padding: 1rem 1.5rem;
+    margin-bottom: 2rem;
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+}
+
+.status-indicator {
+    display: flex;
+    align-items: center;
+    gap: 0.75rem;
+}
+
+.pulse-dot {
+    width: 10px;
+    height: 10px;
+    background: #22c55e;
+    border-radius: 50%;
+    box-shadow: 0 0 0 0 rgba(34, 197, 94, 1);
+    animation: pulse-animation 2s infinite;
+}
+
+@keyframes pulse-animation {
+    0% {
+        box-shadow: 0 0 0 0 rgba(34, 197, 94, 0.7);
+    }
+    70% {
+        box-shadow: 0 0 0 10px rgba(34, 197, 94, 0);
+    }
+    100% {
+        box-shadow: 0 0 0 0 rgba(34, 197, 94, 0);
+    }
+}
+
+/* Chat Interface */
+#chatbot {
+    background: #18181b !important;
+    border: 1px solid #27272a !important;
+    border-radius: 20px !important;
+    box-shadow: 0 10px 30px rgba(0,0,0,0.5) !important;
+}
+
+#chatbot .message-wrap {
+    padding: 0 !important;
+}
+
+#chatbot .message {
+    padding: 1.5rem !important;
+    margin: 0.5rem !important;
+    border-radius: 16px !important;
+    font-size: 1.05rem !important;
+    line-height: 1.7 !important;
+}
+
+#chatbot .user {
+    background: linear-gradient(135deg, #4c1d95 0%, #5b21b6 100%) !important;
+    color: #f3f4f6 !important;
+    margin-left: 20% !important;
+    box-shadow: 0 4px 12px rgba(139, 92, 246, 0.3) !important;
+}
+
+#chatbot .bot {
+    background: #27272a !important;
+    color: #e4e4e7 !important;
+    margin-right: 20% !important;
+    border: 1px solid #3f3f46 !important;
+}
+
+#chatbot .bot h1, #chatbot .bot h2, #chatbot .bot h3 {
+    color: #a78bfa !important;
+    margin-top: 1.5rem !important;
+    margin-bottom: 1rem !important;
+}
+
+#chatbot .bot h1 { font-size: 2rem !important; }
+#chatbot .bot h2 { font-size: 1.5rem !important; }
+#chatbot .bot h3 { font-size: 1.25rem !important; }
+
+#chatbot .bot strong {
+    color: #c4b5fd !important;
+}
+
+#chatbot .bot code {
+    background: #374151 !important;
+    color: #fbbf24 !important;
+    padding: 2px 6px !important;
+    border-radius: 4px !important;
+    font-family: 'JetBrains Mono', monospace !important;
+}
+
+#chatbot .bot a {
+    color: #60a5fa !important;
+    text-decoration: none !important;
+    border-bottom: 1px solid transparent !important;
+    transition: border-color 0.2s !important;
+}
+
+#chatbot .bot a:hover {
+    border-bottom-color: #60a5fa !important;
+}
+
+/* Progress Messages */
+.progress-message {
+    background: #1e1b4b;
+    border-left: 4px solid #6366f1;
+    padding: 1rem 1.5rem;
+    margin: 1rem 0;
+    border-radius: 0 8px 8px 0;
+}
+
+/* Input Area */
+.input-group {
+    background: #18181b;
+    border: 1px solid #27272a;
+    border-radius: 16px;
+    padding: 1.5rem;
+    margin-top: 2rem;
+}
+
+#user-input textarea {
+    background: #27272a !important;
+    border: 2px solid #3f3f46 !important;
+    color: #f3f4f6 !important;
+    border-radius: 12px !important;
+    padding: 1rem !important;
+    font-size: 1.05rem !important;
+    min-height: 80px !important;
+}
+
+#user-input textarea:focus {
+    border-color: #6366f1 !important;
+    outline: none !important;
+    box-shadow: 0 0 0 3px rgba(99, 102, 241, 0.1) !important;
+}
+
+/* Buttons */
+.submit-btn {
+    background: linear-gradient(135deg, #6366f1 0%, #8b5cf6 100%) !important;
+    color: white !important;
+    font-weight: 600 !important;
+    font-size: 1.1rem !important;
+    padding: 0.875rem 2rem !important;
+    border: none !important;
+    border-radius: 12px !important;
+    cursor: pointer !important;
+    box-shadow: 0 4px 12px rgba(99, 102, 241, 0.4) !important;
+}
+
+.submit-btn:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 6px 20px rgba(99, 102, 241, 0.5) !important;
+}
+
+.submit-btn:active {
+    transform: translateY(0) !important;
+}
+
+/* Examples */
+.examples-container {
+    background: #18181b;
+    border: 1px solid #27272a;
+    border-radius: 16px;
+    padding: 1.5rem;
+    margin-top: 2rem;
+}
+
+.examples-container h3 {
+    color: #a78bfa;
+    margin-bottom: 1rem;
+    font-size: 1.25rem;
+    font-weight: 600;
+}
+
+.examples-grid {
+    display: grid;
+    grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
+    gap: 1rem;
+}
+
+.example-card {
+    background: #27272a;
+    border: 1px solid #3f3f46;
+    padding: 1rem;
+    border-radius: 12px;
+    cursor: pointer;
+    transition: all 0.2s;
+}
+
+.example-card:hover {
+    background: #374151;
+    border-color: #6366f1;
+    transform: translateY(-2px);
+    box-shadow: 0 4px 12px rgba(0,0,0,0.5);
+}
+
+/* Loading Animation */
+.loading-wave {
+    display: inline-flex;
+    gap: 4px;
+}
+
+.loading-wave span {
+    width: 4px;
+    height: 20px;
+    background: #6366f1;
+    border-radius: 2px;
+    animation: wave 1.2s linear infinite;
+}
+
+.loading-wave span:nth-child(2) { animation-delay: -1.1s; }
+.loading-wave span:nth-child(3) { animation-delay: -1s; }
+.loading-wave span:nth-child(4) { animation-delay: -0.9s; }
+.loading-wave span:nth-child(5) { animation-delay: -0.8s; }
+
+@keyframes wave {
+    0%, 40%, 100% { transform: scaleY(0.4); }
+    20% { transform: scaleY(1); }
+}
+
+/* Report Sections */
+.report-section {
+    background: #1e1b4b;
+    border: 1px solid #312e81;
+    border-radius: 12px;
+    padding: 1.5rem;
+    margin: 1rem 0;
+}
+
+.source-badge {
+    display: inline-block;
+    background: #374151;
+    color: #60a5fa;
+    padding: 4px 12px;
+    border-radius: 16px;
+    font-size: 0.875rem;
+    margin: 0.25rem;
+}
+
+/* Export Section */
+.export-section {
+    background: #18181b;
+    border: 1px solid #27272a;
+    border-radius: 16px;
+    padding: 2rem;
+    margin-top: 2rem;
+    text-align: center;
+}
+
+.export-buttons {
+    display: flex;
+    gap: 1rem;
+    justify-content: center;
+    margin-top: 1.5rem;
+}
+
+.export-btn {
+    background: #27272a !important;
+    border: 1px solid #3f3f46 !important;
+    color: #e4e4e7 !important;
+    padding: 0.75rem 1.5rem !important;
+    border-radius: 10px !important;
+    font-weight: 500 !important;
+    cursor: pointer !important;
+    display: flex !important;
+    align-items: center !important;
+    gap: 0.5rem !important;
+}
+
+.export-btn:hover {
+    background: #374151 !important;
+    border-color: #6366f1 !important;
+}
+
+/* Responsive */
+@media (max-width: 768px) {
+    .header h1 { font-size: 2.5rem; }
+    #chatbot .user { margin-left: 10% !important; }
+    #chatbot .bot { margin-right: 10% !important; }
+    .examples-grid { grid-template-columns: 1fr; }
+}
+"""
+
+# Initialize models
+config = AgentConfig()
+writer_model, planner_model, embedding_model, reranker, tavily_client = None, None, None, None, None
+IS_PROCESSING = False
+
+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)
+        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')
+    except Exception as e:
+        print(f"FATAL: Failed to initialize models. Error: {str(e)}")
+        raise gr.Error(f"Failed to initialize models. Please check the logs. Error: {str(e)}")
+    IS_PROCESSING = False
+    print("Models and clients initialized successfully.")
+
+# Initialize models on startup
+initialize_models()
+
+# Store the last generated report for export
+LAST_REPORT = {"content": "", "timestamp": "", "topic": ""}
+
+def format_timestamp():
+    """Get formatted timestamp"""
+    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+def save_report_to_file(report_content, topic):
+    """Save report as markdown file"""
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"research_report_{topic.replace(' ', '_')[:30]}_{timestamp}.md"
     
-    verification_result = run_gemini_text_completion(writer_model, verification_prompt, 0.0)
-    
-    if "OK" in verification_result.upper():
-        return section_text
-    else:
-        return f"{section_text}\n\n---\n*Self-Correction Note: An issue was found during verification. The model suggested the following correction: {verification_result}*"
-
-def get_clarifying_questions(planner_model, initial_topic: str):
-    """Generates clarifying questions for the user."""
-    prompt = prompts.clarification_prompt_template.format(initial_topic=initial_topic)
-    questions = run_gemini_text_completion(planner_model, prompt, 0.5)
-    return questions
-
-def research_and_plan(config: AgentConfig, planner_model, tavily_client, initial_topic: str, user_answers: str):
-    """Constructs the research brief and generates the report outline."""
-    print("\n--- Step 1: Constructing Detailed Research Brief ---")
-    brief_constructor_prompt = prompts.brief_constructor_prompt_template.format(
-        initial_topic=initial_topic,
-        user_answers=user_answers
-    )
-    detailed_topic = run_gemini_text_completion(planner_model, brief_constructor_prompt, config.PLANNER_TEMPERATURE).strip()
+    # Add metadata to the report
+    full_content = f"""---
+title: {topic}
+generated: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
+generator: DeepSearch Research Agent
+---
+
+{report_content}
+"""
     
-    if detailed_topic.startswith("[Error:"):
-        raise ValueError(f"Could not construct a research brief from the provided topic. The LLM returned an error: {detailed_topic}")
+    with open(filename, 'w', encoding='utf-8') as f:
+        f.write(full_content)
     
-    print(f"\n--- Step 2: Performing Broad Initial Research for Outline ---")
-    initial_research = gather_research(tavily_client, [detailed_topic], config.INITIAL_SEARCH_RESULTS)
-    planning_context = "\n\n".join(item['content'] for item in initial_research)
+    return filename
+
+def create_loading_message(stage):
+    """Create animated loading message"""
+    return f"""
+<div style="display: flex; align-items: center; gap: 12px; color: #a78bfa;">
+    <div class="loading-wave">
+        <span></span><span></span><span></span><span></span><span></span>
+    </div>
+    <span style="font-weight: 500;">{stage}</span>
+</div>
+"""
+
+def format_progress_update(message):
+    """Format progress messages with icons"""
+    icons = {
+        "Step": "🎯",
+        "Searching": "🔍",
+        "Found": "✅",
+        "Processing": "⚡",
+        "Writing": "✍️",
+        "Synthesizing": "🧪",
+        "Fact-checking": "🔎",
+        "Indexing": "📚",
+        "Expanding": "🔄"
+    }
     
-    planner_prompt = prompts.planner_prompt.format(topic=detailed_topic, context=planning_context[:20000])
-    plan_response = run_gemini_json_completion(planner_model, planner_prompt, config.PLANNER_TEMPERATURE)
+    for key, icon in icons.items():
+        if key in message:
+            return f'<div class="progress-message">{icon} {message}</div>'
     
-    try:
-        initial_sections = [Section(**s) for s in plan_response.get("sections", [])]
-    except Exception as e:
-        raise ValueError(f"Could not create a valid report plan. Error: {e}")
-    if not initial_sections:
-        raise ValueError("Planner returned no sections.")
+    return f'<div class="progress-message">▶ {message}</div>'
 
-    print("\n--- Step 3: Expanding Outline for Deep Research ---")
-    expanded_sections = []
-    for section in initial_sections:
-        expansion_prompt = prompts.expansion_prompt_template.format(section_title=section.title, section_description=section.description)
-        sub_topics_text = run_gemini_text_completion(planner_model, expansion_prompt, 0.6)
-        section.description += "\n\nKey areas to investigate:\n" + sub_topics_text
-        expanded_sections.append(section)
+def chat_step_wrapper(user_input, history, current_agent_state, topic_state):
+    """Enhanced wrapper with better error handling"""
+    global IS_PROCESSING, LAST_REPORT
     
-    return {"detailed_topic": detailed_topic, "sections": expanded_sections}
+    if IS_PROCESSING:
+        print("Ignoring duplicate request while processing.")
+        if False: yield
+        return
 
-def write_report_stream(config: AgentConfig, writer_model, tavily_client, embedding_model, reranker, plan: dict):
-    """Writes the report section by section, yielding progress updates."""
-    
-    detailed_topic = plan["detailed_topic"]
-    sections = plan["sections"]
+    IS_PROCESSING = True
+    start_time = time.time()
     
-    yield f"### Starting Report Generation for: {detailed_topic}\n\n---\n"
+    try:
+        for update in chat_step(user_input, history, current_agent_state, topic_state):
+            yield update
+    except Exception as e:
+        error_msg = f"""
+<div style="background: #7f1d1d; border: 1px solid #dc2626; padding: 1rem; border-radius: 8px; margin: 1rem 0;">
+    <strong>❌ Error:</strong> {str(e)}
+    <br><br>
+    <small>Please check your API keys and try again.</small>
+</div>
+"""
+        history.append((None, error_msg))
+        yield history, "INITIAL", "", gr.update(interactive=True, placeholder="What would you like to research?"), None, gr.update(visible=False)
+    finally:
+        IS_PROCESSING = False
+        elapsed = time.time() - start_time
+        print(f"Processing completed in {elapsed:.2f} seconds")
+
+def chat_step(user_input, history, current_agent_state, topic_state):
+    """Enhanced chat step with rich visual feedback"""
+    global LAST_REPORT
     
-    report_state = {
-        "full_report_text": f"# Deep Research Report: {detailed_topic}\n\n", 
-        "all_source_urls": set(),
-        "completed_section_titles": []
-    }
-    rag_pipeline = RAGPipeline(embedding_model, reranker)
+    history = history or []
+    history.append((user_input, None))
     
-    for i, section in enumerate(sections):
-        yield f"### Processing Section {i+1}/{len(sections)}: {section.title}...\n\n"
-        
-        previous_sections_context = "\n".join(f"- {title}" for title in report_state["completed_section_titles"])
-        if previous_sections_context:
-            previous_sections_context = "The following sections have already been written:\n" + previous_sections_context
-        else:
-            previous_sections_context = "This is the first section."
-        
-        section_queries = [f"{detailed_topic} - {section.title}"] + section.description.split('\n')[-3:]
-        section_queries = [q.strip() for q in section_queries if q.strip()]
-        section_queries = [q[:400] for q in section_queries]
-        
-        yield f"-> Searching the web for: `{'`, `'.join(section_queries)}`\n"
-        section_research = gather_research(tavily_client, section_queries, config.DEEP_DIVE_SEARCH_RESULTS)
+    if current_agent_state == "INITIAL":
+        # Initial topic analysis
+        yield history, "CLARIFYING", user_input, gr.update(interactive=False, placeholder="Analyzing topic..."), None, gr.update(visible=False)
         
-        if not section_research:
-            section_content = f"## {section.title}\n\nNo research material could be gathered for this section.\n\n"
-            report_state["full_report_text"] += section_content
-            continue
+        history[-1] = (user_input, create_loading_message("Analyzing your research topic"))
+        yield history, "CLARIFYING", user_input, gr.update(interactive=False), None, gr.update(visible=False)
         
-        yield f"-> Found {len(section_research)} sources. Indexing for RAG...\n"
-        rag_pipeline.index_research(section_research)
-        top_chunks_with_meta = rag_pipeline.retrieve_and_rerank(section.description, top_k=config.CHUNKS_TO_USE_FOR_WRITING)
+        time.sleep(0.5)  # Brief pause for effect
         
-        context_for_llm = ""
-        cited_sources_for_section = {}
-        citation_counter = 1
-        for item in top_chunks_with_meta:
-            source_url = item['source']
-            report_state["all_source_urls"].add(source_url)
-            if source_url not in cited_sources_for_section:
-                cited_sources_for_section[source_url] = citation_counter
-                citation_counter += 1
-            citation_num = cited_sources_for_section[source_url]
-            context_for_llm += f"Source [{citation_num}]: {item['content']}\n\n"
+        questions = get_clarifying_questions(planner_model, user_input)
         
-        bibliography = "\n".join(f"[{num}] {url}" for url, num in cited_sources_for_section.items())
-
-        yield f"-> Synthesizing and writing section content...\n"
-        yield f"--> Calling LLM for section '{section.title}'. This may take a moment...\n"
-        writer_prompt = prompts.writer_prompt_template.format(
-            writer_system_instruction=prompts.writer_system_instruction,
-            previous_sections_context=previous_sections_context,
-            section_title=section.title,
-            context_for_llm=context_for_llm
-        )
-        
-        draft_content = run_gemini_text_completion(writer_model, writer_prompt, config.WRITER_TEMPERATURE)
+        formatted_response = f"""
+## 🎯 Research Topic: {user_input}
+
+I'll help you create a comprehensive research report on this topic. To ensure I cover exactly what you're looking for, please answer these clarifying questions:
+
+{questions}
+
+💡 **Tip**: The more detailed your answers, the more tailored and comprehensive your report will be!
+"""
         
-        yield "-> Fact-checking and verifying section...\n\n---\n"
-        final_content = run_verification_step(writer_model, draft_content, context_for_llm)
-        final_content_with_sources = f"## {section.title}\n\n{final_content}\n\n**Sources Used in this Section**\n{bibliography}\n\n"
+        history[-1] = (user_input, formatted_response)
+        yield history, "CLARIFYING", user_input, gr.update(interactive=True, placeholder="Type your answers here..."), None, gr.update(visible=False)
+
+    elif current_agent_state == "CLARIFYING":
+        # Start research process
+        history[-1] = (user_input, "## 🚀 Starting Deep Research Process\n\nYour answers have been recorded. Initiating comprehensive research...")
+        yield history, "GENERATING", topic_state, gr.update(interactive=False, placeholder="Generating report..."), None, gr.update(visible=False)
         
-        report_state["full_report_text"] += final_content_with_sources
-        report_state["completed_section_titles"].append(section.title)
+        try:
+            # Phase 1: Planning
+            history[-1] = (user_input, history[-1][1] + "\n\n" + format_progress_update("Step 1: Creating research plan and outline"))
+            yield history, "GENERATING", topic_state, gr.update(interactive=False), None, gr.update(visible=False)
+            
+            plan = research_and_plan(config, planner_model, tavily_client, topic_state, user_input)
+            
+            # Show plan
+            sections_list = "\n".join([f"   {i+1}. **{s.title}**" for i, s in enumerate(plan['sections'])])
+            plan_display = f"""
+## 🚀 Starting Deep Research Process
 
-    final_bibliography = "\n".join(f"- {url}" for url in sorted(list(report_state["all_source_urls"])))
-    report_state["full_report_text"] += f"## Master Bibliography\n\n{final_bibliography}"
+### 📋 Research Plan Created
+
+**Research Focus**: {plan['detailed_topic']}
+
+**Report Structure**:
+{sections_list}
+
+**Total Sections**: {len(plan['sections'])}
+"""
+            history[-1] = (user_input, plan_display)
+            yield history, "GENERATING", topic_state, gr.update(interactive=False), None, gr.update(visible=False)
+            
+            # Phase 2: Research and Writing
+            report_generator = write_report_stream(config, writer_model, tavily_client, embedding_model, reranker, plan)
+            
+            full_report = ""
+            progress_display = plan_display
+            
+            for update in report_generator:
+                if update.startswith("#"):
+                    # This is the actual report content
+                    full_report = update
+                    LAST_REPORT = {
+                        "content": full_report,
+                        "timestamp": format_timestamp(),
+                        "topic": plan['detailed_topic']
+                    }
+                else:
+                    # This is a progress update
+                    progress_display = plan_display + "\n\n---\n\n### 📊 Current Progress\n\n" + format_progress_update(update)
+                    
+                    # Show preview of report if we have content
+                    if full_report:
+                        preview = full_report[:500] + "..." if len(full_report) > 500 else full_report
+                        progress_display += f"\n\n---\n\n### 📄 Report Preview\n\n{preview}"
+                
+                history[-1] = (user_input, progress_display)
+                yield history, "GENERATING", topic_state, gr.update(interactive=False), None, gr.update(visible=False)
+            
+            # Final report display
+            completion_msg = f"""
+## ✅ Research Complete!
+
+**Topic**: {plan['detailed_topic']}  
+**Generated**: {format_timestamp()}  
+**Sections**: {len(plan['sections'])}  
+
+---
+
+{full_report}
+
+---
+
+### 💾 Export Options
+
+Your report has been saved and is ready for download. Use the export buttons below to:
+- 📄 Download as Markdown file
+- 📋 Copy to clipboard
+- 🔄 Start a new research topic
+"""
+            
+            history[-1] = (user_input, completion_msg)
+            
+            # Save report to file
+            report_file = save_report_to_file(full_report, plan['detailed_topic'])
+            
+            yield history, "INITIAL", "", gr.update(interactive=True, placeholder="What would you like to research next?"), report_file, gr.update(visible=True)
+            
+        except Exception as e:
+            error_display = f"""
+<div style="background: #7f1d1d; border: 1px solid #dc2626; padding: 1.5rem; border-radius: 12px;">
+    <h3 style="color: #fca5a5; margin-top: 0;">❌ Research Error</h3>
+    <p style="color: #fecaca;">{str(e)}</p>
+    <hr style="border-color: #dc2626; margin: 1rem 0;">
+    <p style="color: #fca5a5; font-size: 0.9rem;">Please try again with a different topic or check your API configuration.</p>
+</div>
+"""
+            history.append((None, error_display))
+            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 interface
+with gr.Blocks(css=CSS, theme=gr.themes.Base()) as app:
+    # Header
+    gr.HTML("""
+        <div class="header">
+            <h1>DeepSearch Research Agent</h1>
+            <p>AI-Powered Comprehensive Research & Analysis</p>
+        </div>
+    """)
+    
+    # Status Bar
+    gr.HTML("""
+        <div class="status-bar">
+            <div class="status-indicator">
+                <div class="pulse-dot"></div>
+                <span style="font-weight: 500;">System Active</span>
+            </div>
+            <div style="display: flex; gap: 2rem; align-items: center;">
+                <span style="color: #71717a;">Powered by Gemini Pro</span>
+                <span style="color: #71717a;">Enhanced with Tavily Search</span>
+            </div>
+        </div>
+    """)
+    
+    # State management
+    agent_state = gr.State("INITIAL")
+    initial_topic_state = gr.State("")
+    
+    # Chat interface
+    chatbot = gr.Chatbot(
+        elem_id="chatbot",
+        bubble_full_width=False,
+        height=650,
+        visible=True,
+        value=[(None, """
+## 👋 Welcome to DeepSearch!
+
+I'm your AI research assistant, capable of creating comprehensive, well-researched reports on any topic.
+
+### 🎯 How It Works:
+1. **Choose a Topic** - Tell me what you want to research
+2. **Clarify Details** - I'll ask a few questions to understand your needs
+3. **Deep Research** - I'll search, analyze, and synthesize information
+4. **Get Your Report** - Receive a detailed, sourced report ready for download
+
+### 💡 Example Topics:
+- Impact of AI on healthcare
+- Climate change mitigation strategies  
+- History of quantum computing
+- Economic effects of remote work
+- Future of renewable energy
+
+**Ready to start? Enter your research topic below!**
+""")],
+        avatar_images=(None, "🔬")
+    )
+    
+    # Input section
+    with gr.Group(elem_classes="input-group"):
+        with gr.Row():
+            chat_input = gr.Textbox(
+                placeholder="What would you like to research today?",
+                interactive=True,
+                show_label=False,
+                scale=8,
+                elem_id="user-input"
+            )
+            submit_button = gr.Button(
+                "🚀 Start Research",
+                scale=2,
+                elem_classes="submit-btn"
+            )
+    
+    # Examples section
+    with gr.Group(elem_classes="examples-container"):
+        gr.HTML("<h3>💡 Popular Research Topics</h3>")
+        gr.Examples(
+            examples=[
+                "Impact of artificial intelligence on job markets",
+                "Sustainable urban development strategies",
+                "The psychology of social media addiction",
+                "Blockchain technology in supply chain management",
+                "Future of space exploration and colonization",
+                "Gene editing and CRISPR technology implications"
+            ],
+            inputs=chat_input,
+            label=""
+        )
+    
+    # Export section (hidden initially)
+    with gr.Group(visible=False, elem_classes="export-section") as export_group:
+        gr.HTML("""
+            <h3 style="color: #a78bfa; margin-bottom: 0.5rem;">📥 Download Your Report</h3>
+            <p style="color: #71717a;">Your research report has been generated and saved.</p>
+        """)
+        report_file = gr.File(
+            label="Download Report",
+            visible=True,
+            file_types=[".md"],
+            elem_classes="export-btn"
+        )
+    
+    # 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, report_file, export_group],
+    ).then(
+        fn=lambda: "",
+        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, report_file, export_group],
+    ).then(
+        fn=lambda: "",
+        outputs=[chat_input],
+        queue=False
+    )
 
-    yield report_state["full_report_text"]
+# Launch
+if __name__ == "__main__":
+    app.queue()
+    app.launch(
+        debug=True,
+        share=False,
+        server_name="0.0.0.0",
+        server_port=7860
+    )