agents/image_agent.py

from tools.generate import generate_image
from tools.edit import edit_image
from tools.analyze import analyze_image
from tools.upload import save_uploads


class ImageAgent:

    def __init__(self, router, memory, planner):
        self.router = router
        self.memory = memory
        self.planner = planner

    # --------------------------
    # Upload Files
    # --------------------------
    def upload(self, files):

        uploaded = save_uploads(files)

        self.memory.uploaded_files = uploaded

        # Automatically choose first uploaded image
        for file in uploaded:
            if file["type"] and file["type"].startswith("image/"):
                self.memory.last_used_image = file["path"]
                break

        return uploaded

    # --------------------------
    # Analyze Image
    # --------------------------
    def auto_analyze(self):

        image = self.memory.last_used_image

        if not image:
            return None

        analysis = analyze_image(image)

        self.memory.analysis = analysis
        self.memory.image_context = analysis

        return analysis

    # --------------------------
    # Main Agent
    # --------------------------
    def run(self, user_input):

        task = self.router.route(user_input)

        steps = self.planner.plan(task)

        result = None

        image = self.memory.last_used_image

        # Automatically analyze uploaded image once
        if image and self.memory.image_context is None:
            self.auto_analyze()

        analysis = self.memory.image_context

        context = ""

        if analysis:

            context = f"""
Image Description:
{analysis["description"]}

Objects:
{", ".join(analysis.get("objects", []))}

Style:
{analysis.get("style", "")}

Lighting:
{analysis.get("lighting", "")}

OCR:
{analysis.get("text", "")}
"""

        prompt = f"""
{context}

User Request:
{user_input}
"""

        for step in steps:

            if step == "generate":

                result = generate_image(prompt)

            elif step == "edit":

                result = edit_image(
                    image_path=image,
                    prompt=prompt
                )

                # Remember edited image
                if result:
                    self.memory.last_used_image = result

            elif step == "analyze":

                result = self.auto_analyze()

        return result