image_detector.py

"""
Image quality detector based on generation time.
"""

import logging
from typing import Dict, Any
from agentic_reliability_framework.runtime.agents.base import BaseAgent, AgentSpecialization
from ai_event import AIEvent

logger = logging.getLogger(__name__)

class ImageQualityDetector(BaseAgent):
    """
    Detects potential quality issues in generated images.
    Currently uses generation time as a proxy; can be extended with CLIP or aesthetic models.
    """

    def __init__(self):
        super().__init__(AgentSpecialization.DETECTIVE)
        self._thresholds = {'generation_time': 2.0}  # seconds

    async def analyze(self, event: AIEvent) -> Dict[str, Any]:
        """
        Analyze image generation event.
        Expects `retrieval_scores` to contain [retrieval_score, generation_time].
        """
        gen_time = 1.0
        if event.retrieval_scores and len(event.retrieval_scores) > 1:
            gen_time = event.retrieval_scores[1]

        flags = []
        if gen_time > self._thresholds['generation_time']:
            flags.append('slow_generation')

        confidence = max(0, 1.0 - (gen_time / 5.0))

        return {
            'specialization': 'image_quality',
            'confidence': confidence,
            'findings': {
                'flags': flags,
                'generation_time': gen_time,
                'quality_issues': flags
            },
            'recommendations': [
                "Use a simpler prompt",
                "Upgrade hardware",
                "Reduce resolution"
            ] if flags else []
        }