Update research_agent/agent.py

research_agent/agent.py

@@ -7,6 +7,9 @@ from datetime import datetime
 import json
 import time
 
+# Import the research agent components
+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")
@@ -17,22 +20,18 @@ if not google_key or not tavily_key:
 # 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%);
@@ -42,7 +41,6 @@ body, .gradio-container {
     text-align: center;
     box-shadow: 0 20px 40px rgba(0,0,0,0.5);
 }
-
 .header h1 {
     font-size: 3.5rem;
     font-weight: 800;
@@ -52,14 +50,12 @@ body, .gradio-container {
     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;
@@ -71,13 +67,11 @@ body, .gradio-container {
     justify-content: space-between;
     align-items: center;
 }
-
 .status-indicator {
     display: flex;
     align-items: center;
     gap: 0.75rem;
 }
-
 .pulse-dot {
     width: 10px;
     height: 10px;
@@ -86,7 +80,6 @@ body, .gradio-container {
     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);
@@ -98,7 +91,6 @@ body, .gradio-container {
         box-shadow: 0 0 0 0 rgba(34, 197, 94, 0);
     }
 }
-
 /* Chat Interface */
 #chatbot {
     background: #18181b !important;
@@ -106,11 +98,9 @@ body, .gradio-container {
     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;
@@ -118,35 +108,29 @@ body, .gradio-container {
     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;
@@ -154,18 +138,15 @@ body, .gradio-container {
     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;
@@ -174,7 +155,6 @@ body, .gradio-container {
     margin: 1rem 0;
     border-radius: 0 8px 8px 0;
 }
-
 /* Input Area */
 .input-group {
     background: #18181b;
@@ -183,7 +163,6 @@ body, .gradio-container {
     padding: 1.5rem;
     margin-top: 2rem;
 }
-
 #user-input textarea {
     background: #27272a !important;
     border: 2px solid #3f3f46 !important;
@@ -193,13 +172,11 @@ body, .gradio-container {
     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;
@@ -212,16 +189,13 @@ body, .gradio-container {
     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;
@@ -230,20 +204,17 @@ body, .gradio-container {
     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;
@@ -252,20 +223,17 @@ body, .gradio-container {
     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;
@@ -273,17 +241,14 @@ body, .gradio-container {
     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;
@@ -292,7 +257,6 @@ body, .gradio-container {
     padding: 1.5rem;
     margin: 1rem 0;
 }
-
 .source-badge {
     display: inline-block;
     background: #374151;
@@ -302,7 +266,6 @@ body, .gradio-container {
     font-size: 0.875rem;
     margin: 0.25rem;
 }
-
 /* Export Section */
 .export-section {
     background: #18181b;
@@ -312,14 +275,12 @@ body, .gradio-container {
     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;
@@ -332,12 +293,10 @@ body, .gradio-container {
     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; }
@@ -503,7 +462,7 @@ I'll help you create a comprehensive research report on this topic. To ensure I
             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'])])
+            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
 
@@ -605,7 +564,7 @@ with gr.Blocks(css=CSS, theme=gr.themes.Base()) as app:
                 <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;">Powered by Gemini Flash</span>
                 <span style="color: #71717a;">Enhanced with Tavily Search</span>
             </div>
         </div>