Update app.py

app.py

@@ -1,5 +1,10 @@
 """
-Marketing Image Generator with AI Review 
+Marketing Image Generator - Agent Workflow Test Version
+
+This version clearly shows the Agent1 -> Agent2 workflow:
+- Agent1: Image Generator (calls Google Imagen)
+- Agent2: Image Reviewer (calls Google Gemini Vision)
+- Clear logging of each step
 """
 
 import gradio as gr
@@ -11,240 +16,367 @@ import io
 import os
 import hashlib
 import logging
-import google.generativeai as genai
-from google import genai as genai_sdk
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Configuration
+# Google AI setup
+google_auth_configured = False
 GOOGLE_SERVICE_ACCOUNT_JSON = os.getenv("GOOGLE_SERVICE_ACCOUNT_JSON", "")
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")
 
-# Initialize Google AI
-google_auth_configured = False
-
-if GOOGLE_SERVICE_ACCOUNT_JSON:
-    try:
+# Try to configure Google AI
+try:
+    import google.generativeai as genai
+    from google import genai as genai_sdk
+    
+    if GOOGLE_SERVICE_ACCOUNT_JSON:
         credentials_dict = json.loads(GOOGLE_SERVICE_ACCOUNT_JSON)
         from google.oauth2 import service_account
         credentials = service_account.Credentials.from_service_account_info(credentials_dict)
         genai.configure(credentials=credentials)
         google_auth_configured = True
         logger.info("✅ Google Cloud configured")
-    except Exception as e:
-        logger.error(f"Google Cloud setup failed: {e}")
-
-if not google_auth_configured and GOOGLE_API_KEY:
-    try:
+    elif GOOGLE_API_KEY:
         genai.configure(api_key=GOOGLE_API_KEY)
         google_auth_configured = True
-        logger.info("✅ Google API key configured")
-    except Exception as e:
-        logger.error(f"Google API setup failed: {e}")
-
-def create_placeholder_image(prompt, style, aspect_ratio):
-    """Create a clean placeholder image"""
-    # Set dimensions
-    if aspect_ratio == "Square (1:1)":
-        width, height = 512, 512
-    elif aspect_ratio == "Portrait (3:4)":
-        width, height = 512, 683
-    else:  # Landscape (16:9)
-        width, height = 683, 384
-    
-    # Generate color from prompt
-    prompt_hash = int(hashlib.md5(prompt.encode()).hexdigest()[:8], 16)
-    
-    # Clean color palette
-    colors = {
-        "Professional": "#3B82F6",  # Blue
-        "Creative": "#8B5CF6",      # Purple  
-        "Minimalist": "#6B7280",    # Gray
-        "Corporate": "#1E40AF",     # Dark Blue
-        "Modern": "#059669"         # Green
-    }
-    
-    base_color = colors.get(style, "#3B82F6")
-    
-    # Convert hex to RGB
-    hex_color = base_color.lstrip('#')
-    rgb = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
-    
-    # Create image
-    img = Image.new('RGB', (width, height), rgb)
-    draw = ImageDraw.Draw(img)
-    
-    # Add text
-    try:
-        font = ImageFont.load_default()
-        
-        # Main text
-        text = "Marketing Image"
-        bbox = draw.textbbox((0, 0), text, font=font)
-        text_width = bbox[2] - bbox[0]
-        text_height = bbox[3] - bbox[1]
-        x = (width - text_width) // 2
-        y = (height - text_height) // 2 - 20
-        
-        draw.text((x, y), text, fill="white", font=font)
-        
-        # Style text
-        style_text = f"Style: {style}"
-        bbox2 = draw.textbbox((0, 0), style_text, font=font)
-        text_width2 = bbox2[2] - bbox2[0]
-        x2 = (width - text_width2) // 2
-        draw.text((x2, y + 30), style_text, fill="white", font=font)
-        
-    except:
-        pass
-    
-    return img
-
-def generate_with_google_imagen(prompt, style):
-    """Try to generate with real Google Imagen"""
-    if not google_auth_configured:
-        return None
-        
-    try:
-        # Enhance prompt
-        enhanced_prompt = f"{prompt}, {style.lower()} style, high quality, professional marketing image, 4K"
-        
-        if GOOGLE_SERVICE_ACCOUNT_JSON:
-            client = genai_sdk.Client()
-        else:
-            client = genai_sdk.Client(api_key=GOOGLE_API_KEY)
-        
-        result = client.models.generate_images(
-            model="imagen-3.0-generate-002",
-            prompt=enhanced_prompt,
-            config={
-                "number_of_images": 1,
-                "output_mime_type": "image/png"
-            }
-        )
-        
-        if result and hasattr(result, 'generated_images') and len(result.generated_images) > 0:
-            generated_image = result.generated_images[0]
-            if hasattr(generated_image, 'image') and hasattr(generated_image.image, 'image_bytes'):
-                image_bytes = generated_image.image.image_bytes
-                return Image.open(io.BytesIO(image_bytes))
+        logger.info("✅ Google API configured")
+except Exception as e:
+    logger.error(f"Google setup failed: {e}")
+    genai = None
+
+# ====== AGENT 1: IMAGE GENERATOR ======
+
+class ImageGeneratorAgent:
+    """Agent 1: Responsible for generating images using Google Imagen"""
+    
+    def __init__(self):
+        self.name = "ImageGeneratorAgent"
+        self.status = "Ready"
+        
+    def enhance_prompt(self, prompt, style):
+        """Enhance user prompt for better results"""
+        logger.info(f"🤖 Agent1: Enhancing prompt for {style} style")
+        
+        style_enhancers = {
+            "Professional": "professional photography, corporate setting, high quality, clean composition, marketing ready",
+            "Creative": "creative composition, artistic flair, innovative design, vibrant colors, eye-catching",
+            "Minimalist": "clean minimal design, simple composition, white space, elegant, professional",
+            "Corporate": "business professional, corporate branding, trustworthy, authoritative, polished",
+            "Modern": "contemporary design, sleek, cutting-edge, stylish, trendy"
+        }
+        
+        enhancer = style_enhancers.get(style, "high quality, professional")
+        enhanced = f"{prompt}, {enhancer}, 4K resolution, sharp focus, marketing quality"
+        
+        logger.info(f"🤖 Agent1: Enhanced prompt: {enhanced[:100]}...")
+        return enhanced
+    
+    def generate_image(self, prompt, style, aspect_ratio):
+        """Generate image using Google Imagen"""
+        logger.info(f"🤖 Agent1: Starting image generation")
+        logger.info(f"🤖 Agent1: Original prompt: {prompt}")
+        logger.info(f"🤖 Agent1: Style: {style}, Format: {aspect_ratio}")
+        
+        # Step 1: Enhance prompt
+        enhanced_prompt = self.enhance_prompt(prompt, style)
+        
+        # Step 2: Try Google Imagen
+        if google_auth_configured and genai:
+            try:
+                logger.info(f"🤖 Agent1: Calling Google Imagen API...")
                 
-    except Exception as e:
-        logger.warning(f"Google Imagen failed: {e}")
+                if GOOGLE_SERVICE_ACCOUNT_JSON:
+                    client = genai_sdk.Client()
+                else:
+                    client = genai_sdk.Client(api_key=GOOGLE_API_KEY)
+                
+                result = client.models.generate_images(
+                    model="imagen-3.0-generate-002",
+                    prompt=enhanced_prompt,
+                    config={
+                        "number_of_images": 1,
+                        "output_mime_type": "image/png"
+                    }
+                )
+                
+                if result and hasattr(result, 'generated_images') and result.generated_images:
+                    img_data = result.generated_images[0]
+                    if hasattr(img_data, 'image') and hasattr(img_data.image, 'image_bytes'):
+                        image = Image.open(io.BytesIO(img_data.image.image_bytes))
+                        logger.info(f"✅ Agent1: Successfully generated real image with Google Imagen!")
+                        return {
+                            "success": True,
+                            "image": image,
+                            "method": "Google Imagen",
+                            "enhanced_prompt": enhanced_prompt,
+                            "agent": "ImageGeneratorAgent"
+                        }
+                        
+                logger.warning(f"⚠️ Agent1: Google Imagen returned no images")
+                        
+            except Exception as e:
+                logger.error(f"❌ Agent1: Google Imagen failed: {e}")
+        
+        # Step 3: Fallback to demo image
+        logger.info(f"🤖 Agent1: Creating demo image (fallback)")
+        demo_image = self.create_demo_image(prompt, style, aspect_ratio)
+        
+        return {
+            "success": True,
+            "image": demo_image,
+            "method": "Demo Mode",
+            "enhanced_prompt": enhanced_prompt,
+            "agent": "ImageGeneratorAgent"
+        }
     
-    return None
-
-def review_with_gemini(image, prompt):
-    """Review image with Gemini Vision"""
-    if not google_auth_configured:
-        return "Demo mode - no AI review available"
+    def create_demo_image(self, prompt, style, aspect_ratio):
+        """Create demo image when API isn't available"""
+        # Set dimensions
+        if "Square" in aspect_ratio:
+            width, height = 400, 400
+        elif "Portrait" in aspect_ratio:
+            width, height = 400, 500
+        else:
+            width, height = 500, 300
+        
+        # Style colors
+        colors = {
+            "Professional": (59, 130, 246),
+            "Creative": (139, 92, 246),
+            "Minimalist": (107, 114, 128),
+            "Corporate": (30, 64, 175),
+            "Modern": (5, 150, 105)
+        }
+        
+        color = colors.get(style, (59, 130, 246))
+        img = Image.new('RGB', (width, height), color)
+        draw = ImageDraw.Draw(img)
         
-    try:
-        model = genai.GenerativeModel('gemini-2.0-flash-exp')
+        try:
+            font = ImageFont.load_default()
+            
+            # Add text overlay
+            title = "Marketing Image"
+            bbox = draw.textbbox((0, 0), title, font=font)
+            text_width = bbox[2] - bbox[0]
+            x = (width - text_width) // 2
+            y = height // 2 - 30
+            draw.text((x, y), title, fill="white", font=font)
+            
+            style_text = f"{style} Style"
+            bbox2 = draw.textbbox((0, 0), style_text, font=font)
+            text_width2 = bbox2[2] - bbox2[0]
+            x2 = (width - text_width2) // 2
+            draw.text((x2, y + 25), style_text, fill="white", font=font)
+            
+        except:
+            pass
+        
+        return img
+
+# ====== AGENT 2: IMAGE REVIEWER ======
+
+class ImageReviewerAgent:
+    """Agent 2: Responsible for reviewing images using Google Gemini Vision"""
+    
+    def __init__(self):
+        self.name = "ImageReviewerAgent"
+        self.status = "Ready"
+    
+    def review_image(self, image, original_prompt, enhanced_prompt, generation_method):
+        """Review generated image for quality and marketing effectiveness"""
+        logger.info(f"🔍 Agent2: Starting image review")
+        logger.info(f"🔍 Agent2: Reviewing image generated by: {generation_method}")
+        logger.info(f"🔍 Agent2: Original prompt: {original_prompt}")
         
-        review_prompt = f"""
-        Review this marketing image generated from: "{prompt}"
+        if google_auth_configured and genai and generation_method == "Google Imagen":
+            try:
+                logger.info(f"🔍 Agent2: Calling Google Gemini Vision for AI review...")
+                
+                model = genai.GenerativeModel('gemini-2.0-flash-exp')
+                
+                review_prompt = f"""You are an expert marketing image reviewer. Analyze this image that was generated from the prompt: "{original_prompt}"
+
+                Rate the image on a scale of 1-10 and provide specific feedback on:
+                1. Technical Quality (clarity, composition, lighting)
+                2. Marketing Effectiveness (professional appeal, brand suitability)  
+                3. Prompt Matching (how well it matches the original request)
+
+                Provide your response in this format:
+                SCORE: X/10
+                QUALITY: [brief assessment]
+                MARKETING: [marketing effectiveness]
+                MATCHING: [how well it matches prompt]
+                RECOMMENDATIONS: [specific improvements]
+
+                Keep total response under 200 words."""
+                
+                response = model.generate_content([review_prompt, image])
+                review_text = response.text
+                
+                logger.info(f"✅ Agent2: AI review completed successfully")
+                logger.info(f"🔍 Agent2: Review preview: {review_text[:100]}...")
+                
+                return {
+                    "success": True,
+                    "review": review_text,
+                    "method": "Gemini Vision AI",
+                    "agent": "ImageReviewerAgent"
+                }
+                
+            except Exception as e:
+                logger.error(f"❌ Agent2: Gemini Vision review failed: {e}")
         
-        Rate it 1-10 and provide brief feedback on:
-        - Visual quality and clarity
-        - Marketing effectiveness 
-        - How well it matches the prompt
+        # Fallback review for demo mode
+        logger.info(f"🔍 Agent2: Performing basic review (fallback)")
         
-        Keep response under 100 words.
-        """
+        word_count = len(original_prompt.split())
+        if word_count > 15:
+            quality = "Excellent"
+            score = "8/10"
+        elif word_count > 8:
+            quality = "Good" 
+            score = "7/10"
+        else:
+            quality = "Basic"
+            score = "6/10"
         
-        response = model.generate_content([review_prompt, image])
-        return response.text[:200]
+        basic_review = f"""SCORE: {score}
+QUALITY: {quality} prompt with {word_count} descriptive words
+MARKETING: Suitable for {generation_method.lower()} marketing use
+MATCHING: Generated image reflects the requested style and content
+RECOMMENDATIONS: Add more specific details for enhanced results
+
+Note: This is a basic review. Connect Google Cloud for full AI-powered analysis."""
         
-    except Exception as e:
-        logger.warning(f"Gemini review failed: {e}")
-        return "AI review unavailable"
+        return {
+            "success": True,
+            "review": basic_review,
+            "method": "Basic Analysis",
+            "agent": "ImageReviewerAgent"
+        }
 
-def generate_image(prompt, style, aspect_ratio):
-    """Main generation function"""
-    if not prompt.strip():
-        return None, "Please enter a description"
+# ====== ORCHESTRATOR: MANAGES AGENT WORKFLOW ======
+
+def orchestrate_workflow(prompt, style, aspect_ratio):
+    """Main workflow orchestrator - manages Agent1 -> Agent2 pipeline"""
+    logger.info(f"🎭 ORCHESTRATOR: Starting agent workflow")
+    logger.info(f"🎭 ORCHESTRATOR: Input - Prompt: {prompt}, Style: {style}, Format: {aspect_ratio}")
+    
+    if not prompt or not prompt.strip():
+        return None, "❌ ORCHESTRATOR: No prompt provided", "Workflow failed"
+    
+    workflow_log = []
+    
+    # Step 1: Initialize agents
+    logger.info(f"🎭 ORCHESTRATOR: Initializing agents...")
+    agent1 = ImageGeneratorAgent()
+    agent2 = ImageReviewerAgent()
+    workflow_log.append("✅ Agents initialized")
+    
+    # Step 2: Agent1 generates image
+    logger.info(f"🎭 ORCHESTRATOR: Calling Agent1 (ImageGenerator)...")
+    workflow_log.append("🤖 Agent1: Starting image generation...")
+    
+    generation_result = agent1.generate_image(prompt, style, aspect_ratio)
     
-    # Try real AI generation first
-    if google_auth_configured:
-        real_image = generate_with_google_imagen(prompt, style)
-        if real_image:
-            review = review_with_gemini(real_image, prompt)
-            return real_image, f"✅ Generated with Google Imagen\n\nAI Review: {review}"
+    if not generation_result["success"]:
+        logger.error(f"🎭 ORCHESTRATOR: Agent1 failed")
+        return None, "❌ Image generation failed", "Workflow failed"
     
-    # Fallback to placeholder
-    placeholder = create_placeholder_image(prompt, style, aspect_ratio)
+    workflow_log.append(f"✅ Agent1: Generated image using {generation_result['method']}")
     
-    # Simple quality estimate
-    word_count = len(prompt.split())
-    quality = "Good" if word_count > 10 else "Basic"
+    # Step 3: Agent2 reviews image
+    logger.info(f"🎭 ORCHESTRATOR: Calling Agent2 (ImageReviewer)...")
+    workflow_log.append("🔍 Agent2: Starting image review...")
     
-    feedback = f"⚠️ Demo mode active\n\nPrompt quality: {quality}\nStyle: {style}\n\nTo use real AI generation, configure Google Cloud credentials."
+    review_result = agent2.review_image(
+        generation_result["image"],
+        prompt,
+        generation_result["enhanced_prompt"],
+        generation_result["method"]
+    )
+    
+    workflow_log.append(f"✅ Agent2: Review completed using {review_result['method']}")
+    
+    # Step 4: Compile final results
+    logger.info(f"🎭 ORCHESTRATOR: Workflow completed successfully")
     
-    return placeholder, feedback
+    final_review = f"""🎭 AGENT WORKFLOW COMPLETED
+
+📋 WORKFLOW LOG:
+{chr(10).join(workflow_log)}
+
+🤖 AGENT1 RESULT: 
+Generated using: {generation_result['method']}
+Enhanced prompt: {generation_result['enhanced_prompt'][:100]}...
 
-# Create clean Gradio interface
-with gr.Blocks(title="Marketing Image Generator", theme=gr.themes.Default()) as demo:
+🔍 AGENT2 REVIEW:
+{review_result['review']}
+
+🎯 SYSTEM STATUS: {"✅ Full AI Pipeline" if google_auth_configured else "⚠️ Demo Mode"}
+"""
     
-    gr.Markdown("""
-    # 🎨 Marketing Image Generator
-    Create professional marketing images with AI-powered review
-    """)
+    return generation_result["image"], final_review, "✅ Workflow Success"
+
+# ====== GRADIO INTERFACE ======
+
+with gr.Blocks(title="Marketing Image Generator - Agent Workflow") as demo:
+    
+    gr.Markdown("# 🎨 Marketing Image Generator - Agent Workflow")
+    gr.Markdown("**Agent1** (ImageGenerator) → **Agent2** (ImageReviewer) → **Results**")
     
     # Status
-    if google_auth_configured:
-        gr.Markdown("✅ **Status**: Google AI Connected - Real image generation active")
-    else:
-        gr.Markdown("⚠️ **Status**: Demo Mode - Configure Google Cloud for real AI generation")
+    status_msg = "✅ Google AI Connected - Full Agent Pipeline" if google_auth_configured else "⚠️ Demo Mode - Simulated Agent Workflow"
+    gr.Markdown(f"**Status:** {status_msg}")
     
     with gr.Row():
-        with gr.Column():
-            # Inputs
+        with gr.Column(scale=1):
             prompt = gr.Textbox(
-                label="Describe your marketing image",
-                placeholder="A professional team collaborating in a modern office space...",
-                lines=3
+                label="Marketing Image Description",
+                placeholder="A professional team meeting in a modern office...",
+                lines=2
             )
             
-            with gr.Row():
-                style = gr.Dropdown(
-                    ["Professional", "Creative", "Minimalist", "Corporate", "Modern"],
-                    value="Professional",
-                    label="Style"
-                )
-                
-                aspect_ratio = gr.Dropdown(
-                    ["Landscape (16:9)", "Square (1:1)", "Portrait (3:4)"],
-                    value="Landscape (16:9)",
-                    label="Aspect Ratio"
-                )
+            style = gr.Dropdown(
+                ["Professional", "Creative", "Minimalist", "Corporate", "Modern"],
+                value="Professional",
+                label="Style"
+            )
+            
+            aspect_ratio = gr.Dropdown(
+                ["Landscape (16:9)", "Square (1:1)", "Portrait (3:4)"],
+                value="Landscape (16:9)",
+                label="Format"
+            )
             
-            generate_btn = gr.Button("Generate Image", variant="primary", size="lg")
+            generate_btn = gr.Button("🚀 Start Agent Workflow", variant="primary")
         
-        with gr.Column():
-            # Outputs
-            image_output = gr.Image(label="Generated Image", height=400)
-            feedback_output = gr.Textbox(label="AI Review", lines=4)
+        with gr.Column(scale=1):
+            image_output = gr.Image(label="Generated Image", height=300)
+            workflow_status = gr.Textbox(label="Workflow Status", lines=1, max_lines=1)
+    
+    # Detailed workflow log
+    workflow_output = gr.Textbox(label="Agent Workflow Log", lines=8, max_lines=12)
     
     # Examples
-    gr.Markdown("### Example Prompts")
-    examples = gr.Examples(
+    gr.Examples(
         examples=[
-            ["A diverse team of professionals in a bright modern office", "Professional", "Landscape (16:9)"],
-            ["A sleek product photo with clean white background", "Minimalist", "Square (1:1)"],
-            ["A friendly customer service representative with headset", "Corporate", "Portrait (3:4)"],
-            ["An inspiring workspace with laptop and coffee", "Modern", "Landscape (16:9)"]
+            ["Professional team collaboration in modern office", "Professional", "Landscape (16:9)"],
+            ["Clean product showcase on white background", "Minimalist", "Square (1:1)"],
+            ["Customer service representative with headset", "Corporate", "Portrait (3:4)"]
         ],
         inputs=[prompt, style, aspect_ratio]
     )
     
-    # Connect function
+    # Connect workflow
     generate_btn.click(
-        fn=generate_image,
+        fn=orchestrate_workflow,
         inputs=[prompt, style, aspect_ratio],
-        outputs=[image_output, feedback_output]
+        outputs=[image_output, workflow_output, workflow_status],
+        show_progress=True
     )
 
 if __name__ == "__main__":