agents/specialized/image_agent.py

# agents/specialized/image_agent.py
import os
import re
import json
import base64
import requests
from pathlib import Path
from datetime import datetime
from agents.specialized.base import build_result, call_llm
from agents.registry import AGENT_REGISTRY

ROLE = """
Eres el agente de imágenes.
Cuando se te pida generar, crear, mostrar o visualizar una imagen, responde SOLO con un JSON:
{"image_queries": ["término en inglés 1", "término en inglés 2", "término en inglés 3"]}
Los términos deben ser específicos, detallados y en inglés para obtener mejores resultados.
Si la tarea no es sobre imágenes, responde: {"skip":"no image task"}
"""

DOCS_DIR = Path("data/docs")


async def generate_with_gemini(queries: list, model_name: str) -> list:
    # ⚠️ pendiente de implementar con Gemini imagen real
    return []


async def generate_with_hf(queries: list) -> list:
    urls = []
    for prompt in queries[:2]:
        response = requests.post(
            "https://router.huggingface.co/v1/images/generations",
            headers={
                "Authorization": f"Bearer {os.getenv('HF_API_TOKEN')}",
                "Content-Type": "application/json",
            },
            json={
                "model": "black-forest-labs/FLUX.1-schnell",
                "prompt": prompt,
                "size": "1024x1024",
            },
        )

        if response.status_code == 200:
            data = response.json()

            # Algunos modelos devuelven b64_json, otros url directa
            item = data["data"][0]
            if "url" in item:
                urls.append(item["url"])
            elif "b64_json" in item:
                image_bytes = base64.b64decode(item["b64_json"])
                DOCS_DIR.mkdir(parents=True, exist_ok=True)
                file_name = f"hf_image_{datetime.now().timestamp()}.png"
                file_path = DOCS_DIR / file_name
                with open(file_path, "wb") as f:
                    f.write(image_bytes)
                urls.append(f"/docs/{file_name}")
        else:
            print("HF image error:", response.text)

    return urls


async def run(task: str, context: dict = None) -> dict:
    result = build_result("image_agent")
    try:
        response = await call_llm("image_agent", ROLE, task, context)
        result["response"] = response

        agent_config = AGENT_REGISTRY.get("image_agent", {})
        provider = agent_config.get("provider")
        model_name = agent_config.get("models", [None])[0]

        print("USANDO PROVIDER:", provider)
        print("USANDO MODELO:", model_name)

        match = re.search(r'\{.*"image_queries".*\}', response, re.DOTALL)

        if match:
            try:
                data = json.loads(match.group(0))
                queries = data.get("image_queries", [])
                result["queries"] = queries
                result["response"] = "Ideas de imágenes generadas:\n" + "\n".join(queries)

                image_urls = []

                # Intento 1: provider principal
                try:
                    if provider == "gemini":
                        image_urls = await generate_with_gemini(queries, model_name)
                    elif provider == "huggingface":
                        image_urls = await generate_with_hf(queries)
                except Exception as e:
                    print("Provider principal falló:", str(e))

                # Fallback a HF
                if not image_urls:
                    try:
                        print("Intentando fallback HF...")
                        image_urls = await generate_with_hf(queries)
                    except Exception as e:
                        print("Fallback HF falló:", str(e))

                # Fallback visual final
                if not image_urls:
                    image_urls = ["https://picsum.photos/400/300"]

                result["image_urls"] = image_urls
                result["response"] += "\n\nImagen generada ✔" if image_urls else "\n\n⚠️ No se pudo generar imagen"

            except Exception as e:
                result["response"] = f"Error procesando JSON: {str(e)}"
        else:
            result["response"] = "No se detectaron instrucciones claras para imágenes."

    except Exception as e:
        result["success"] = False
        result["error"] = str(e)

    return result