Update app.py

app.py

@@ -123,6 +123,53 @@ def main():
         except Exception as e:
             return None, f"Error retrieving outline data: {str(e)}"
     
+    # Function to fetch research data from AWS
+    def fetch_research_data(outline_id, bucket_name="dediro-backup"):
+        s3_client, message = init_aws_client()
+        if not s3_client:
+            return None, message
+        
+        try:
+            # Construct the S3 key with the correct path for Research Agent
+            s3_key = f"Research-Agent/agent_steps/{outline_id}_results.json"
+            
+            # Debug info
+            log_message = f"Looking for research file at: {bucket_name}/{s3_key}"
+            print(log_message)
+            
+            # Try to get the object from S3
+            try:
+                response = s3_client.get_object(Bucket=bucket_name, Key=s3_key)
+                research_data = json.loads(response['Body'].read().decode('utf-8'))
+                return research_data, f"Successfully retrieved research data for ID: {outline_id}"
+            except s3_client.exceptions.NoSuchKey:
+                # If exact key not found, try to list objects with similar prefix
+                prefix = f"Research-Agent/agent_steps/{outline_id}"
+                response = s3_client.list_objects_v2(
+                    Bucket=bucket_name,
+                    Prefix=prefix
+                )
+                
+                if 'Contents' in response:
+                    # Find the first research JSON file
+                    research_files = [obj['Key'] for obj in response['Contents'] 
+                                    if '_results.json' in obj['Key']]
+                    
+                    if research_files:
+                        # Get the first matching research file
+                        research_key = research_files[0]
+                        print(f"Found alternative file: {research_key}")
+                        response = s3_client.get_object(Bucket=bucket_name, Key=research_key)
+                        research_data = json.loads(response['Body'].read().decode('utf-8'))
+                        return research_data, f"Found similar research results: {research_key}"
+                
+                return None, f"No research data found for ID: {outline_id}"
+                
+        except NoCredentialsError:
+            return None, "AWS credentials not available"
+        except Exception as e:
+            return None, f"Error retrieving research data: {str(e)}"
+    
     # Helper function to extract outline from potentially nested structure
     def extract_outline_from_data(data):
         """Extract outline data from potentially nested structures"""
@@ -142,6 +189,26 @@ def main():
         
         # Return the original if we can't determine the structure
         return data
+    
+    # Helper function to extract research results from potentially nested structure
+    def extract_research_from_data(data):
+        """Extract research data from potentially nested structures"""
+        if not data:
+            return None
+        
+        # Try different possible paths to the research data
+        if isinstance(data, dict):
+            if "results" in data:
+                return data["results"]
+            elif "data" in data and "results" in data["data"]:
+                return data["data"]["results"]
+            
+            # Check if this looks like research results directly
+            if "theme" in data or "articles" in data:
+                return data
+        
+        # Return the original if we can't determine the structure
+        return data
 
     with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange")) as demo:
         # Header
@@ -173,9 +240,34 @@ def main():
                 with gr.Row():
                     load_button = gr.Button("Load Outline")
         
+        # New: Research section
+        with gr.Row():
+            with gr.Column(scale=2):
+                # Research Agent section
+                gr.Markdown("## Generate Research")
+                research_query_input = gr.Textbox(
+                    label="Enter Outline ID for Research:",
+                    placeholder="Enter the Outline ID to generate research for...",
+                    lines=1
+                )
+                
+                with gr.Row():
+                    generate_research_button = gr.Button("Generate Research", variant="primary")
+            
+            with gr.Column(scale=2):
+                # Continue existing research section
+                gr.Markdown("## Continue Existing Research")
+                existing_research_id_input = gr.Textbox(
+                    label="Enter Existing Research ID:",
+                    placeholder="Enter ID to view existing research..."
+                )
+                
+                with gr.Row():
+                    load_research_button = gr.Button("Load Research")
+        
         # Synthesis status display
         with gr.Row(visible=False) as status_panel:
-            gr.Markdown("## Current Outline Status")
+            gr.Markdown("## Current Project Status")
             
             with gr.Column():
                 status_query = gr.Textbox(label="Query")
@@ -191,7 +283,7 @@ def main():
         # Step Status Card
         steps_container = gr.Row(visible=False)
         with steps_container:
-            gr.Markdown("## Outline Status")
+            gr.Markdown("## Process Status")
             
             with gr.Group():
                 gr.Markdown("### Outline Generator")
@@ -200,8 +292,17 @@ def main():
                     outline_retry = gr.Button("Retry", size="sm", variant="primary")
                     outline_view = gr.Button("View Results", size="sm")
                     outline_download = gr.Button("Download", size="sm")
+            
+            # Add Research Agent status section
+            with gr.Group():
+                gr.Markdown("### Research Agent")
+                research_status = gr.Markdown("⏳ **Waiting**...")
+                with gr.Row():
+                    research_retry = gr.Button("Retry", size="sm", variant="primary")
+                    research_view = gr.Button("View Results", size="sm")
+                    research_download = gr.Button("Download", size="sm")
         
-        # Result Preview
+        # Result Preview for Outline
         with gr.Row(visible=False) as result_panel:
             gr.Markdown("## Outline Preview")
             with gr.Column():
@@ -223,6 +324,28 @@ def main():
                     refresh_results = gr.Button("Refresh Results")
                     download_results = gr.Button("Download Results")
         
+        # Result Preview for Research
+        with gr.Row(visible=False) as research_panel:
+            gr.Markdown("## Research Results")
+            with gr.Column():
+                # Add a text area to display any research errors or messages
+                research_message = gr.Textbox(
+                    label="Status",
+                    value="",
+                    visible=False,
+                    lines=2
+                )
+                # JSON component for the research data
+                research_json = gr.JSON(
+                    label="Research Results",
+                    value={}
+                )
+                
+                # Add buttons for refresh and download
+                with gr.Row():
+                    refresh_research = gr.Button("Refresh Research")
+                    download_research_btn = gr.Button("Download Research")
+        
         # Log panel
         with gr.Row(visible=False) as log_panel:
             gr.Markdown("## Process Log")
@@ -236,7 +359,7 @@ def main():
         # Function for generating a new outline
         def generate_new_outline(query):
             if not query.strip():
-                return "Please enter a query first", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), ""
+                return "Please enter a query first", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), "", "", "", "", "", 0.00, "", ""
             
             # Generate a timestamp-based ID (in real implementation, this would match your AWS naming pattern)
             timestamp = str(int(time.time()))
@@ -256,13 +379,31 @@ def main():
                 datetime.now().strftime("%Y-%m-%d %H:%M"),
                 datetime.now().strftime("%Y-%m-%d %H:%M"),
                 0.00,
-                "🟡 **In Progress**..."
+                "🟡 **In Progress**...",
+                "⏳ **Waiting**..."  # Research status starts as waiting
             )
         
-        # Function for loading existing outline - FIX: Return all required values even in error case
+        # Function for generating research
+        def generate_research(outline_id):
+            if not outline_id.strip():
+                return gr.update(visible=True), f"Please enter a valid outline ID", gr.update(visible=False)
+            
+            # First check if the outline exists
+            outline_data, outline_message = fetch_outline_data(outline_id)
+            
+            if not outline_data:
+                return gr.update(visible=True), f"Error: {outline_message}", gr.update(visible=False)
+            
+            # Update UI
+            log_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            log_message = f"[{log_timestamp}] Starting research generation for outline ID: {outline_id}\n[{log_timestamp}] This would trigger the ResearchAgent in a real implementation...\n"
+            
+            return gr.update(visible=True), log_message, gr.update(visible=True)
+        
+        # Function for loading existing outline
         def load_existing(existing_id):
             if not existing_id.strip():
-                # Return all 10 required outputs even in error case
+                # Return all required outputs even in error case
                 return (
                     gr.update(visible=False),  # status_panel
                     gr.update(visible=False),  # steps_container
@@ -273,25 +414,29 @@ def main():
                     "",  # status_created
                     "",  # status_updated
                     0.0,  # status_cost
-                    "❌ **Error**"  # outline_status
+                    "❌ **Error**",  # outline_status
+                    "❌ **Error**"   # research_status (new)
                 )
             
             # Now we actually fetch the data from AWS S3
-            outline_data, message = fetch_outline_data(existing_id)
+            outline_data, outline_message = fetch_outline_data(existing_id)
+            # Also check if research data exists
+            research_data, research_message = fetch_research_data(existing_id)
             
             if not outline_data:
-                # Return all 10 required outputs even in error case
+                # Return all required outputs even in error case
                 return (
                     gr.update(visible=False),  # status_panel
                     gr.update(visible=False),  # steps_container
                     gr.update(visible=False),  # log_panel
-                    f"Error: {message}",  # log_output
+                    f"Error: {outline_message}",  # log_output
                     "",  # status_query
                     existing_id,  # status_id
                     "",  # status_created
                     "",  # status_updated
                     0.0,  # status_cost
-                    "❌ **Error**"  # outline_status
+                    "❌ **Error**",  # outline_status
+                    "❌ **Not Started**"  # research_status
                 )
             
             # Extract the actual outline 
@@ -300,7 +445,7 @@ def main():
             # Extract query from ID (assuming ID format: timestamp_formatted_query)
             parts = existing_id.split('_', 1)
             if len(parts) < 2:
-                # Return all 10 required outputs even in error case
+                # Return all required outputs even in error case
                 return (
                     gr.update(visible=False),  # status_panel
                     gr.update(visible=False),  # steps_container
@@ -311,7 +456,8 @@ def main():
                     "",  # status_created
                     "",  # status_updated
                     0.0,  # status_cost
-                    "❌ **Error**"  # outline_status
+                    "❌ **Error**",  # outline_status
+                    "❌ **Error**"  # research_status
                 )
             
             # Reconstruct the original query
@@ -323,6 +469,14 @@ def main():
             # Determine actual status based on outline data
             outline_status_text = "✅ **Completed**"  # Default to completed if we found data
             
+            # Determine research status
+            if research_data:
+                research_status_text = "✅ **Completed**"
+                log_message = f"[{log_timestamp}] Loaded outline with ID: {existing_id}\n[{log_timestamp}] {outline_message}\n[{log_timestamp}] Research data found: {research_message}\n"
+            else:
+                research_status_text = "⏳ **Not Started**"
+                log_message = f"[{log_timestamp}] Loaded outline with ID: {existing_id}\n[{log_timestamp}] {outline_message}\n[{log_timestamp}] No research data found yet.\n"
+            
             # Get creation time from timestamp in ID
             try:
                 creation_time = datetime.fromtimestamp(int(parts[0])).strftime("%Y-%m-%d %H:%M")
@@ -339,50 +493,35 @@ def main():
                 gr.update(visible=True),
                 gr.update(visible=True),
                 gr.update(visible=True),
-                f"[{log_timestamp}] Loaded outline with ID: {existing_id}\n[{log_timestamp}] {message}\n",
+                log_message,
                 original_query,
                 existing_id,
                 creation_time,
                 datetime.now().strftime("%Y-%m-%d %H:%M"),  # Last updated is now
                 cost,
-                outline_status_text
+                outline_status_text,
+                research_status_text
             )
         
-        # Connect the handlers
-        create_button.click(
-            fn=generate_new_outline,
-            inputs=[query_input],
-            outputs=[
-                synthesis_id, 
-                status_panel, 
-                steps_container, 
-                log_panel, 
-                log_output,
-                status_query,
-                status_id,
-                status_created,
-                status_updated,
-                status_cost,
-                outline_status
-            ]
-        )
-        
-        load_button.click(
-            fn=load_existing,
-            inputs=[existing_id_input],
-            outputs=[
-                status_panel, 
-                steps_container, 
-                log_panel, 
-                log_output,
-                status_query,
-                status_id,
-                status_created,
-                status_updated,
-                status_cost,
-                outline_status
-            ]
-        )
+        # Function for loading existing research
+        def load_existing_research(existing_id):
+            if not existing_id.strip():
+                return gr.update(visible=False), "Please enter a valid research ID", gr.update(visible=True)
+            
+            # Fetch research data from AWS
+            research_data, message = fetch_research_data(existing_id)
+            
+            if not research_data:
+                return gr.update(visible=False), f"Error: {message}", gr.update(visible=True)
+            
+            # Extract the actual research results
+            actual_research = extract_research_from_data(research_data)
+            
+            # Update UI components
+            log_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            log_message = f"[{log_timestamp}] Loaded research with ID: {existing_id}\n[{log_timestamp}] {message}\n"
+            
+            return gr.update(visible=True), actual_research, log_message
         
         # Function for retrying outline generation
         def retry_outline(outline_id):
@@ -394,6 +533,16 @@ def main():
                 f"[{log_timestamp}] Retrying outline generation for ID: {outline_id}\n[{log_timestamp}] Connecting to AWS...\n"
             )
         
+        # Function for retrying research generation
+        def retry_research(outline_id):
+            # This would call your ResearchAgent with force_rerun=True
+            # For now, just update the UI
+            log_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            return (
+                "🟡 **Retrying**...",
+                f"[{log_timestamp}] Retrying research generation for ID: {outline_id}\n[{log_timestamp}] Connecting to AWS...\n"
+            )
+        
         # Function to view outline results
         def view_outline_results(outline_id):
             # Fetch real data from AWS
@@ -458,6 +607,117 @@ def main():
                     "raw_data": str(outline_data)[:1000]
                 }
 
+        # Function to view research results
+        def view_research_results(outline_id):
+            # Fetch real data from AWS
+            research_data, message = fetch_research_data(outline_id)
+            
+            if not research_data:
+                # Return a message if data not found
+                return gr.update(visible=True), gr.update(visible=True), {"error": message}
+            
+            # Print the raw data for debugging
+            print(f"Raw research data type: {type(research_data)}")
+            
+            # Step by step extraction with verbose logging
+            try:
+                # Extract the actual research results
+                actual_research = extract_research_from_data(research_data)
+                
+                # Ensure we're working with a non-empty structure
+                if not actual_research:
+                    print(f"Warning: Extracted research is not valid: {actual_research}")
+                    # If it's not, pass the original data
+                    actual_research = {"warning": "Could not extract proper research structure", "data": str(research_data)[:1000]}
+                
+                # Make the research panel visible and return the data for the JSON component
+                return gr.update(visible=True), gr.update(visible=False), actual_research
+                
+            except Exception as e:
+                # If there's an error processing the data, log it and return the raw data
+                error_msg = f"Error processing research data: {str(e)}"
+                print(error_msg)
+                traceback_info = __import__('traceback').format_exc()
+                print(f"Traceback: {traceback_info}")
+                
+                # Return both the error and the raw data for debugging
+                return gr.update(visible=True), gr.update(visible=True), {
+                    "error": error_msg,
+                    "raw_data": str(research_data)[:1000]
+                }
+        
+        # Function for downloading outline data
+        def download_outline_file(outline_id):
+            # For now, just log that we would download
+            log_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            return f"[{log_timestamp}] Downloading outline with ID: {outline_id}...\n"
+        
+        # Function for downloading research data
+        def download_research_file(outline_id):
+            # For now, just log that we would download
+            log_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            return f"[{log_timestamp}] Downloading research with ID: {outline_id}...\n"
+        
+        # Connect the handlers
+        create_button.click(
+            fn=generate_new_outline,
+            inputs=[query_input],
+            outputs=[
+                synthesis_id, 
+                status_panel, 
+                steps_container, 
+                log_panel, 
+                log_output,
+                status_query,
+                status_id,
+                status_created,
+                status_updated,
+                status_cost,
+                outline_status,
+                research_status
+            ]
+        )
+        
+        load_button.click(
+            fn=load_existing,
+            inputs=[existing_id_input],
+            outputs=[
+                status_panel, 
+                steps_container, 
+                log_panel, 
+                log_output,
+                status_query,
+                status_id,
+                status_created,
+                status_updated,
+                status_cost,
+                outline_status,
+                research_status
+            ]
+        )
+        
+        # Connect research generation button
+        generate_research_button.click(
+            fn=generate_research,
+            inputs=[research_query_input],
+            outputs=[
+                log_panel,
+                log_output,
+                steps_container
+            ]
+        )
+        
+        # Connect load research button
+        load_research_button.click(
+            fn=load_existing_research,
+            inputs=[existing_research_id_input],
+            outputs=[
+                research_panel,
+                research_json,
+                log_output
+            ]
+        )
+        
         # Connect outline action buttons
         outline_retry.click(
             fn=retry_outline,
@@ -468,35 +728,55 @@ def main():
         outline_view.click(
             fn=view_outline_results,
             inputs=[status_id],
-            outputs=[result_panel, result_json]  # Now returning both visibility update and data
+            outputs=[result_panel, result_json]
         )
         
-        # Function for actual download (in production)
-        def download_outline(outline_id):
-            # For now, just log that we would download
-            log_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-            return f"[{log_timestamp}] Downloading outline with ID: {outline_id}...\n"
-        
         outline_download.click(
-            fn=download_outline,
+            fn=download_outline_file,
+            inputs=[status_id],
+            outputs=[log_output]
+        )
+        
+        # Connect research action buttons
+        research_retry.click(
+            fn=retry_research,
+            inputs=[status_id],
+            outputs=[research_status, log_output]
+        )
+        
+        research_view.click(
+            fn=view_research_results,
+            inputs=[status_id],
+            outputs=[research_panel, research_message, research_json]
+        )
+        
+        research_download.click(
+            fn=download_research_file,
             inputs=[status_id],
             outputs=[log_output]
         )
         
-        # Connect refresh button to view function
+        # Connect refresh button to view function for outline
         refresh_results.click(
             fn=view_outline_results,
             inputs=[status_id],
-            outputs=[result_panel, result_json]  # Updated to match the view function outputs
+            outputs=[result_panel, result_json]
         )
         
-        # Connect download results button
+        # Connect download results button for outline
         download_results.click(
-            fn=download_outline,
+            fn=download_outline_file,
             inputs=[status_id],
             outputs=[log_output]
         )
-
+        
+        # Connect refresh button to view function for research
+        refresh_research.click(
+            fn=view_research_results,
+            inputs=[status_id],
+            outputs=[research_panel, research_message, research_json]
+        )
+        
     return demo
 
 if __name__ == "__main__":