Update research_agent/agent.py

research_agent/agent.py

@@ -1,676 +1,201 @@
 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
 import time
+from datetime import datetime
+from typing import List, Dict, Any, Generator
+import google.generativeai as genai
 
 
-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; }
-}
-"""
+def get_clarifying_questions(model, topic: str) -> str:
+    """Generate clarifying questions for a research topic"""
+    prompt = f"""
+You are a research assistant. Given the topic "{topic}", generate 3-5 clarifying questions that would help create a more focused and comprehensive research report.
 
+The questions should help understand:
+- Specific aspects the user is most interested in
+- The intended audience or use case  
+- The scope and depth required
+- Any particular angle or perspective needed
+- Timeline or geographical focus if relevant
 
-writer_model, planner_model, embedding_model, reranker, tavily_client = None, None, None, None, None
-IS_PROCESSING = False
+Format your response as a numbered list of questions.
+Be concise but thorough.
 
-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
+Topic: {topic}
+"""
+    
     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')
+        response = model.generate_content(prompt)
+        return response.text
     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")
+        return f"""
+1. What specific aspects of {topic} would you like me to focus on?
+2. Who is the intended audience for this research report?
+3. Are you looking for recent developments, historical context, or both?
+4. What level of detail would be most helpful for your needs?
+5. Are there any particular perspectives or viewpoints you'd like me to include?
+
+Please provide your answers to help me create the most relevant research report for you.
+"""
 
-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"
-    
-    # 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}
+def research_and_plan(config, model, tavily_client, topic: str, clarifications: str) -> Dict[str, Any]:
+    """Create a research plan based on topic and clarifications"""
+    
+    prompt = f"""
+You are a research planner. Create a comprehensive research plan for the topic: "{topic}"
+
+User clarifications: {clarifications}
+
+Please provide your response in the following JSON format:
+{{
+    "detailed_topic": "A refined, specific topic statement based on the clarifications",
+    "sections": [
+        {{"title": "Introduction", "description": "Overview and context setting"}},
+        {{"title": "Background", "description": "Historical context and foundational information"}},
+        {{"title": "Current State", "description": "Present situation and recent developments"}},
+        {{"title": "Key Findings", "description": "Main research findings and data"}},
+        {{"title": "Analysis", "description": "Critical analysis and insights"}},
+        {{"title": "Future Implications", "description": "Trends and future outlook"}},
+        {{"title": "Conclusion", "description": "Summary and key takeaways"}}
+    ],
+    "research_questions": [
+        "What is the current state of {topic}?",
+        "What are the key challenges and opportunities?",
+        "What are the future implications and trends?"
+    ]
+}}
+
+Make sure the JSON is valid and properly formatted.
 """
     
-    with open(filename, 'w', encoding='utf-8') as f:
-        f.write(full_content)
-    
-    return filename
+    try:
+        response = model.generate_content(prompt)
+        response_text = response.text.strip()
+        
+        # Try to extract JSON from the response
+        json_start = response_text.find('{')
+        json_end = response_text.rfind('}') + 1
+        
+        if json_start != -1 and json_end != -1:
+            json_text = response_text[json_start:json_end]
+            plan_data = json.loads(json_text)
+        else:
+            raise ValueError("No valid JSON found in response")
+            
+        return plan_data
+        
+    except Exception as e:
+        print(f"Error in research_and_plan: {str(e)}")
+        # Fallback plan
+        return {
+            "detailed_topic": f"Comprehensive Analysis of {topic}",
+            "sections": [
+                {"title": "Introduction", "description": "Overview and context"},
+                {"title": "Background", "description": "Historical context and current state"},
+                {"title": "Key Findings", "description": "Main research findings and data"},
+                {"title": "Current Trends", "description": "Recent developments and patterns"},
+                {"title": "Analysis", "description": "Critical analysis and insights"},
+                {"title": "Future Outlook", "description": "Predictions and implications"},
+                {"title": "Conclusion", "description": "Summary and recommendations"}
+            ],
+            "research_questions": [
+                f"What is the current state of {topic}?",
+                f"What are the key challenges in {topic}?",
+                f"What are the future implications of {topic}?"
+            ]
+        }
 
-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": "🔄"
-    }
+def write_report_stream(config, model, tavily_client, embedding_model, reranker, plan: Dict[str, Any]) -> Generator[str, None, None]:
+    """Generate a research report section by section with progress updates"""
     
-    for key, icon in icons.items():
-        if key in message:
-            return f'<div class="progress-message">{icon} {message}</div>'
+    sections = plan.get('sections', [])
+    detailed_topic = plan.get('detailed_topic', 'Research Topic')
     
-    return f'<div class="progress-message">▶ {message}</div>'
-
-def chat_step_wrapper(user_input, history, current_agent_state, topic_state):
-    """Enhanced wrapper with better error handling"""
-    global IS_PROCESSING, LAST_REPORT
+    # Start with report header
+    report_content = f"# {detailed_topic}\n\n"
+    report_content += f"*Research Report Generated on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*\n\n"
     
-    if IS_PROCESSING:
-        print("Ignoring duplicate request while processing.")
-        if False: yield
-        return
-
-    IS_PROCESSING = True
-    start_time = time.time()
+    yield "Step 1: Initializing research framework"
+    time.sleep(0.5)
     
+    yield "Step 2: Conducting preliminary search"
+    time.sleep(0.3)
+    
+    # Conduct initial broad search
     try:
-        for update in chat_step(user_input, history, current_agent_state, topic_state):
-            yield update
+        search_results = tavily_client.search(
+            query=detailed_topic,
+            max_results=5,
+            search_depth="advanced"
+        )
+        yield f"Found {len(search_results.get('results', []))} initial sources"
     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
+        yield f"Search warning: {str(e)}"
+        search_results = {"results": []}
     
-    history = history or []
-    history.append((user_input, None))
+    time.sleep(0.3)
     
-    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)
-        
-        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)
+    # Process each section
+    for i, section in enumerate(sections):
+        section_title = section.get('title', f'Section {i+1}')
+        section_desc = section.get('description', '')
         
-        time.sleep(0.5)  # Brief pause for effect
+        yield f"Writing section {i+1}/{len(sections)}: {section_title}"
         
-        questions = get_clarifying_questions(planner_model, user_input)
-        
-        formatted_response = f"""
-## 🎯 Research Topic: {user_input}
+        # Generate content for this section
+        section_prompt = f"""
+Write a comprehensive section titled "{section_title}" for a research report on "{detailed_topic}".
 
-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:
+Section focus: {section_desc}
 
-{questions}
+Requirements:
+- Write 3-5 well-structured paragraphs
+- Include relevant facts, data, or examples
+- Use professional academic writing style
+- Ensure content is accurate and informative
+- Connect this section logically to the overall topic
 
-💡 **Tip**: The more detailed your answers, the more tailored and comprehensive your report will be!
+Write only the section content without the title (it will be added separately).
 """
         
-        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)
-        
         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.get('title', f'Section {i+1}')}**" for i, s in enumerate(plan['sections'])])
-            plan_display = f"""
-## 🚀 Starting Deep Research Process
-
-### 📋 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)
-            
+            section_response = model.generate_content(
+                section_prompt,
+                generation_config=genai.types.GenerationConfig(
+                    temperature=0.4,
+                    max_output_tokens=1000
+                )
+            )
+            section_content = section_response.text.strip()
         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 Flash</span>
-                <span style="color: #71717a;">Enhanced with Tavily Search</span>
-            </div>
-        </div>
-    """)
-    
-    # State management
-    agent_state = gr.State("INITIAL")
-    initial_topic_state = gr.State("")
+            section_content = f"This section would cover {section_desc.lower()}. Further research would be needed for detailed information."
+        
+        # Add section to report
+        report_content += f"## {section_title}\n\n{section_content}\n\n"
+        
+        # Yield current progress with updated report
+        yield report_content
+        time.sleep(0.4)
     
-    # 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, "🔬")
-    )
+    # Add sources section
+    yield "Adding sources and references"
     
-    # 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"
-            )
+    sources_content = "## Sources and References\n\n"
     
-    # 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=""
-        )
+    if search_results.get('results'):
+        for i, result in enumerate(search_results['results'][:5], 1):
+            title = result.get('title', 'Source')
+            url = result.get('url', '#')
+            sources_content += f"{i}. [{title}]({url})\n"
+    else:
+        sources_content += "*Note: This report was generated using AI knowledge. For academic use, please verify with current sources.*\n"
     
-    # 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"
-        )
+    report_content += sources_content
     
-    # 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
-    )
+    yield "Research report completed successfully"
+    time.sleep(0.2)
     
-    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
-    )
-
-# Launch
-if __name__ == "__main__":
-    app.queue()
-    app.launch(
-        debug=True,
-        share=False,
-        server_name="0.0.0.0",
-        server_port=7860
-    )
+    # Final yield with complete report
+    yield report_content