Update app.py

app.py

@@ -1,9 +1,71 @@
 import gradio as gr
 from PIL import Image, ImageDraw, ImageFont
-import random
+import requests
+import json
+import base64
+import io
+import asyncio
+import aiohttp
+import os
+import time
+import traceback
 
-def create_image(prompt, style):
-    """Create a demo image"""
+# Configuration - matching full Streamlit functionality
+ORCHESTRATOR_URL = os.getenv("ORCHESTRATOR_URL", "http://localhost:8000")
+GENERATOR_URL = os.getenv("GENERATOR_URL", "http://localhost:8001")
+REVIEWER_URL = os.getenv("REVIEWER_URL", "http://localhost:8002")
+
+# Backward compatibility
+AGENT1_URL = os.getenv("AGENT1_URL", GENERATOR_URL)
+AGENT2_URL = os.getenv("AGENT2_URL", REVIEWER_URL)
+
+async def call_agent1_generate(prompt, style):
+    """Call Agent1 to generate an image"""
+    try:
+        payload = {
+            "prompt": prompt,
+            "style": style.lower(),
+            "size": "1024x1024",
+            "quality": "standard"
+        }
+        
+        async with aiohttp.ClientSession() as session:
+            async with session.post(f"{AGENT1_URL}/generate", json=payload, timeout=60) as response:
+                if response.status == 200:
+                    result = await response.json()
+                    return result
+                else:
+                    error_text = await response.text()
+                    return {"error": f"Agent1 error {response.status}: {error_text}"}
+    except Exception as e:
+        return {"error": f"Failed to connect to Agent1: {str(e)}"}
+
+async def call_agent2_review(image_url, prompt, review_guidelines=""):
+    """Call Agent2 to review the generated image"""
+    try:
+        payload = {
+            "image_url": image_url,
+            "original_prompt": prompt,
+            "review_criteria": ["quality", "relevance", "safety"]
+        }
+        
+        # Add review guidelines if provided
+        if review_guidelines.strip():
+            payload["review_guidelines"] = review_guidelines
+        
+        async with aiohttp.ClientSession() as session:
+            async with session.post(f"{AGENT2_URL}/review", json=payload, timeout=30) as response:
+                if response.status == 200:
+                    result = await response.json()
+                    return result
+                else:
+                    error_text = await response.text()
+                    return {"error": f"Agent2 error {response.status}: {error_text}"}
+    except Exception as e:
+        return {"error": f"Failed to connect to Agent2: {str(e)}"}
+
+def create_fallback_image(prompt, style):
+    """Create a fallback demo image when agents are unavailable"""
     colors = {"Professional": "#3B82F6", "Creative": "#8B5CF6", "Minimalist": "#6B7280", "Corporate": "#1E40AF", "Modern": "#059669"}
     color = colors.get(style, "#3B82F6")
     rgb = tuple(int(color[i:i+2], 16) for i in (1, 4, 7))
@@ -15,59 +77,553 @@ def create_image(prompt, style):
     draw.text((50, 200), "Marketing Image", fill="white", font=font)
     draw.text((50, 230), f"{style} Style", fill="white", font=font)
     draw.text((50, 260), prompt[:35] + "..." if len(prompt) > 35 else prompt, fill=(200, 200, 200), font=font)
+    draw.text((50, 350), "(Fallback Mode)", fill=(150, 150, 150), font=font)
     
     return img
 
-def generate_and_review(prompt, style):
-    """Generate image and review"""
+def generate_image_with_review(prompt, style, max_retries=3, review_threshold=0.8, review_guidelines=""):
+    """Generate an image with automated review using Agent1 and Agent2 with retry logic - mirrors Streamlit function"""
+    workflow_history = []
+    
+    try:
+        for iteration in range(1, max_retries + 1):
+            print(f"🔄 Iteration {iteration} of {max_retries}")
+            
+            # Step 1: Call Agent1 to generate the image
+            agent1_payload = {
+                "prompt": prompt,
+                "style": style.lower(),
+                "size": "1024x1024",
+                "quality": "standard"
+            }
+            
+            agent1_response = requests.post(f"{GENERATOR_URL}/generate", json=agent1_payload, timeout=60)
+            
+            if agent1_response.status_code != 200:
+                return {"success": False, "error": f"Agent1 failed: {agent1_response.text}"}
+            
+            agent1_result = agent1_response.json()
+            image_url = agent1_result.get("image_url", "")
+            
+            if not image_url:
+                return {"success": False, "error": "Agent1 did not return an image URL"}
+            
+            # Step 2: Call Agent2 to review the generated image
+            agent2_payload = {
+                "image_url": image_url,
+                "original_prompt": prompt,
+                "review_criteria": ["quality", "relevance", "safety"]
+            }
+            
+            # Add review guidelines if provided
+            if review_guidelines.strip():
+                agent2_payload["review_guidelines"] = review_guidelines
+            
+            agent2_response = requests.post(f"{REVIEWER_URL}/review", json=agent2_payload, timeout=30)
+            
+            if agent2_response.status_code != 200:
+                # Continue with just the image if review fails
+                workflow_history.append({
+                    "iteration": iteration,
+                    "agent1_status": "success",
+                    "agent2_status": "failed",
+                    "review_score": 0.7,
+                    "feedback": {"error": "Agent2 review failed"}
+                })
+                
+                return {
+                    "success": True,
+                    "image": {
+                        "url": image_url,
+                        "data": image_url,
+                        "prompt": prompt,
+                        "style": style
+                    },
+                    "review": {
+                        "quality_score": 0.7,
+                        "final_status": "review_failed",
+                        "iterations": iteration,
+                        "passed": True,
+                        "recommendations": ["Agent2 review failed - using generated image"],
+                        "workflow_history": workflow_history
+                    },
+                    "metadata": agent1_result.get("metadata", {})
+                }
+            
+            agent2_result = agent2_response.json()
+            review_score = agent2_result.get("review_score", 0.7)
+            
+            # Add to workflow history
+            workflow_history.append({
+                "iteration": iteration,
+                "agent1_status": "success",
+                "agent2_status": "success",
+                "review_score": review_score,
+                "feedback": agent2_result.get("feedback", {}),
+                "recommendations": agent2_result.get("recommendations", [])
+            })
+            
+            # Check if quality threshold is met
+            if review_score >= review_threshold:
+                print(f"✅ Quality threshold met on iteration {iteration}!")
+                
+                return {
+                    "success": True,
+                    "image": {
+                        "url": image_url,
+                        "data": image_url,
+                        "prompt": prompt,
+                        "style": style
+                    },
+                    "review": {
+                        "quality_score": review_score,
+                        "final_status": "passed",
+                        "iterations": iteration,
+                        "passed": True,
+                        "recommendations": agent2_result.get("recommendations", []),
+                        "workflow_history": workflow_history
+                    },
+                    "metadata": agent1_result.get("metadata", {})
+                }
+            else:
+                print(f"⚠️ Quality score {review_score:.2f} below threshold {review_threshold}. Retrying...")
+                if iteration < max_retries:
+                    # Enhance prompt for next iteration based on feedback
+                    feedback = agent2_result.get("feedback", {})
+                    if "relevance" in feedback and "missing_elements" in feedback["relevance"]:
+                        missing = feedback["relevance"]["missing_elements"]
+                        if missing:
+                            prompt += f" Including: {', '.join(missing[:3])}"
+        
+        # If we get here, all retries failed to meet threshold
+        return {
+            "success": True,
+            "image": {
+                "url": image_url,
+                "data": image_url,
+                "prompt": prompt,
+                "style": style
+            },
+            "review": {
+                "quality_score": review_score,
+                "final_status": "needs_improvement",
+                "iterations": max_retries,
+                "passed": False,
+                "recommendations": agent2_result.get("recommendations", []) + [f"Failed to meet quality threshold {review_threshold} after {max_retries} attempts"],
+                "workflow_history": workflow_history
+            },
+            "metadata": agent1_result.get("metadata", {})
+        }
+        
+    except Exception as e:
+        return {"success": False, "error": f"Unexpected error: {str(e)}"}
+
+def process_generated_image_and_results(api_response):
+    """Process API response and return image and review text for Gradio display"""
+    try:
+        # Parse the response if it's a string
+        if isinstance(api_response, str):
+            response_data = json.loads(api_response)
+        else:
+            response_data = api_response
+        
+        # Check if the response was successful
+        if not response_data.get('success', False):
+            return None, f"❌ API call failed: {response_data.get('error', 'Unknown error')}"
+        
+        # Data can be at top level or nested under 'data' key
+        data = response_data.get('data', response_data)
+        
+        # Extract image data
+        image_info = data.get('image', {})
+        image_data_b64 = ""
+        image = None
+        
+        # Try different possible image data locations
+        if 'data' in image_info:
+            image_data_b64 = image_info['data']
+        elif 'url' in image_info and image_info['url'].startswith('data:image'):
+            image_data_b64 = image_info['url']
+        elif 'image_url' in data:
+            image_data_b64 = data['image_url']
+        
+        if image_data_b64:
+            try:
+                # Handle data URL format
+                if image_data_b64.startswith('data:image'):
+                    base64_data = image_data_b64.split(',')[1]
+                elif image_data_b64.startswith('http'):
+                    # Handle regular URL (like picsum.photos)
+                    response = requests.get(image_data_b64, timeout=10)
+                    if response.status_code == 200:
+                        image = Image.open(io.BytesIO(response.content))
+                    else:
+                        image = None
+                else:
+                    base64_data = image_data_b64
+                
+                if image is None and 'base64_data' in locals():
+                    # Decode base64 image
+                    image_bytes = base64.b64decode(base64_data)
+                    image = Image.open(io.BytesIO(image_bytes))
+                
+            except Exception as e:
+                print(f"Error processing image: {str(e)}")
+                image = None
+        
+        # Extract review data
+        review_data = data.get('review', {})
+        if not review_data:
+            review_data = response_data.get('review', {})
+        
+        # Format review text
+        if review_data:
+            quality_score = review_data.get('quality_score', 0)
+            passed = review_data.get('passed', False)
+            final_status = review_data.get('final_status', 'unknown')
+            iterations = review_data.get('iterations', 0)
+            recommendations = review_data.get('recommendations', [])
+            
+            status_emoji = "🟢" if passed else "🔴"
+            
+            review_text = f"""**🔍 AI Review Results**
+
+**Quality Score:** {quality_score:.2f}/1.0
+**Status:** {status_emoji} {final_status.upper()}
+**Iterations:** {iterations}
+
+**💡 Recommendations:**
+"""
+            
+            if recommendations:
+                for i, rec in enumerate(recommendations[:5], 1):
+                    review_text += f"{i}. {rec}\n"
+            else:
+                review_text += "• Image meets quality standards\n"
+            
+            # Add workflow history if available
+            workflow_history = review_data.get('workflow_history', [])
+            if workflow_history and len(workflow_history) > 1:
+                review_text += "\n**📋 Workflow History:**\n"
+                for item in workflow_history:
+                    iteration = item.get('iteration', 'N/A')
+                    score = item.get('review_score', 'N/A')
+                    review_text += f"• Iteration {iteration}: Score {score:.2f}\n"
+        else:
+            review_text = "⚠️ Review data not available"
+        
+        return image, review_text
+        
+    except Exception as e:
+        error_text = f"❌ Error processing results: {str(e)}\n\n**Debug Info:**\n{traceback.format_exc()}"
+        return None, error_text
+
+async def generate_and_review_async(prompt, style, review_guidelines=""):
+    """Generate image with Agent1 and review with Agent2 - legacy function for backward compatibility"""
     if not prompt.strip():
         return None, "Please enter a prompt"
     
-    image = create_image(prompt, style)
-    word_count = len(prompt.split())
-    quality = "Excellent" if word_count > 15 else "Good" if word_count > 8 else "Basic"
+    # Step 1: Generate image with Agent1
+    agent1_result = await call_agent1_generate(prompt, style)
+    
+    if "error" in agent1_result:
+        # Fallback to demo image
+        fallback_image = create_fallback_image(prompt, style)
+        return fallback_image, f"**Agent1 Unavailable**: {agent1_result['error']}\n\nUsing fallback demo image."
+    
+    # Extract image from Agent1 response
+    image_url = agent1_result.get("image_url", "")
+    image = None
     
-    review = f"""**Agent Review Complete**
+    if image_url:
+        try:
+            if image_url.startswith("data:image"):
+                # Handle base64 data URL
+                base64_data = image_url.split(',')[1]
+                image_bytes = base64.b64decode(base64_data)
+                image = Image.open(io.BytesIO(image_bytes))
+            elif image_url.startswith("http"):
+                # Handle regular URL (like picsum.photos)
+                async with aiohttp.ClientSession() as session:
+                    async with session.get(image_url) as response:
+                        if response.status == 200:
+                            image_bytes = await response.read()
+                            image = Image.open(io.BytesIO(image_bytes))
+        except Exception as e:
+            print(f"Error loading image: {e}")
+    
+    if image is None:
+        image = create_fallback_image(prompt, style)
+        image_url = "fallback://demo"
+    
+    # Step 2: Review image with Agent2
+    agent2_result = await call_agent2_review(image_url, prompt, review_guidelines)
+    
+    if "error" in agent2_result:
+        # Simple review fallback
+        word_count = len(prompt.split())
+        quality = "Excellent" if word_count > 15 else "Good" if word_count > 8 else "Basic"
+        
+        review = f"""**Agent Review (Fallback Mode)**
 
-Agent1: Generated {style.lower()} marketing image
-Agent2: Quality analysis complete
+Agent1: Generated {style.lower()} marketing image ✅
+Agent2: Unavailable - {agent2_result['error']}
 
 Prompt Quality: {quality} ({word_count} words)
 Style: {style}
-Status: Ready for use
+Status: Image generated but not reviewed
 
 Recommendation: {f"Great detail level!" if word_count > 15 else "Consider adding more descriptive details"}"""
+    else:
+        # Format Agent2 review results
+        review_score = agent2_result.get("review_score", 0)
+        feedback = agent2_result.get("feedback", {})
+        recommendations = agent2_result.get("recommendations", [])
+        
+        quality_feedback = feedback.get("quality", {})
+        relevance_feedback = feedback.get("relevance", {})
+        safety_feedback = feedback.get("safety", {})
+        
+        review = f"""**🤖 AI Agent Review Complete**
+
+**Agent1**: Generated {style.lower()} marketing image ✅
+**Agent2**: Quality analysis complete ✅
+
+**Overall Score**: {review_score:.2f}/1.0
+**Quality**: {quality_feedback.get('score', 'N/A')}
+**Relevance**: {relevance_feedback.get('score', 'N/A')}
+**Safety**: {safety_feedback.get('score', 'N/A')}
+
+**Status**: {'✅ Approved' if review_score > 0.7 else '⚠️ Needs Improvement'}
+
+**Recommendations**:
+{chr(10).join(f"• {rec}" for rec in recommendations[:3]) if recommendations else "• Image meets quality standards"}"""
     
     return image, review
 
-# Create simple interface
-with gr.Blocks(title="AI Marketing Generator") as demo:
-    gr.Markdown("# 🎨 AI Marketing Image Generator")
-    gr.Markdown("Agent1 creates → Agent2 reviews → Professional results")
+def check_service_health():
+    """Check the health of all agent services"""
+    health_status = {}
+    
+    # Check Agent1 (Image Generator)
+    try:
+        response = requests.get(f"{GENERATOR_URL}/health", timeout=5)
+        health_status["Agent1 (Generator)"] = response.status_code == 200
+    except:
+        health_status["Agent1 (Generator)"] = False
+    
+    # Check Agent2 (Reviewer)
+    try:
+        response = requests.get(f"{REVIEWER_URL}/health", timeout=5)
+        health_status["Agent2 (Marketing Reviewer)"] = response.status_code == 200
+    except:
+        health_status["Agent2 (Marketing Reviewer)"] = False
+    
+    # Check Orchestrator (if available)
+    try:
+        response = requests.get(f"{ORCHESTRATOR_URL}/health", timeout=5)
+        health_status["Orchestrator"] = response.status_code == 200
+    except:
+        health_status["Orchestrator"] = False
+    
+    return health_status
+
+def get_system_status():
+    """Get system status for display"""
+    health_status = check_service_health()
+    
+    status_text = "**🔧 System Status:**\n\n"
+    for service_name, is_healthy in health_status.items():
+        status_emoji = "✅" if is_healthy else "❌"
+        status_text += f"{status_emoji} {service_name}\n"
+    
+    all_healthy = all(health_status.values())
+    if all_healthy:
+        status_text += "\n🎉 All services are running!"
+    else:
+        status_text += "\n⚠️ Some services are not responding."
+    
+    return status_text
+
+def generate_marketing_image(prompt, style, max_retries, review_threshold, review_guidelines):
+    """Main function called by Gradio interface - enhanced version"""
+    if not prompt.strip():
+        return None, "⚠️ Please enter a prompt to generate an image."
+    
+    try:
+        # Call the same backend function as Streamlit
+        result = generate_image_with_review(prompt, style, max_retries, review_threshold, review_guidelines)
+        
+        if result["success"]:
+            # Process the results for display
+            image, review_text = process_generated_image_and_results(result)
+            
+            success_message = f"✅ Image generated successfully!\n\n{review_text}"
+            return image, success_message
+        else:
+            error_message = f"❌ Generation failed: {result.get('error', 'Unknown error')}"
+            return None, error_message
+            
+    except Exception as e:
+        error_message = f"❌ Error generating image: {str(e)}\n\n**Debug Info:**\n{traceback.format_exc()}"
+        return None, error_message
+
+def generate_and_review(prompt, style):
+    """Sync wrapper for async function - legacy compatibility"""
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    try:
+        return loop.run_until_complete(generate_and_review_async(prompt, style))
+    finally:
+        loop.close()
+
+def use_suggested_prompt(suggested_prompt, suggested_style):
+    """Update prompt and style with suggested values"""
+    return suggested_prompt, suggested_style
+
+# Define suggested prompts (matching Streamlit app)
+SUGGESTED_PROMPTS = {
+    "Modern office team collaboration": ("A modern office space with diverse professionals collaborating around a sleek conference table, natural lighting, professional attire, English signage visible", "realistic"),
+    "Executive boardroom meeting": ("Professional executive boardroom with polished conference table, city skyline view, business documents, English presentations on screens", "realistic"),
+    "Customer service excellence": ("Professional customer service representative with headset in modern call center, English signage, clean corporate environment", "realistic"),
+    "Product showcase display": ("Clean product showcase on white background with professional lighting, English product labels, minimalist marketing aesthetic", "realistic"),
+    "Creative workspace design": ("Creative workspace with colorful design elements, inspirational English quotes on walls, modern furniture, artistic marketing materials", "artistic"),
+    "Brand presentation setup": ("Professional brand presentation setup with English branded materials, corporate colors, marketing displays, conference room setting", "realistic")
+}
+
+# Create enhanced Gradio interface
+with gr.Blocks(title="Marketing Image Generator", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🎨 Marketing Image Generator with Marketing Review
+    ### Create stunning marketing images with AI-powered Marketing Reviewer
+    
+    Agent1 creates → Agent2 reviews → Professional results with automated quality assurance
+    """)
     
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
+            gr.Markdown("### ⚙️ Configuration")
+            
+            # Main inputs
             prompt = gr.Textbox(
-                label="Marketing Image Description",
-                placeholder="A professional team meeting in a modern office",
-                lines=3
+                label="Describe your marketing image",
+                placeholder="e.g., A modern office space with natural lighting, featuring diverse professionals collaborating around a sleek conference table",
+                lines=4,
+                info="Be specific about the scene, style, mood, and any marketing elements you want to include"
             )
+            
             style = gr.Dropdown(
-                choices=["Professional", "Creative", "Minimalist", "Corporate", "Modern"],
-                value="Professional",
-                label="Brand Style"
+                choices=["realistic", "artistic", "cartoon", "photographic", "illustration"],
+                value="realistic",
+                label="Art Style",
+                info="Choose the artistic style for your generated image"
+            )
+            
+            review_guidelines = gr.Textbox(
+                label="🔍 Marketing Review Guidelines (Optional)",
+                placeholder="e.g., All text must be in English only, focus on professional appearance, ensure brand colors are prominent, check accessibility compliance, verify readability",
+                lines=3,
+                info="Provide specific marketing guidelines for the Marketing Reviewer to evaluate against your brand standards"
             )
-            generate_btn = gr.Button("Generate Image", variant="primary")
+            
+            # Advanced settings
+            with gr.Accordion("🔧 Advanced Settings", open=False):
+                max_retries = gr.Slider(
+                    minimum=1,
+                    maximum=5,
+                    value=3,
+                    step=1,
+                    label="Max Retries",
+                    info="Maximum number of retry attempts if generation fails"
+                )
+                
+                review_threshold = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.8,
+                    step=0.1,
+                    label="Quality Threshold",
+                    info="Minimum quality score required for auto-approval"
+                )
+            
+            # Generate buttons
+            generate_enhanced_btn = gr.Button("🚀 Generate with Full Review", variant="primary", size="lg")
+            generate_simple_btn = gr.Button("⚡ Quick Generate", variant="secondary", size="sm")
+            
+            # System status
+            with gr.Accordion("📊 System Status", open=False):
+                status_display = gr.Markdown(get_system_status())
+                refresh_status_btn = gr.Button("🔄 Refresh Status", size="sm")
+        
+        with gr.Column(scale=2):
+            # Results display
+            gr.Markdown("### 🖼️ Generated Image & Review")
+            
+            image_output = gr.Image(
+                label="Generated Marketing Image",
+                type="pil",
+                height=400,
+                show_download_button=True
+            )
+            
+            review_output = gr.Markdown(
+                value="Click **Generate** to create your marketing image with AI review",
+                label="AI Review Results"
+            )
+    
+    # Suggested prompts section
+    gr.Markdown("---")
+    gr.Markdown("### 💡 Suggested Marketing Prompts")
+    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("**🏢 Professional/Corporate**")
+            for prompt_name in ["Modern office team collaboration", "Executive boardroom meeting", "Customer service excellence"]:
+                suggested_prompt, suggested_style = SUGGESTED_PROMPTS[prompt_name]
+                btn = gr.Button(prompt_name, size="sm")
+                btn.click(
+                    fn=lambda p=suggested_prompt, s=suggested_style: (p, s),
+                    outputs=[prompt, style]
+                )
         
         with gr.Column():
-            image_output = gr.Image(label="Generated Image")
-            review_output = gr.Markdown(value="Click Generate to see review")
+            gr.Markdown("**🎨 Creative/Marketing**")
+            for prompt_name in ["Product showcase display", "Creative workspace design", "Brand presentation setup"]:
+                suggested_prompt, suggested_style = SUGGESTED_PROMPTS[prompt_name]
+                btn = gr.Button(prompt_name, size="sm")
+                btn.click(
+                    fn=lambda p=suggested_prompt, s=suggested_style: (p, s),
+                    outputs=[prompt, style]
+                )
+    
+    # Event handlers
+    generate_enhanced_btn.click(
+        fn=generate_marketing_image,
+        inputs=[prompt, style, max_retries, review_threshold, review_guidelines],
+        outputs=[image_output, review_output],
+        show_progress=True
+    )
     
-    generate_btn.click(
+    generate_simple_btn.click(
         fn=generate_and_review,
         inputs=[prompt, style],
-        outputs=[image_output, review_output]
+        outputs=[image_output, review_output],
+        show_progress=True
+    )
+    
+    refresh_status_btn.click(
+        fn=get_system_status,
+        outputs=status_display
     )
+    
+    # Footer
+    gr.Markdown("""
+    ---
+    <div style='text-align: center; color: #666; font-size: 0.9rem;'>
+        <p>🎨 Marketing Image Generator with Agent Review | Powered by Google Imagen3 & AI Agents</p>
+        <p>Create professional marketing images with automated quality assurance</p>
+    </div>
+    """)
 
 if __name__ == "__main__":
     demo.launch()