src/agents/image_agent.py

#!/usr/bin/env python3
"""
Image Agent — generates images via HuggingFace Inference API.

Routes image-generation requests from the email inbox through
FLUX.1-schnell (fast, free tier) to the HF Inference API.
Saves generated images to the Obsidian vault under Email/Images/.
"""
from __future__ import annotations

import json, os, re, datetime, io, base64, sys
from pathlib import Path
from typing import Any

import httpx

CONFIG_PATH = Path(__file__).resolve().parents[2] / "config" / "email.json"
HF_TOKEN = os.getenv("HF_TOKEN", "") or os.getenv("HF_ACCESS_TOKEN", "")
HF_API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
OBSIDIAN_VAULT = Path(os.getenv("OBSIDIAN_VAULT", "/tmp/email_vault"))
IMAGE_DIR = OBSIDIAN_VAULT / "Email" / "Images"


def load_config() -> dict:
    with open(CONFIG_PATH) as f:
        return json.load(f)


def generate_image(prompt: str, width: int = 1024, height: int = 1024,
                   num_steps: int = 4, guidance_scale: float = 0.0) -> dict[str, Any]:
    """Generate an image via HF Inference API."""
    if not HF_TOKEN:
        return {"error": "HF_TOKEN not set"}

    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    payload = {
        "inputs": prompt,
        "parameters": {
            "num_inference_steps": num_steps,
            "guidance_scale": guidance_scale,
            "width": width,
            "height": height,
            "max_sequence_length": 512,
        },
    }

    try:
        with httpx.Client(timeout=120, headers=headers) as client:
            resp = client.post(HF_API_URL, json=payload)
            if resp.status_code != 200:
                return {"error": f"HF API {resp.status_code}: {resp.text[:200]}"}

            # HF returns image bytes for image-generation tasks
            img_bytes = resp.content
            if not img_bytes:
                return {"error": "Empty response from HF"}
    except httpx.ConnectError as e:
        return {"error": f"Connection failed: {str(e)}"}
    except httpx.TimeoutException:
        return {"error": "HF API request timed out"}
    except Exception as e:
        return {"error": f"Request failed: {str(e)}"}

        # Save to Obsidian vault
        IMAGE_DIR.mkdir(parents=True, exist_ok=True)
        ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        safe_prompt = re.sub(r"[^\w\s-]", "", prompt)[:40].replace(" ", "_")
        fname = f"{ts}_{safe_prompt}.png"
        fpath = IMAGE_DIR / fname

        with open(fpath, "wb") as f:
            f.write(img_bytes)

        # Also save as base64 for markdown embedding
        b64 = base64.b64encode(img_bytes).decode()

        # Write metadata
        meta = {
            "prompt": prompt,
            "width": width,
            "height": height,
            "steps": num_steps,
            "guidance_scale": guidance_scale,
            "file": str(fpath),
            "size_bytes": len(img_bytes),
            "timestamp": datetime.datetime.now().isoformat(),
        }
        meta_path = IMAGE_DIR / f"{ts}_{safe_prompt}.json"
        with open(meta_path, "w") as f:
            json.dump(meta, f, indent=2)

        return {
            "status": "success",
            "prompt": prompt,
            "file": str(fpath),
            "size_bytes": len(img_bytes),
            "base64_preview": b64[:200] + "...",
            "metadata": meta,
        }


def generate_from_email(email_data: dict) -> dict[str, Any]:
    """Process an email requesting image generation."""
    subject = email_data.get("subject", "")
    body = email_data.get("body", "")
    sender = email_data.get("sender", "")

    # Extract image prompt from email body
    # Look for explicit prompt markers
    prompt_match = re.search(
        r"(?:image|picture|generate|create|draw|make)\s*:?\s*(.+?)(?:\n|$|(?:\.|!|\?)$)",
        body, re.IGNORECASE,
    )
    if prompt_match:
        prompt = prompt_match.group(1).strip().rstrip(".")
    else:
        # Use subject + first line of body as prompt
        prompt = f"{subject}. {body[:200]}"

    result = generate_image(prompt)

    # Store in vault with email context
    if "error" not in result:
        vault_entry = {
            "source_email": {
                "subject": subject,
                "sender": sender,
                "date": email_data.get("date", ""),
            },
            "image_result": result,
        }
        meta_path = Path(result["file"]).with_suffix(".json")
        with open(meta_path, "w") as f:
            json.dump(vault_entry, f, indent=2)

    return result


def list_generated_images(limit: int = 20) -> list[dict]:
    """List generated images from the vault."""
    images = []
    if not IMAGE_DIR.exists():
        return images

    for f in sorted(IMAGE_DIR.glob("*.png"), reverse=True)[:limit]:
        meta_path = f.with_suffix(".json")
        meta = {}
        if meta_path.exists():
            with open(meta_path) as mf:
                meta = json.load(mf)
        images.append({
            "file": f.name,
            "size_bytes": f.stat().st_size,
            "metadata": meta,
        })
    return images


if __name__ == "__main__":
    # CLI test mode
    if len(sys.argv) > 1:
        result = generate_image(sys.argv[1])
        print(json.dumps(result, indent=2, default=str))
    else:
        print(json.dumps(list_generated_images(), indent=2))