Update app.py

app.py

@@ -4,150 +4,372 @@ import gradio as gr
 import google.generativeai as genai
 from PIL import Image
 import json
+import asyncio
+import threading
+from typing import Optional, List, Dict, Any
+import time
+from functools import lru_cache
 
 # Load environment variables
 load_dotenv()
 
-# ### 1. Configure the AI Model
-def setup_api_key():
-    """Set up the API key for the generative AI model."""
-    try:
-        api_key = os.environ.get('GEMINI_API_KEY')
-        if not api_key:
-            raise ValueError("GEMINI_API_KEY not found in environment variables.")
-        genai.configure(api_key=api_key)
-        print("🔑 API Key configured successfully.")
-        return True
-    except Exception as e:
-        print(f"🔴 Error during API configuration: {e}")
-        return False
+# ### 1. Enhanced Configuration with Error Handling
+class APIManager:
+    """Centralized API management with connection pooling and error handling."""
+    
+    def __init__(self):
+        self.api_key = None
+        self.is_configured = False
+        self.models = {}
+        self.setup_api()
+    
+    def setup_api(self):
+        """Enhanced API setup with better error handling."""
+        try:
+            self.api_key = os.environ.get('GEMINI_API_KEY')
+            if not self.api_key:
+                raise ValueError("GEMINI_API_KEY not found in environment variables.")
+            
+            genai.configure(api_key=self.api_key)
+            self.is_configured = True
+            print("🔑 API Key configured successfully.")
+            
+            # Pre-initialize models for better performance
+            self.initialize_models()
+            
+        except Exception as e:
+            print(f"🔴 Error during API configuration: {e}")
+            self.is_configured = False
+    
+    def initialize_models(self):
+        """Pre-initialize models for better performance."""
+        try:
+            self.models = {
+                'vision': genai.GenerativeModel('gemini-2.5-pro', 
+                                              system_instruction=VISION_SYSTEM_INSTRUCTION),
+                'initial': genai.GenerativeModel('gemini-2.5-flash-lite-preview-06-17', 
+                                               system_instruction=PROMPT_ENGINEER_SYSTEM_INSTRUCTION),
+                'refiner': genai.GenerativeModel('gemini-2.5-pro', 
+                                               system_instruction=PROMPT_REFINER_SYSTEM_INSTRUCTION),
+                'rewriter': genai.GenerativeModel('gemini-2.5-flash', 
+                                                system_instruction=META_PROMPT_SYSTEM_INSTRUCTION)
+            }
+        except Exception as e:
+            print(f"⚠️ Warning: Could not pre-initialize models: {e}")
+
+# Global API manager instance
+api_manager = APIManager()
+
+# ### 2. Enhanced System Instructions with Better Prompting
+VISION_SYSTEM_INSTRUCTION = """
+You are an expert UI/UX and content analyst with deep knowledge of design patterns and user experience principles.
 
-API_KEY_IS_SET = setup_api_key()
+TASK: Analyze the provided screenshot or content and extract actionable insights.
 
-# ### 2. Model Constants
-MODEL_VISION = 'gemini-2.5-pro'
-MODEL_INITIAL = 'gemini-2.5-flash'
-MODEL_REFINER = 'gemini-2.5-pro'
-MODEL_REWRITER = 'gemini-2.5-pro'
+ANALYSIS FRAMEWORK:
+1. **Visual Hierarchy**: Identify key elements, layout structure, and information architecture
+2. **Content Strategy**: Analyze text, messaging, and content organization
+3. **User Flow**: Understand the intended user journey and interaction patterns
+4. **Design Patterns**: Recognize UI components, design systems, and visual elements
+5. **Context Clues**: Infer purpose, target audience, and business objectives
 
-# ### 3. System Instructions
-VISION_SYSTEM_INSTRUCTION = """
-You are an expert UI and content analyst. Analyze the provided screenshot and describe its key elements, including layout, text, and notable features. If no screenshot is provided, use the textual description as context. Provide a concise summary suitable for further processing, even if the input is minimal.
+OUTPUT FORMAT:
+- Be specific and actionable
+- Focus on elements that would inform prompt creation
+- Highlight pain points or opportunities for improvement
+- Use clear, professional language
+- Keep analysis concise but comprehensive (200-400 words)
+
+If no screenshot is provided, analyze the textual description with the same framework.
 """
 
 PROMPT_ENGINEER_SYSTEM_INSTRUCTION = """
-You are a meticulous AI Prompt Engineer. Based on the provided analysis and goal, craft a clear, specific, and high-quality prompt optimized for instructing an AI model. If the analysis or goal is minimal, infer a reasonable task and create a concise, actionable prompt that maximizes clarity and effectiveness.
+You are a world-class AI Prompt Engineer specializing in creating high-performance prompts that maximize AI model capabilities.
+
+EXPERTISE AREAS:
+- Prompt architecture and optimization
+- AI model behavior and limitations
+- Task decomposition and instruction design
+- Output formatting and constraints
+- Performance optimization techniques
+
+PROMPT CREATION PRINCIPLES:
+1. **Clarity**: Use precise, unambiguous language
+2. **Structure**: Organize instructions logically with clear sections
+3. **Specificity**: Define exact requirements and expectations
+4. **Context**: Provide necessary background information
+5. **Constraints**: Set appropriate boundaries and limitations
+6. **Examples**: Include relevant examples when helpful
+7. **Output Format**: Specify desired response structure
+
+TASK: Create a professional, optimized prompt based on the provided analysis and goal.
+
+OUTPUT REQUIREMENTS:
+- Start with a clear role definition for the AI
+- Include specific task instructions
+- Add relevant context and constraints
+- Specify output format if needed
+- Optimize for the intended AI model's capabilities
+- Ensure the prompt is actionable and measurable
+
+Length: 100-300 words (unless complexity requires more)
 """
 
 PROMPT_REFINER_SYSTEM_INSTRUCTION = """
-You are an AI Prompt Refiner. Given an original prompt and user feedback, refine the prompt to address the feedback while maintaining its original intent. Output a single, improved prompt that is clear, concise, and actionable.
+You are an expert AI Prompt Optimization specialist focused on iterative improvement.
+
+TASK: Refine the given prompt based on user feedback while preserving core intent.
+
+REFINEMENT APPROACH:
+1. **Analyze Feedback**: Understand specific improvement requests
+2. **Identify Issues**: Pinpoint areas needing enhancement
+3. **Preserve Intent**: Maintain original purpose and goals
+4. **Enhance Clarity**: Improve language and structure
+5. **Optimize Performance**: Adjust for better AI model response
+
+REFINEMENT TECHNIQUES:
+- Add missing context or constraints
+- Improve instruction clarity
+- Enhance output specifications
+- Adjust tone and style
+- Optimize prompt structure
+- Add or modify examples
+
+OUTPUT: A single, improved prompt that addresses the feedback while maintaining effectiveness.
 """
 
 META_PROMPT_SYSTEM_INSTRUCTION = """
-You are an AI Prompt Optimization expert. Take a single input prompt and generate 3 distinct, improved variations that are clear, concise, and optimized for instructing an AI model. Each variation should enhance clarity, specificity, and effectiveness while maintaining the original intent.
-## CRITICAL OUTPUT FORMAT
-Your output MUST be a valid JSON array of 3 strings, each a unique prompt variation. For example:
-["Variation 1", "Variation 2", "Variation 3"]
+You are an AI Prompt Optimization expert specializing in creating multiple high-quality variations.
+
+TASK: Generate 3 distinct, improved variations of the input prompt.
+
+VARIATION STRATEGY:
+1. **Variation 1**: Enhanced clarity and structure
+2. **Variation 2**: Different approach or perspective
+3. **Variation 3**: Optimized for specific use case or context
+
+OPTIMIZATION TECHNIQUES:
+- Reframe instructions for better comprehension
+- Adjust specificity levels
+- Modify tone and style
+- Enhance context and examples
+- Improve output formatting
+- Optimize for different AI model strengths
+
+CRITICAL OUTPUT FORMAT:
+Return ONLY a valid JSON array of exactly 3 strings:
+["Variation 1 text here", "Variation 2 text here", "Variation 3 text here"]
+
+Each variation should be complete, standalone, and ready to use.
 """
 
-# ### 4. AI Processing Functions
+# ### 3. Enhanced Processing Functions with Better Error Handling
+@lru_cache(maxsize=128)
+def analyze_screenshot_cached(image_hash: str, pil_image: Image.Image) -> str:
+    """Cached screenshot analysis for better performance."""
+    return analyze_screenshot(pil_image)
+
 def analyze_screenshot(pil_image: Image.Image) -> str:
-    """Analyze a screenshot and return a textual description."""
+    """Enhanced screenshot analysis with better error handling."""
     if not isinstance(pil_image, Image.Image):
         return "Error: Invalid image provided."
-    model = genai.GenerativeModel(MODEL_VISION, system_instruction=VISION_SYSTEM_INSTRUCTION)
+    
+    if not api_manager.is_configured:
+        return "Error: API not configured. Please check your API key."
+    
     try:
-        response = model.generate_content(pil_image)
-        return response.text.strip() or "No meaningful content detected in the screenshot."
+        model = api_manager.models.get('vision') or genai.GenerativeModel(
+            'gemini-2.0-flash-exp', 
+            system_instruction=VISION_SYSTEM_INSTRUCTION
+        )
+        
+        response = model.generate_content([
+            "Please analyze this screenshot following the framework provided in your system instructions.",
+            pil_image
+        ])
+        
+        result = response.text.strip()
+        return result if result else "No meaningful content detected in the screenshot."
+        
     except Exception as e:
-        return f"Error in vision analysis: {e}"
+        error_msg = f"Error in vision analysis: {str(e)}"
+        print(error_msg)
+        return error_msg
 
-def initial_prompt_stream(analysis_text, goal):
-    """Generate an initial prompt based on analysis and goal, streaming the output."""
-    model = genai.GenerativeModel(MODEL_INITIAL, system_instruction=PROMPT_ENGINEER_SYSTEM_INSTRUCTION)
-    prompt = f"**Analysis:**\n{analysis_text}\n\n**Goal:**\n{goal or 'Infer from analysis.'}"
-    final_prompt_full = ""
+def initial_prompt_stream(analysis_text: str, goal: str):
+    """Enhanced streaming prompt generation with better error handling."""
+    if not api_manager.is_configured:
+        yield "Error: API not configured. Please check your API key."
+        return
+    
     try:
+        model = api_manager.models.get('initial') or genai.GenerativeModel(
+            'gemini-2.0-flash-exp', 
+            system_instruction=PROMPT_ENGINEER_SYSTEM_INSTRUCTION
+        )
+        
+        # Enhanced prompt construction
+        prompt_parts = [
+            "**CONTEXT ANALYSIS:**",
+            analysis_text,
+            "\n**USER GOAL:**",
+            goal or "Create an optimized prompt based on the provided analysis.",
+            "\n**TASK:** Create a professional, optimized prompt that addresses the context and achieves the specified goal."
+        ]
+        
+        prompt = "\n".join(prompt_parts)
+        final_prompt_full = ""
+        
         for chunk in model.generate_content(prompt, stream=True):
-            final_prompt_full += chunk.text
-            yield final_prompt_full.strip()
+            if chunk.text:
+                final_prompt_full += chunk.text
+                yield final_prompt_full.strip()
+        
         if not final_prompt_full.strip():
-            final_prompt_full = "Generated prompt: Create a general-purpose AI instruction based on minimal context."
-        yield final_prompt_full
+            fallback = "Create a comprehensive prompt based on the provided context to achieve optimal AI performance."
+            yield fallback
+            
     except Exception as e:
-        yield f"Error in prompt generation: {e}"
+        error_msg = f"Error in prompt generation: {str(e)}"
+        print(error_msg)
+        yield error_msg
 
-def refinement_prompt_stream(original_prompt, feedback):
-    """Refine a prompt based on user feedback, streaming the output."""
-    model = genai.GenerativeModel(MODEL_REFINER, system_instruction=PROMPT_REFINER_SYSTEM_INSTRUCTION)
-    prompt = f"**Original Prompt:**\n{original_prompt}\n\n**Feedback:**\n{feedback}"
-    final_prompt_full = ""
+def refinement_prompt_stream(original_prompt: str, feedback: str):
+    """Enhanced prompt refinement with better streaming."""
+    if not api_manager.is_configured:
+        yield "Error: API not configured. Please check your API key."
+        return
+    
     try:
-        for chunk in model.generate_content(prompt, stream=True):
-            final_prompt_full += chunk.text
-            yield final_prompt_full.strip()
+        model = api_manager.models.get('refiner') or genai.GenerativeModel(
+            'gemini-2.0-flash-exp', 
+            system_instruction=PROMPT_REFINER_SYSTEM_INSTRUCTION
+        )
+        
+        refinement_prompt = f"""
+**ORIGINAL PROMPT:**
+{original_prompt}
+
+**USER FEEDBACK:**
+{feedback}
+
+**TASK:** Refine the original prompt based on the feedback while preserving its core intent and improving its effectiveness.
+"""
+        
+        final_prompt_full = ""
+        for chunk in model.generate_content(refinement_prompt, stream=True):
+            if chunk.text:
+                final_prompt_full += chunk.text
+                yield final_prompt_full.strip()
+        
         if not final_prompt_full.strip():
-            final_prompt_full = original_prompt  # Fallback to original if refinement fails
-        yield final_prompt_full
+            yield original_prompt  # Fallback to original
+            
     except Exception as e:
-        yield f"Error in prompt refinement: {e}"
+        error_msg = f"Error in prompt refinement: {str(e)}"
+        print(error_msg)
+        yield error_msg
 
-def rewrite_prompt_with_prewrite(original_prompt):
-    """Generate three improved variations of a prompt."""
-    model = genai.GenerativeModel(MODEL_REWRITER, system_instruction=META_PROMPT_SYSTEM_INSTRUCTION)
+def rewrite_prompt_with_prewrite(original_prompt: str) -> List[str]:
+    """Enhanced prompt rewriting with better JSON parsing."""
+    if not api_manager.is_configured:
+        return ["Error: API not configured. Please check your API key.", "", ""]
+    
     try:
-        response = model.generate_content(original_prompt)
-        cleaned_response = response.text.strip().replace("```json", "").replace("```", "")
-        variations = json.loads(cleaned_response)
+        model = api_manager.models.get('rewriter') or genai.GenerativeModel(
+            'gemini-2.0-flash-exp', 
+            system_instruction=META_PROMPT_SYSTEM_INSTRUCTION
+        )
+        
+        rewrite_prompt = f"""
+Generate 3 improved variations of this prompt:
+
+{original_prompt}
+
+Remember: Output ONLY a valid JSON array of 3 strings.
+"""
+        
+        response = model.generate_content(rewrite_prompt)
+        
+        # Enhanced JSON parsing
+        response_text = response.text.strip()
+        
+        # Clean up common formatting issues
+        response_text = response_text.replace("```json", "").replace("```", "").strip()
+        
+        # Try to extract JSON if it's wrapped in other text
+        if not response_text.startswith('['):
+            import re
+            json_match = re.search(r'\[.*\]', response_text, re.DOTALL)
+            if json_match:
+                response_text = json_match.group(0)
+        
+        variations = json.loads(response_text)
+        
         if isinstance(variations, list) and len(variations) >= 1:
-            return variations + [""] * (3 - len(variations))  # Pad with empty strings if needed
+            # Ensure we have exactly 3 variations
+            while len(variations) < 3:
+                variations.append("")
+            return variations[:3]
+        
         return ["Error: AI returned an invalid format.", "", ""]
+        
+    except json.JSONDecodeError:
+        return ["Error: Could not parse AI response as JSON.", "", ""]
     except Exception as e:
-        return [f"Error: {e}", "The AI's response was not valid JSON.", ""]
+        error_msg = f"Error in prompt rewriting: {str(e)}"
+        print(error_msg)
+        return [error_msg, "", ""]
 
-# ### 5. Gradio App Logic
-def run_analysis_step(pil_image, situation_text):
-    """Analyze the input context (screenshot or text) and update the UI."""
+# ### 4. Enhanced Gradio Interface Functions
+def run_analysis_step(pil_image: Optional[Image.Image], situation_text: str):
+    """Enhanced analysis step with better state management."""
+    # Reset UI state
     yield {
         satisfaction_row: gr.update(visible=False),
         feedback_col: gr.update(visible=False),
-        prewrite_col: gr.update(visible=False)
+        prewrite_col: gr.update(visible=False),
+        analysis_accordion: gr.update(open=True),
+        analysis_output: "🔍 Analyzing context...",
+        final_prompt_output: ""
     }
-
-    if not API_KEY_IS_SET:
-        return {
+    
+    # Validation
+    if not api_manager.is_configured:
+        yield {
             analysis_accordion: gr.update(open=True),
-            analysis_output: "Error: API Key not set.",
+            analysis_output: "❌ Error: API Key not configured. Please check your GEMINI_API_KEY environment variable.",
             final_prompt_output: "",
             satisfaction_row: gr.update(visible=False),
             feedback_col: gr.update(visible=False),
             prewrite_col: gr.update(visible=False),
             analysis_state: None
         }
-
+        return
+    
     if pil_image is None and not situation_text.strip():
-        return {
+        yield {
             analysis_accordion: gr.update(open=True),
-            analysis_output: "Error: Please provide a screenshot or a situation description.",
+            analysis_output: "⚠️ Please provide either a screenshot or a situation description to proceed.",
             final_prompt_output: "",
             satisfaction_row: gr.update(visible=False),
             feedback_col: gr.update(visible=False),
             prewrite_col: gr.update(visible=False),
             analysis_state: None
         }
-
-    yield {
-        analysis_accordion: gr.update(open=True),
-        analysis_output: "Processing context...",
-        final_prompt_output: "",
-    }
-
+        return
+    
+    # Perform analysis
     try:
-        analysis_text = analyze_screenshot(pil_image) if pil_image else situation_text.strip()
-        print("Generated Analysis Text:", analysis_text)
-        if not analysis_text:
-            analysis_text = "No detailed analysis available; using raw input as context."
-        return {
+        if pil_image:
+            analysis_text = analyze_screenshot(pil_image)
+        else:
+            analysis_text = situation_text.strip()
+        
+        if not analysis_text or analysis_text.startswith("Error"):
+            analysis_text = analysis_text or "Unable to generate analysis. Please try again."
+        
+        yield {
             analysis_accordion: gr.update(open=True),
             analysis_output: analysis_text,
             final_prompt_output: "",
@@ -156,10 +378,11 @@ def run_analysis_step(pil_image, situation_text):
             prewrite_col: gr.update(visible=False),
             analysis_state: analysis_text
         }
+        
     except Exception as e:
-        error_msg = f"Error during analysis: {str(e)}"
+        error_msg = f"❌ Error during analysis: {str(e)}"
         print(error_msg)
-        return {
+        yield {
             analysis_accordion: gr.update(open=True),
             analysis_output: error_msg,
             final_prompt_output: "",
@@ -169,39 +392,70 @@ def run_analysis_step(pil_image, situation_text):
             analysis_state: None
         }
 
-def run_streaming_generation(analysis, goal):
-    """Generate and stream a prompt based on the analysis and goal."""
-    print("Analysis State in Streaming:", analysis)
+def run_streaming_generation(analysis: str, goal: str):
+    """Enhanced streaming generation with better progress indication."""
     if not analysis:
         yield {
-            final_prompt_output: "Error: No analysis available for prompt generation.",
+            final_prompt_output: "❌ Error: No analysis available for prompt generation.",
             first_prompt_state: None,
             satisfaction_row: gr.update(visible=False)
         }
         return
-    yield {final_prompt_output: "Generating prompt..."}
+    
+    yield {
+        final_prompt_output: "🚀 Generating optimized prompt...",
+        satisfaction_row: gr.update(visible=False)
+    }
+    
     final_prompt_full = ""
     for chunk in initial_prompt_stream(analysis, goal):
         final_prompt_full = chunk
         yield {final_prompt_output: final_prompt_full}
+    
     yield {
         final_prompt_output: final_prompt_full,
         first_prompt_state: final_prompt_full,
         satisfaction_row: gr.update(visible=True)
     }
 
-def handle_auto_refine(original_prompt):
-    """Generate and display three prompt variations for user selection."""
+def handle_auto_refine(original_prompt: str):
+    """Enhanced auto-refinement with better user feedback."""
+    if not original_prompt:
+        return {
+            prewrite_col: gr.update(visible=False),
+            satisfaction_row: gr.update(visible=True),
+            feedback_col: gr.update(visible=False)
+        }
+    
     variations = rewrite_prompt_with_prewrite(original_prompt)
+    
+    # Filter out empty variations
+    valid_variations = [v for v in variations if v.strip()]
+    
+    if not valid_variations:
+        return {
+            prewrite_col: gr.update(visible=False),
+            satisfaction_row: gr.update(visible=True),
+            feedback_col: gr.update(visible=False)
+        }
+    
     return {
         prewrite_col: gr.update(visible=True),
-        prewrite_choices: gr.update(choices=variations, value=variations[0] if variations else ""),
+        prewrite_choices: gr.update(choices=valid_variations, value=valid_variations[0]),
         satisfaction_row: gr.update(visible=False),
         feedback_col: gr.update(visible=False)
     }
 
-def select_rewritten_prompt(selected_prompt):
-    """Update the final prompt with the selected variation."""
+def select_rewritten_prompt(selected_prompt: str):
+    """Enhanced prompt selection with validation."""
+    if not selected_prompt or not selected_prompt.strip():
+        return {
+            final_prompt_output: "❌ Error: No prompt selected.",
+            first_prompt_state: None,
+            satisfaction_row: gr.update(visible=False),
+            prewrite_col: gr.update(visible=False)
+        }
+    
     return {
         final_prompt_output: selected_prompt,
         first_prompt_state: selected_prompt,
@@ -210,7 +464,7 @@ def select_rewritten_prompt(selected_prompt):
     }
 
 def handle_manual_feedback():
-    """Show the feedback input area."""
+    """Show feedback input area."""
     return {
         feedback_col: gr.update(visible=True),
         satisfaction_row: gr.update(visible=False),
@@ -218,23 +472,29 @@ def handle_manual_feedback():
     }
 
 def handle_like():
-    """Hide all refinement options when the user is satisfied."""
+    """Hide refinement options when user is satisfied."""
     return {
         satisfaction_row: gr.update(visible=False),
         feedback_col: gr.update(visible=False),
         prewrite_col: gr.update(visible=False)
     }
 
-def refine_with_manual_feedback(original_prompt, feedback):
-    """Refine the prompt based on manual feedback."""
+def refine_with_manual_feedback(original_prompt: str, feedback: str):
+    """Enhanced manual refinement with better streaming."""
     if not feedback.strip():
         yield {
             final_prompt_output: original_prompt,
             first_prompt_state: original_prompt,
-            satisfaction_row: gr.update(visible=True)
+            satisfaction_row: gr.update(visible=True),
+            feedback_col: gr.update(visible=False)
         }
         return
-    yield {final_prompt_output: "Refining manually..."}
+    
+    yield {
+        final_prompt_output: "🛠️ Refining prompt based on your feedback...",
+        satisfaction_row: gr.update(visible=False)
+    }
+    
     final_prompt_full = ""
     for chunk in refinement_prompt_stream(original_prompt, feedback):
         final_prompt_full = chunk
@@ -242,10 +502,14 @@ def refine_with_manual_feedback(original_prompt, feedback):
             final_prompt_output: final_prompt_full,
             first_prompt_state: final_prompt_full
         }
-    yield {satisfaction_row: gr.update(visible=True)}
+    
+    yield {
+        satisfaction_row: gr.update(visible=True),
+        feedback_col: gr.update(visible=False)
+    }
 
 def clear_all():
-    """Reset the entire interface to its initial state."""
+    """Enhanced reset function with complete state clearing."""
     return {
         image_input: None,
         situation_input: "",
@@ -257,100 +521,296 @@ def clear_all():
         feedback_col: gr.update(visible=False),
         prewrite_col: gr.update(visible=False),
         prewrite_choices: gr.update(choices=[], value=None),
+        feedback_input: "",
         analysis_state: None,
         first_prompt_state: None
     }
 
-# ### 6. Build and Launch Gradio Interface
-theme = gr.themes.Base(
-    primary_hue=gr.themes.colors.blue,
-    secondary_hue=gr.themes.colors.neutral,
-    font=(gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif")
-).set(
-    button_primary_background_fill_hover='*primary_600',
-    button_primary_border_color_hover='*primary_600'
-)
-
-with gr.Blocks(theme=theme, analytics_enabled=False) as demo:
-    # State variables
-    analysis_state = gr.State()
-    first_prompt_state = gr.State()
-
-    # Header
-    gr.Markdown("# 🤖 Context-Aware Prompt Engineer AI (v22 - Revised)")
-
-    # Layout
-    with gr.Row(variant="panel"):
-        with gr.Column(scale=1):
-            with gr.Group():
-                gr.Markdown("### 1. Provide Context")
-                image_input = gr.Image(type="pil", label="Upload Screenshot")
-                gr.HTML("<div style='text-align: center; font-weight: 500;'>OR</div>")
-                situation_input = gr.Textbox(label="Describe Your Situation", lines=6)
-            with gr.Accordion("2. (Optional) Clarify Goal", open=False):
-                goal_input = gr.Textbox(label=None, show_label=False)
-            with gr.Row():
-                submit_btn = gr.Button("🚀 Generate Prompt", variant="primary", scale=3)
-                clear_btn = gr.Button("🔄 Start Over", scale=1)
-        with gr.Column(scale=2):
-            with gr.Accordion("Show AI's Context Analysis", open=False) as analysis_accordion:
-                analysis_output = gr.Textbox(label=None, lines=7, interactive=False, show_label=False, value="")
-            final_prompt_output = gr.Textbox(label="✅ Final Engineered Prompt", lines=12, interactive=False, show_copy_button=True)
-
-    # Refinement UI
-    with gr.Row(visible=False) as satisfaction_row:
-        gr.Markdown("#### Are you satisfied, or would you like to refine the prompt?")
-        like_btn = gr.Button("👍 Yes, Looks Good")
-        auto_refine_btn = gr.Button("🤖 Auto-Refine (PRewrite)")
-        dislike_btn = gr.Button("👎 Manual Feedback")
-
-    with gr.Column(visible=False) as prewrite_col:
-        gr.Markdown("### PRewrite: AI-Generated Variations")
-        prewrite_choices = gr.Radio(label="Select the best variation:", type="value")
-        select_version_btn = gr.Button("✅ Use This Version", variant="primary")
-
-    with gr.Column(visible=False) as feedback_col:
-        feedback_input = gr.Textbox(label="What would you like to change or improve?", lines=3)
-        refine_btn = gr.Button("🛠️ Refine Manually", variant="primary")
-
-    # Event Handlers
-    analysis_outputs = [analysis_accordion, analysis_output, final_prompt_output, satisfaction_row, feedback_col, prewrite_col, analysis_state]
-    streaming_outputs = [final_prompt_output, first_prompt_state, satisfaction_row]
-
-    submit_btn.click(
-        fn=run_analysis_step,
-        inputs=[image_input, situation_input],
-        outputs=analysis_outputs,
-        show_progress="hidden"
-    ).then(
-        fn=run_streaming_generation,
-        inputs=[analysis_state, goal_input],
-        outputs=streaming_outputs,
-        show_progress="hidden"
-    )
-
-    goal_input.submit(
-        fn=run_analysis_step,
-        inputs=[image_input, situation_input],
-        outputs=analysis_outputs,
-        show_progress="hidden"
-    ).then(
-        fn=run_streaming_generation,
-        inputs=[analysis_state, goal_input],
-        outputs=streaming_outputs,
-        show_progress="hidden"
-    )
-
-    like_btn.click(fn=handle_like, outputs=[satisfaction_row, feedback_col, prewrite_col])
-    auto_refine_btn.click(fn=handle_auto_refine, inputs=[first_prompt_state], outputs=[prewrite_col, prewrite_choices, satisfaction_row, feedback_col])
-    dislike_btn.click(fn=handle_manual_feedback, outputs=[feedback_col, satisfaction_row, prewrite_col])
-    select_version_btn.click(fn=select_rewritten_prompt, inputs=[prewrite_choices], outputs=[final_prompt_output, first_prompt_state, satisfaction_row, prewrite_col])
-    refine_btn.click(fn=refine_with_manual_feedback, inputs=[first_prompt_state, feedback_input], outputs=[final_prompt_output, first_prompt_state, satisfaction_row])
-    feedback_input.submit(fn=refine_with_manual_feedback, inputs=[first_prompt_state, feedback_input], outputs=[final_prompt_output, first_prompt_state, satisfaction_row])
-
-    clear_btn.click(
-        fn=clear_all,
-        outputs=[image_input, situation_input, goal_input, analysis_accordion, analysis_output, final_prompt_output, satisfaction_row, feedback_col, prewrite_col, prewrite_choices, analysis_state, first_prompt_state]
+# ### 5. Enhanced UI with Modern Design
+def create_enhanced_interface():
+    """Create the enhanced Gradio interface."""
+    
+    # Custom CSS for better styling
+    custom_css = """
+    .gradio-container {
+        max-width: 1200px !important;
+        margin: auto;
+    }
+    
+    .generate-btn {
+        background: linear-gradient(45deg, #007bff, #0056b3) !important;
+        border: none !important;
+        color: white !important;
+        font-weight: 600 !important;
+        transition: all 0.3s ease !important;
+    }
+    
+    .generate-btn:hover {
+        transform: translateY(-2px) !important;
+        box-shadow: 0 4px 12px rgba(0, 123, 255, 0.3) !important;
+    }
+    
+    .status-indicator {
+        padding: 8px 16px;
+        border-radius: 20px;
+        font-size: 14px;
+        font-weight: 500;
+        margin: 8px 0;
+    }
+    
+    .status-success {
+        background: #d4edda;
+        color: #155724;
+        border: 1px solid #c3e6cb;
+    }
+    
+    .status-error {
+        background: #f8d7da;
+        color: #721c24;
+        border: 1px solid #f5c6cb;
+    }
+    
+    .status-processing {
+        background: #d1ecf1;
+        color: #0c5460;
+        border: 1px solid #bee5eb;
+    }
+    """
+    
+    # Enhanced theme
+    theme = gr.themes.Soft(
+        primary_hue=gr.themes.colors.blue,
+        secondary_hue=gr.themes.colors.neutral,
+        neutral_hue=gr.themes.colors.slate,
+        font=(gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif")
+    ).set(
+        button_primary_background_fill='*primary_500',
+        button_primary_background_fill_hover='*primary_600',
+        button_primary_border_color='*primary_500',
+        button_primary_border_color_hover='*primary_600',
+        button_secondary_background_fill='*neutral_100',
+        button_secondary_text_color='*neutral_800'
     )
-
-demo.launch(debug=True, pwa=True)
+    
+    with gr.Blocks(theme=theme, css=custom_css, title="AI Prompt Engineer") as interface:
+        # State management
+        analysis_state = gr.State()
+        first_prompt_state = gr.State()
+        
+        # Header
+        gr.Markdown("""
+        # 🚀 AI Prompt Engineer Pro
+        ### Transform your ideas into optimized AI prompts with intelligent analysis and refinement
+        """)
+        
+        # Status indicator
+        with gr.Row():
+            with gr.Column():
+                if api_manager.is_configured:
+                    gr.HTML('<div class="status-indicator status-success">✅ API Connected & Ready</div>')
+                else:
+                    gr.HTML('<div class="status-indicator status-error">❌ API Configuration Error</div>')
+        
+        # Main interface
+        with gr.Row(equal_height=True):
+            # Input column
+            with gr.Column(scale=1, min_width=400):
+                with gr.Group():
+                    gr.Markdown("### 📝 Provide Context")
+                    
+                    with gr.Tabs():
+                        with gr.Tab("📸 Screenshot"):
+                            image_input = gr.Image(
+                                type="pil", 
+                                label="Upload Screenshot",
+                                sources=['upload'],
+                                interactive=True
+                            )
+                        
+                        with gr.Tab("✍️ Text Description"):
+                            situation_input = gr.Textbox(
+                                label="Describe Your Situation",
+                                placeholder="Describe what you need help with, the context, or the task you want to accomplish...",
+                                lines=6,
+                                max_lines=10
+                            )
+                
+                with gr.Accordion("🎯 Optional: Specify Goal", open=False):
+                    goal_input = gr.Textbox(
+                        label="What do you want to achieve?",
+                        placeholder="e.g., 'Create a prompt for generating marketing copy' or 'Help me write better documentation'",
+                        lines=2,
+                        show_label=False
+                    )
+                
+                with gr.Row():
+                    submit_btn = gr.Button(
+                        "🚀 Generate Prompt",
+                        variant="primary",
+                        scale=3,
+                        elem_classes=["generate-btn"]
+                    )
+                    clear_btn = gr.Button("🔄 Reset", scale=1)
+            
+            # Output column
+            with gr.Column(scale=2, min_width=600):
+                with gr.Accordion("🔍 AI Analysis", open=False) as analysis_accordion:
+                    analysis_output = gr.Textbox(
+                        label="Context Analysis",
+                        lines=6,
+                        interactive=False,
+                        show_copy_button=True,
+                        placeholder="AI analysis will appear here..."
+                    )
+                
+                final_prompt_output = gr.Textbox(
+                    label="✅ Optimized Prompt",
+                    lines=12,
+                    interactive=False,
+                    show_copy_button=True,
+                    placeholder="Your optimized prompt will appear here..."
+                )
+        
+        # Refinement interface
+        with gr.Row(visible=False) as satisfaction_row:
+            with gr.Column():
+                gr.Markdown("### 🎨 Refinement Options")
+                gr.Markdown("*How would you like to improve this prompt?*")
+                
+                with gr.Row():
+                    like_btn = gr.Button("👍 Perfect!", variant="secondary")
+                    auto_refine_btn = gr.Button("🤖 Auto-Refine", variant="primary")
+                    dislike_btn = gr.Button("✏️ Manual Feedback", variant="secondary")
+        
+        # Auto-refinement section
+        with gr.Column(visible=False) as prewrite_col:
+            gr.Markdown("### 🔄 Auto-Generated Variations")
+            gr.Markdown("*Select the variation that best fits your needs:*")
+            
+            prewrite_choices = gr.Radio(
+                label="Choose your preferred version:",
+                type="value",
+                interactive=True
+            )
+            select_version_btn = gr.Button("✅ Use This Version", variant="primary")
+        
+        # Manual feedback section
+        with gr.Column(visible=False) as feedback_col:
+            gr.Markdown("### 💬 Manual Refinement")
+            feedback_input = gr.Textbox(
+                label="What would you like to improve?",
+                placeholder="e.g., 'Make it more specific', 'Add examples', 'Change the tone to be more professional'...",
+                lines=3
+            )
+            refine_btn = gr.Button("🛠️ Refine Prompt", variant="primary")
+        
+        # Event handlers with better error handling
+        def safe_wrapper(func):
+            """Wrapper for safe function execution."""
+            def wrapper(*args, **kwargs):
+                try:
+                    return func(*args, **kwargs)
+                except Exception as e:
+                    print(f"Error in {func.__name__}: {e}")
+                    return {"error": str(e)}
+            return wrapper
+        
+        # Main workflow
+        analysis_outputs = [
+            analysis_accordion, analysis_output, final_prompt_output,
+            satisfaction_row, feedback_col, prewrite_col, analysis_state
+        ]
+        
+        streaming_outputs = [final_prompt_output, first_prompt_state, satisfaction_row]
+        
+        # Event bindings
+        submit_btn.click(
+            fn=run_analysis_step,
+            inputs=[image_input, situation_input],
+            outputs=analysis_outputs,
+            show_progress="minimal"
+        ).then(
+            fn=run_streaming_generation,
+            inputs=[analysis_state, goal_input],
+            outputs=streaming_outputs,
+            show_progress="minimal"
+        )
+        
+        # Auto-submission on goal input
+        goal_input.submit(
+            fn=run_analysis_step,
+            inputs=[image_input, situation_input],
+            outputs=analysis_outputs,
+            show_progress="minimal"
+        ).then(
+            fn=run_streaming_generation,
+            inputs=[analysis_state, goal_input],
+            outputs=streaming_outputs,
+            show_progress="minimal"
+        )
+        
+        # Refinement handlers
+        like_btn.click(
+            fn=handle_like,
+            outputs=[satisfaction_row, feedback_col, prewrite_col]
+        )
+        
+        auto_refine_btn.click(
+            fn=handle_auto_refine,
+            inputs=[first_prompt_state],
+            outputs=[prewrite_col, prewrite_choices, satisfaction_row, feedback_col]
+        )
+        
+        dislike_btn.click(
+            fn=handle_manual_feedback,
+            outputs=[feedback_col, satisfaction_row, prewrite_col]
+        )
+        
+        select_version_btn.click(
+            fn=select_rewritten_prompt,
+            inputs=[prewrite_choices],
+            outputs=[final_prompt_output, first_prompt_state, satisfaction_row, prewrite_col]
+        )
+        
+        refine_btn.click(
+            fn=refine_with_manual_feedback,
+            inputs=[first_prompt_state, feedback_input],
+            outputs=[final_prompt_output, first_prompt_state, satisfaction_row, feedback_col]
+        )
+        
+        feedback_input.submit(
+            fn=refine_with_manual_feedback,
+            inputs=[first_prompt_state, feedback_input],
+            outputs=[final_prompt_output, first_prompt_state, satisfaction_row, feedback_col]
+        )
+        
+        # Reset functionality
+        clear_btn.click(
+            fn=clear_all,
+            outputs=[
+                image_input, situation_input, goal_input, analysis_accordion,
+                analysis_output, final_prompt_output, satisfaction_row,
+                feedback_col, prewrite_col, prewrite_choices, feedback_input,
+                analysis_state, first_prompt_state
+            ]
+        )
+        
+        return interface
+
+# ### 6. Launch Configuration
+if __name__ == "__main__":
+    # Create and launch the enhanced interface
+    demo = create_enhanced_interface()
+    
+    # Launch with optimal settings
+    demo.launch(
+        debug=True,
+        share=False,
+        inbrowser=True,
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        favicon_path=None,
+        ssl_verify=False,
+        quiet=False
+    )