ai_qc/backend/visual_qc_apps/flask_app_template.py

from flask import Flask, request, jsonify
import os
import base64
import tempfile
from visual_qc_apps.utils import run_visual_qc

class FlaskAppTemplate:
    """Base template for Visual QC Flask applications"""
    
    def __init__(self, name, prompt):
        """
        Initialize the Flask app with a specific prompt
        
        Args:
            name (str): Name of the app
            prompt (str): The hardcoded prompt to use
        """
        self.app = Flask(name)
        
        # Add scoring instructions to the prompt
        scoring_instructions = """

SCORING INSTRUCTIONS:
After your analysis, provide a numerical score from 1 to 10 (with 10 being perfect/excellent and 1 being poor/inadequate) based on how well the asset meets the criteria.

YOUR OUTPUT FORMAT:
Include a JSON code block with the following fields:
  - "score": a number from 1 to 10
  - "explanation": your detailed reasoning for the score
  - "recommendations": specific suggestions for improvement if applicable

Example:
```json
{
  "score": 8,
  "explanation": "The logo is clearly visible and occupies approximately 12% of the advertisement area, which exceeds the minimum requirement. It has good contrast against the background and would be recognizable from the appropriate viewing distance.",
  "recommendations": "Consider increasing contrast around the edges of the logo for even better visibility in darker environments."
}
```
"""
        # Combine the base prompt with scoring instructions
        self.prompt = prompt + scoring_instructions
        
        # Register routes
        self.register_routes()
        
    def register_routes(self):
        """Register all API routes"""
        
        @self.app.route('/')
        def index():
            return jsonify({
                "status": "ok",
                "message": "Visual QC API active",
                "endpoints": [
                    "/api/analyze (POST)",
                    "/api/health (GET)"
                ]
            })
            
        @self.app.route('/api/health', methods=['GET'])
        def health_check():
            return jsonify({"status": "healthy"})
        
        @self.app.route('/api/analyze', methods=['POST'])
        def analyze_image():
            # Check if request has the required data
            if not request.json:
                return jsonify({"status": "error", "message": "Request must be JSON"}), 400
                
            # Validate required fields
            if 'image' not in request.json:
                return jsonify({"status": "error", "message": "Missing 'image' field"}), 400
                
            # Get parameters
            image_base64 = request.json.get('image')  # Required
            reference_base64 = request.json.get('reference')  # Optional
            model = request.json.get('model', 'Gemini')  # Optional, default to Gemini
            
            # Validate model selection
            if model not in ["Gemini", "OpenAI"]:
                return jsonify({"status": "error", "message": "Invalid model. Choose 'Gemini' or 'OpenAI'"}), 400
            
            try:
                # Create temporary files for images
                with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as main_file:
                    main_file.write(base64.b64decode(image_base64))
                    main_path = main_file.name
                    
                ref_path = None
                if reference_base64:
                    with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as ref_file:
                        ref_file.write(base64.b64decode(reference_base64))
                        ref_path = ref_file.name
                
                # Process the QC request
                result = run_visual_qc(
                    prompt=self.prompt,
                    asset_path=main_path,
                    reference_path=ref_path,
                    model_name=model
                )
                
                # Clean up temp files
                try:
                    os.unlink(main_path)
                    if ref_path:
                        os.unlink(ref_path)
                except Exception as e:
                    print(f"Error cleaning up temp files: {e}")
                
                # Return the result
                if result["status"] == "success":
                    return jsonify({
                        "status": "success",
                        "model": model,
                        "response": result["response"]
                    })
                else:
                    return jsonify({
                        "status": "error",
                        "message": result["message"]
                    }), 500
                    
            except Exception as e:
                return jsonify({
                    "status": "error",
                    "message": f"Server error: {str(e)}"
                }), 500
    
    def run(self, host='0.0.0.0', port=5001, debug=False):
        """Run the Flask application"""
        self.app.run(host=host, port=port, debug=debug)