Update app.py

app.py

@@ -5,133 +5,219 @@ import json
 from datetime import datetime
 import random
 import string
+import boto3
+from botocore.exceptions import NoCredentialsError
+from dotenv import load_dotenv
+
+# Load environment variables with detailed error handling
+def load_environment_variables():
+    # Try to load from .env file
+    load_dotenv()
+    
+    # Check if critical environment variables exist
+    required_vars = ['AWS_ACCESS_KEY_ID', 'AWS_SECRET_ACCESS_KEY', 'AWS_REGION']
+    missing_vars = [var for var in required_vars if not os.getenv(var)]
+    
+    if missing_vars:
+        print(f"WARNING: Missing required environment variables: {', '.join(missing_vars)}")
+        print("Make sure you have a .env file with these variables or they are set in your environment")
+        print(f"Current environment variable keys: {list(os.environ.keys())}")
+        
+        # Set default values for testing/development (DO NOT USE IN PRODUCTION)
+        if 'AWS_REGION' in missing_vars and not os.getenv('AWS_REGION'):
+            os.environ['AWS_REGION'] = 'us-east-1'
+            print("Set default AWS_REGION to 'us-east-1' for testing purposes")
+    else:
+        print("All required environment variables are set")
+
+# Load environment variables at startup
+load_environment_variables()
 
 def main():
+    # Initialize AWS client with more robust error handling
+    def init_aws_client():
+        aws_access_key_id = os.getenv('AWS_ACCESS_KEY_ID')
+        aws_secret_access_key = os.getenv('AWS_SECRET_ACCESS_KEY')
+        aws_region = os.getenv('AWS_REGION')
+        
+        # Print debug info about environment variables (redacted for security)
+        print(f"AWS_ACCESS_KEY_ID: {'*' * 5 + aws_access_key_id[-4:] if aws_access_key_id else 'Not set'}")
+        print(f"AWS_SECRET_ACCESS_KEY: {'*' * 10 if aws_secret_access_key else 'Not set'}")
+        print(f"AWS_REGION: {aws_region or 'Not set'}")
+        
+        if not all([aws_access_key_id, aws_secret_access_key, aws_region]):
+            missing = []
+            if not aws_access_key_id: missing.append("AWS_ACCESS_KEY_ID")
+            if not aws_secret_access_key: missing.append("AWS_SECRET_ACCESS_KEY")
+            if not aws_region: missing.append("AWS_REGION")
+            
+            return None, f"AWS credentials not found in environment variables: Missing {', '.join(missing)}"
+        
+        try:
+            s3_client = boto3.client(
+                's3',
+                aws_access_key_id=aws_access_key_id,
+                aws_secret_access_key=aws_secret_access_key,
+                region_name=aws_region
+            )
+            
+            # Test the connection with a simple operation
+            try:
+                s3_client.list_buckets()
+                print("Successfully connected to AWS S3")
+            except Exception as e:
+                print(f"Warning: AWS credentials may be invalid: {str(e)}")
+                
+            return s3_client, "AWS client initialized successfully"
+        except Exception as e:
+            return None, f"Error initializing AWS client: {str(e)}"
+    
+    # Function to fetch outline data from AWS
+    def fetch_outline_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
+            s3_key = f"Outline-Agent/agent_steps/{outline_id}_outline.json"
+            
+            # Debug info
+            log_message = f"Looking for 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)
+                outline_data = json.loads(response['Body'].read().decode('utf-8'))
+                return outline_data, f"Successfully retrieved outline data for ID: {outline_id}"
+            except s3_client.exceptions.NoSuchKey:
+                # If exact key not found, try to list objects with similar prefix
+                prefix = f"Outline-Agent/agent_steps/{outline_id}"
+                response = s3_client.list_objects_v2(
+                    Bucket=bucket_name,
+                    Prefix=prefix
+                )
+                
+                if 'Contents' in response:
+                    # Find the first outline JSON file
+                    outline_files = [obj['Key'] for obj in response['Contents'] 
+                                    if '_outline.json' in obj['Key']]
+                    
+                    if outline_files:
+                        # Get the first matching outline file
+                        outline_key = outline_files[0]
+                        print(f"Found alternative file: {outline_key}")
+                        response = s3_client.get_object(Bucket=bucket_name, Key=outline_key)
+                        outline_data = json.loads(response['Body'].read().decode('utf-8'))
+                        return outline_data, f"Found similar outline: {outline_key}"
+                
+                return None, f"No outline data found for ID: {outline_id}"
+                
+        except NoCredentialsError:
+            return None, "AWS credentials not available"
+        except Exception as e:
+            return None, f"Error retrieving outline 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"""
+        if not data:
+            return None
+        
+        # Try different possible paths to the outline data
+        if isinstance(data, dict):
+            if "outline" in data:
+                return data["outline"]
+            elif "data" in data and "outline" in data["data"]:
+                return data["data"]["outline"]
+            
+            # Check if this looks like an outline directly (has title and themes)
+            if "title" in data and "themes" 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
-        gr.Markdown("# AI Synthesis Manager")
+        gr.Markdown("# AI Outline Generator")
         
         # ID-based process handling
         with gr.Row():
             with gr.Column(scale=2):
                 # New synthesis section
-                gr.Markdown("## Create New Synthesis")
+                gr.Markdown("## Create New Outline")
                 query_input = gr.Textbox(
                     label="Enter Query:",
-                    placeholder="What topic would you like to generate a synthesis for?",
+                    placeholder="What topic would you like to generate an outline for?",
                     lines=2
                 )
                 
                 with gr.Row():
-                    create_button = gr.Button("Generate New Synthesis ID", variant="primary")
-                    synthesis_id = gr.Textbox(label="Synthesis ID", placeholder="ID will appear here...")
+                    create_button = gr.Button("Generate New Outline", variant="primary")
+                    synthesis_id = gr.Textbox(label="Outline ID", placeholder="ID will appear here...")
             
             with gr.Column(scale=2):
                 # Continue existing synthesis section
-                gr.Markdown("## Continue Existing Synthesis")
+                gr.Markdown("## Continue Existing Outline")
                 existing_id_input = gr.Textbox(
-                    label="Enter Existing Synthesis ID:",
+                    label="Enter Existing Outline ID:",
                     placeholder="Enter ID to continue from a previous session..."
                 )
                 
                 with gr.Row():
-                    load_button = gr.Button("Load Synthesis State")
+                    load_button = gr.Button("Load Outline")
         
         # Synthesis status display
         with gr.Row(visible=False) as status_panel:
-            gr.Markdown("## Current Synthesis Status")
+            gr.Markdown("## Current Outline Status")
             
             with gr.Column():
                 status_query = gr.Textbox(label="Query")
-                status_id = gr.Textbox(label="Synthesis ID")
+                status_id = gr.Textbox(label="Outline ID")
                 
                 with gr.Row():
                     status_created = gr.Textbox(label="Created", scale=1)
                     status_updated = gr.Textbox(label="Last Updated", scale=1)
                 
                 with gr.Row():
-                    status_progress = gr.Slider(
-                        minimum=0, maximum=8, step=1, value=3,
-                        label="Progress (Steps Completed)"
-                    )
-                    status_cost = gr.Number(label="Cost So Far ($)", value=0.00)
+                    status_cost = gr.Number(label="Cost ($)", value=0.00)
         
-        # Step Status Cards
+        # Step Status Card
         steps_container = gr.Row(visible=False)
         with steps_container:
-            gr.Markdown("## Pipeline Steps")
-            
-            with gr.Column():
-                # Step 1: Preoutline
-                with gr.Group():
-                    gr.Markdown("### 1. Preoutline")
-                    step1_status = gr.Markdown("✅ **Completed**")
-                    with gr.Row():
-                        step1_view = gr.Button("View Results", size="sm")
-                        step1_download = gr.Button("Download", size="sm")
-                
-                # Step 2: Themes Generator
-                with gr.Group():
-                    gr.Markdown("### 2. Themes Generator")
-                    step2_status = gr.Markdown("✅ **Completed**")
-                    with gr.Row():
-                        step2_view = gr.Button("View Results", size="sm")
-                        step2_download = gr.Button("Download", size="sm")
-                
-                # Step 3: Outline
-                with gr.Group():
-                    gr.Markdown("### 3. Outline")
-                    step3_status = gr.Markdown("❌ **Failed** (Connection timeout)")
-                    with gr.Row():
-                        step3_retry = gr.Button("Retry Step", size="sm", variant="primary")
-                        step3_view = gr.Button("View Partial Results", size="sm")
-                        step3_download = gr.Button("Download Partial", size="sm")
-                
-                # Step 4: Research
-                with gr.Group():
-                    gr.Markdown("### 4. Research")
-                    step4_status = gr.Markdown("⏳ **Waiting** (Depends on Outline)")
-                    with gr.Row():
-                        step4_force = gr.Button("Force Start", size="sm")
+            gr.Markdown("## Outline Status")
             
-            with gr.Column():
-                # Step 5: Thinking Ladder
-                with gr.Group():
-                    gr.Markdown("### 5. Thinking Ladder")
-                    step5_status = gr.Markdown("⏳ **Waiting** (Depends on Research)")
-                    with gr.Row():
-                        step5_force = gr.Button("Force Start", size="sm")
-                
-                # Step 6: Enhanced Outline
-                with gr.Group():
-                    gr.Markdown("### 6. Enhanced Outline")
-                    step6_status = gr.Markdown("⏳ **Waiting** (Depends on Thinking Ladder)")
-                    with gr.Row():
-                        step6_force = gr.Button("Force Start", size="sm")
-                
-                # Step 7: Synthesis
-                with gr.Group():
-                    gr.Markdown("### 7. Synthesis")
-                    step7_status = gr.Markdown("⏳ **Waiting** (Depends on Enhanced Outline)")
-                    with gr.Row():
-                        step7_force = gr.Button("Force Start", size="sm")
-                
-                # Step 8: Accountability
-                with gr.Group():
-                    gr.Markdown("### 8. Accountability")
-                    step8_status = gr.Markdown("⏳ **Waiting** (Depends on Synthesis)")
-                    with gr.Row():
-                        step8_force = gr.Button("Force Start", size="sm")
+            with gr.Group():
+                gr.Markdown("### Outline Generator")
+                outline_status = gr.Markdown("⏳ **In Progress**...")
+                with gr.Row():
+                    outline_retry = gr.Button("Retry", size="sm", variant="primary")
+                    outline_view = gr.Button("View Results", size="sm")
+                    outline_download = gr.Button("Download", size="sm")
         
         # Result Preview
         with gr.Row(visible=False) as result_panel:
-            gr.Markdown("## Results Preview")
-            # Removed 'interactive=False' as it's not a valid parameter for gr.JSON
-            result_json = gr.JSON(label="Synthesis Results")
-            
-            with gr.Row():
-                refresh_results = gr.Button("Refresh Results")
-                download_final = gr.Button("Download Final Results")
+            gr.Markdown("## Outline Preview")
+            with gr.Column():
+                # Add a text area to display any errors or messages
+                result_message = gr.Textbox(
+                    label="Status",
+                    value="",
+                    visible=False,
+                    lines=2
+                )
+                # Keep the JSON component for the outline data
+                result_json = gr.JSON(
+                    label="Outline Results",
+                    value={}
+                )
+                
+                # Add buttons for refresh and download
+                with gr.Row():
+                    refresh_results = gr.Button("Refresh Results")
+                    download_results = gr.Button("Download Results")
         
         # Log panel
         with gr.Row(visible=False) as log_panel:
@@ -143,56 +229,88 @@ def main():
                 placeholder="Process logs will appear here..."
             )
         
-        # Mock function for generating a new synthesis ID
-        def generate_new_id(query):
+        # 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), ""
             
-            # Generate a random ID (in real implementation, this would be more robust)
-            synthesis_id = ''.join(random.choices(string.ascii_uppercase + string.digits, k=8))
+            # Generate a timestamp-based ID (in real implementation, this would match your AWS naming pattern)
+            timestamp = str(int(time.time()))
+            formatted_query = query.replace(" ", "_")
+            synthesis_id = f"{timestamp}_{formatted_query}"
             
             # Update UI
-            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            log_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
             return (
                 synthesis_id,
                 gr.update(visible=True),
                 gr.update(visible=True), 
                 gr.update(visible=True),
-                f"[{timestamp}] Created new synthesis with ID: {synthesis_id}\n[{timestamp}] Starting preoutline generation for query: {query}...\n",
+                f"[{log_timestamp}] Created new outline with ID: {synthesis_id}\n[{log_timestamp}] Starting outline generation for query: {query}...\n",
                 query,
                 synthesis_id,
                 datetime.now().strftime("%Y-%m-%d %H:%M"),
                 datetime.now().strftime("%Y-%m-%d %H:%M"),
-                0.00
+                0.00,
+                "🟡 **In Progress**..."
             )
         
-        # Mock function for loading existing synthesis
+        # Function for loading existing outline
         def load_existing(existing_id):
             if not existing_id.strip():
-                return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), "Please enter a valid synthesis ID"
+                return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), "Please enter a valid outline ID"
+            
+            # Now we actually fetch the data from AWS S3
+            outline_data, message = fetch_outline_data(existing_id)
+            
+            if not outline_data:
+                return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), f"Error: {message}"
+            
+            # Extract the actual outline 
+            actual_outline = extract_outline_from_data(outline_data)
             
-            # In a real implementation, this would fetch data from AWS based on the ID
-            # For demonstration, we'll use mock data
+            # Extract query from ID (assuming ID format: timestamp_formatted_query)
+            parts = existing_id.split('_', 1)
+            if len(parts) < 2:
+                return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), "Invalid outline ID format"
+            
+            # Reconstruct the original query
+            original_query = parts[1].replace("_", " ")
             
             # Update UI components
-            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            log_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            
+            # Determine actual status based on outline data
+            outline_status_text = "✅ **Completed**"  # Default to completed if we found data
+            
+            # Get creation time from timestamp in ID
+            try:
+                creation_time = datetime.fromtimestamp(int(parts[0])).strftime("%Y-%m-%d %H:%M")
+            except:
+                creation_time = "Unknown"
+            
+            # Extract cost if available in the data
+            cost = 0.00
+            if isinstance(outline_data, dict):
+                if "metrics" in outline_data:
+                    cost = outline_data.get("metrics", {}).get("cost", 0.00)
             
             return (
                 gr.update(visible=True),
                 gr.update(visible=True),
                 gr.update(visible=True),
-                f"[{timestamp}] Loaded synthesis with ID: {existing_id}\n[{timestamp}] Last step completed: Themes Generator\n[{timestamp}] Next step: Outline (previously failed)\n",
-                "Sustainability in urban planning",
+                f"[{log_timestamp}] Loaded outline with ID: {existing_id}\n[{log_timestamp}] {message}\n",
+                original_query,
                 existing_id,
-                "2025-04-25 14:30",
-                "2025-04-25 14:45",
-                2,
-                0.85
+                creation_time,
+                datetime.now().strftime("%Y-%m-%d %H:%M"),  # Last updated is now
+                cost,
+                outline_status_text
             )
         
         # Connect the handlers
         create_button.click(
-            fn=generate_new_id,
+            fn=generate_new_outline,
             inputs=[query_input],
             outputs=[
                 synthesis_id, 
@@ -204,7 +322,8 @@ def main():
                 status_id,
                 status_created,
                 status_updated,
-                status_cost
+                status_cost,
+                outline_status
             ]
         )
         
@@ -220,81 +339,138 @@ def main():
                 status_id,
                 status_created,
                 status_updated,
-                status_progress,
-                status_cost
+                status_cost,
+                outline_status
             ]
         )
         
-        # Mock function for step retry
-        def retry_step(id):
-            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        # Function for retrying outline generation
+        def retry_outline(outline_id):
+            # This would call your OutlineAgent with force_rerun=True
+            # For now, just update the UI
+            log_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
             return (
-                "🟡 **In Progress**...",
-                f"[{timestamp}] Retrying Outline step for synthesis ID: {id}\n[{timestamp}] Connecting to AWS...\n"
+                "🟡 **Retrying**...",
+                f"[{log_timestamp}] Retrying outline generation for ID: {outline_id}\n[{log_timestamp}] Connecting to AWS...\n"
             )
         
-        # Connect retry button
-        step3_retry.click(
-            fn=retry_step,
+        # Function to view outline results
+        def view_outline_results(outline_id):
+            # Fetch real data from AWS
+            outline_data, message = fetch_outline_data(outline_id)
+            
+            if not outline_data:
+                # Return a message if data not found
+                return gr.update(visible=True), {"error": message}
+            
+            # Print the raw data for debugging
+            print(f"Raw outline data: {json.dumps(outline_data, indent=2)[:1000]}...")
+            
+            # Step by step extraction with verbose logging
+            try:
+                # First, check if "outline" is directly in the data
+                if "outline" in outline_data:
+                    print("Found 'outline' key at top level")
+                    actual_outline = outline_data["outline"]
+                # Then check if it's in a nested "data" object
+                elif "data" in outline_data and "outline" in outline_data["data"]:
+                    print("Found 'outline' key inside 'data'")
+                    actual_outline = outline_data["data"]["outline"]
+                # Next, try the most common pattern seen in OutlineAgent
+                elif "models" in outline_data and "prompts" in outline_data and "query" in outline_data:
+                    print("Found OutlineAgent format with 'models', 'prompts', 'query'")
+                    # In this case, we want the actual outline in the data structure
+                    if isinstance(outline_data.get("outline", None), dict):
+                        actual_outline = outline_data["outline"]
+                        print("Using 'outline' field in OutlineAgent format")
+                    else:
+                        # Just use the whole data
+                        actual_outline = outline_data
+                        print("Using complete OutlineAgent data")
+                # If none of those work, check if this looks like an outline directly
+                elif "title" in outline_data and "themes" in outline_data:
+                    print("Found direct outline structure with 'title' and 'themes'")
+                    actual_outline = outline_data
+                else:
+                    # As a last resort, just use the whole data
+                    print("Could not identify specific outline structure, using complete data")
+                    actual_outline = outline_data
+                
+                # Ensure we're working with a non-empty dictionary
+                if not actual_outline or not isinstance(actual_outline, dict):
+                    print(f"Warning: Extracted outline is not a valid dictionary: {actual_outline}")
+                    # If it's not, pass the original data
+                    actual_outline = {"warning": "Could not extract proper outline structure", "data": str(outline_data)[:1000]}
+                
+                # Log what we found
+                print(f"Final outline structure: {json.dumps(actual_outline, indent=2)[:500]}...")
+                
+                # Make the result panel visible and return the data for the JSON component
+                return gr.update(visible=True), actual_outline
+                
+            except Exception as e:
+                # If there's an error processing the data, log it and return the raw data
+                error_msg = f"Error processing outline 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), {
+                    "error": error_msg,
+                    "raw_data": str(outline_data)[:1000]
+                }
+
+        # Connect outline action buttons
+        outline_retry.click(
+            fn=retry_outline,
             inputs=[status_id],
-            outputs=[step3_status, log_output]
+            outputs=[outline_status, log_output]
         )
         
-        # Mock function for force start
-        def force_start(step_name, id):
-            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-            return (
-                "🟡 **Force Started**...",
-                f"[{timestamp}] Force starting {step_name} step for synthesis ID: {id}\n[{timestamp}] Warning: This step depends on previous incomplete steps\n"
-            )
-        
-        # Connect force start buttons
-        step4_force.click(
-            fn=lambda id: force_start("Research", id),
+        outline_view.click(
+            fn=view_outline_results,
             inputs=[status_id],
-            outputs=[step4_status, log_output]
+            outputs=[result_panel, result_json]  # Now returning both visibility update and data
         )
         
-        step5_force.click(
-            fn=lambda id: force_start("Thinking Ladder", id),
+        # 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,
             inputs=[status_id],
-            outputs=[step5_status, log_output]
+            outputs=[log_output]
         )
         
-        # Mock view results for steps
-        def view_step_results(step_name, id):
-            mock_results = {
-                "Preoutline": {
-                    "main_points": ["Sustainable urban development", "Green infrastructure", "Smart city technologies"],
-                    "key_concepts": ["Energy efficiency", "Waste management", "Public transportation"],
-                    "potential_structure": "Introduction → Current challenges → Sustainable solutions → Case studies → Conclusion"
-                },
-                "Themes": {
-                    "themes": ["Environmental Impact", "Economic Considerations", "Social Aspects", "Technological Integration"],
-                    "query": "Sustainability in urban planning",
-                    "focus_areas": ["Renewable energy in cities", "Waste reduction strategies", "Community involvement"]
-                }
-            }
-            
-            if step_name in mock_results:
-                return gr.update(visible=True, value=mock_results[step_name])
-            return gr.update(visible=True, value={"message": f"No results available for {step_name} step"})
-        
-        # Connect view buttons
-        step1_view.click(
-            fn=lambda: view_step_results("Preoutline", ""),
-            inputs=[],
-            outputs=[result_panel]
+        # Connect refresh button to view function
+        refresh_results.click(
+            fn=view_outline_results,
+            inputs=[status_id],
+            outputs=[result_panel, result_json]  # Updated to match the view function outputs
         )
         
-        step2_view.click(
-            fn=lambda: view_step_results("Themes", ""),
-            inputs=[],
-            outputs=[result_panel]
+        # Connect download results button
+        download_results.click(
+            fn=download_outline,
+            inputs=[status_id],
+            outputs=[log_output]
         )
 
     return demo
 
 if __name__ == "__main__":
+    # Check for .env file existence and provide guidance if not found
+    env_file_path = '.env'
+    if not os.path.exists(env_file_path):
+        print("\n" + "="*80)
+        print("WARNING: .env file not found!")
+
+        print("="*80 + "\n")
+    
+    # Initialize the demo
     demo = main()
     demo.launch(debug=True)