settled-reach/tooling/db/image_connector.py

#!/usr/bin/env python3
"""
Gemini image generator connector — direct API wrapper.

Generates images via Google's Gemini 2.0 Flash image generation API.
API key from GEMINI_API_KEY env var or config.json.

Usage:
    python3 image_connector.py health
    python3 image_connector.py generate "prompt" [--output file.png] [--aspect 1:1] [--size 1K] [--input image.png]
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from common import ensure_venv, get_api_key as _get_api_key  # noqa: E402

ensure_venv()

import base64
import json
import os
import urllib.error
import urllib.request

DEFAULT_OUTPUT_DIR = os.path.expanduser("~/Pictures/mcp-images")


def get_api_key() -> str:
    """Get Gemini API key from env or config."""
    return _get_api_key("GEMINI_API_KEY", "gemini_api_key")


def health():
    """Check if the Gemini API is reachable with the configured key."""
    key = get_api_key()
    url = f"https://generativelanguage.googleapis.com/v1beta/models?key={key}"
    try:
        req = urllib.request.Request(url, method="GET")
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = json.loads(resp.read())
            models = [m.get("name", "") for m in data.get("models", [])
                      if "imagen" in m.get("name", "").lower()
                      or "flash" in m.get("name", "").lower()]
            print(json.dumps({
                "ok": True,
                "api": "gemini",
                "image_capable_models": models[:5],
            }, indent=2))
    except Exception as e:
        print(json.dumps({
            "ok": False,
            "error": str(e)
        }, indent=2))
        sys.exit(1)


def generate(prompt, output=None, aspect_ratio="1:1", image_size=None,
             input_image=None):
    """
    Generate an image from a text prompt using Gemini.

    Args:
        prompt: Text description of the image to generate
        output: Output file path (default: auto-named in ~/Pictures/mcp-images/)
        aspect_ratio: Aspect ratio (1:1, 16:9, 3:2, etc.)
        image_size: Resolution hint (1K, 2K, 4K) - may not be honored
        input_image: Optional input image path for image-to-image generation
    """
    key = get_api_key()

    # Gemini image generation model
    model = "gemini-2.5-flash-image"
    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}"

    if output is None:
        safe = "".join(c if c.isalnum() or c in "-_ " else "" for c in prompt[:40])
        safe = safe.strip().replace(" ", "_").lower()
        os.makedirs(DEFAULT_OUTPUT_DIR, exist_ok=True)
        output = os.path.join(DEFAULT_OUTPUT_DIR, f"{safe}.png")

    # Build the request
    parts = []

    # Add input image if provided (image-to-image)
    if input_image:
        if not os.path.isfile(input_image):
            print(json.dumps({"ok": False, "error": f"Input image not found: {input_image}"}), indent=2)
            sys.exit(1)
        with open(input_image, "rb") as f:
            image_data = base64.b64encode(f.read()).decode("utf-8")
        # Detect mime type
        ext = os.path.splitext(input_image)[1].lower()
        mime = {"png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg",
                ".webp": "image/webp"}.get(ext, "image/png")
        parts.append({
            "inlineData": {
                "mimeType": mime,
                "data": image_data
            }
        })

    # Build enhanced prompt with aspect ratio and size hints
    enhanced_prompt = prompt
    if aspect_ratio and aspect_ratio != "1:1":
        enhanced_prompt += f" Aspect ratio: {aspect_ratio}."
    if image_size:
        enhanced_prompt += f" Resolution: {image_size}."

    parts.append({"text": enhanced_prompt})

    payload = json.dumps({
        "contents": [{"parts": parts}],
        "generationConfig": {
            "responseModalities": ["TEXT", "IMAGE"],
        }
    })

    req = urllib.request.Request(
        url,
        data=payload.encode(),
        headers={"Content-Type": "application/json"},
        method="POST"
    )

    print("Generating image...", file=sys.stderr)
    print(f"  Prompt: {prompt}", file=sys.stderr)
    if input_image:
        print(f"  Input image: {input_image}", file=sys.stderr)

    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            result = json.loads(resp.read())
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", errors="replace")
        print(json.dumps({
            "ok": False,
            "error": f"API error {e.code}: {e.reason}",
            "details": body[:500]
        }, indent=2))
        sys.exit(1)
    except Exception as e:
        print(json.dumps({"ok": False, "error": str(e)}), indent=2)
        sys.exit(1)

    # Extract image data from response
    candidates = result.get("candidates", [])
    if not candidates:
        print(json.dumps({
            "ok": False,
            "error": "No candidates in response",
            "response": json.dumps(result)[:500]
        }, indent=2))
        sys.exit(1)

    image_saved = False
    text_response = ""

    for candidate in candidates:
        content = candidate.get("content", {})
        for part in content.get("parts", []):
            if "inlineData" in part:
                # Image data
                image_b64 = part["inlineData"]["data"]
                image_bytes = base64.b64decode(image_b64)
                os.makedirs(os.path.dirname(os.path.abspath(output)), exist_ok=True)
                with open(output, "wb") as f:
                    f.write(image_bytes)
                image_saved = True
            elif "text" in part:
                text_response += part["text"]

    if not image_saved:
        print(json.dumps({
            "ok": False,
            "error": "No image data in response",
            "text_response": text_response[:500],
            "response": json.dumps(result)[:500]
        }, indent=2))
        sys.exit(1)

    file_size = os.path.getsize(output)
    print(json.dumps({
        "ok": True,
        "file": output,
        "size_bytes": file_size,
        "prompt": prompt,
        "aspect_ratio": aspect_ratio,
    }, indent=2))


def main():
    if len(sys.argv) < 2:
        print("Usage:")
        print("  image_connector.py health")
        print("  image_connector.py generate 'prompt' [--output file.png] [--aspect 1:1] [--size 1K] [--input image.png]")
        sys.exit(1)

    cmd = sys.argv[1]

    if cmd == "health":
        health()
    elif cmd == "generate":
        if len(sys.argv) < 3:
            print("Error: prompt required", file=sys.stderr)
            sys.exit(1)

        prompt = sys.argv[2]
        output = None
        aspect_ratio = "1:1"
        image_size = None
        input_image = None

        i = 3
        while i < len(sys.argv):
            if sys.argv[i] == "--output" and i + 1 < len(sys.argv):
                output = sys.argv[i + 1]
                i += 2
            elif sys.argv[i] == "--aspect" and i + 1 < len(sys.argv):
                aspect_ratio = sys.argv[i + 1]
                i += 2
            elif sys.argv[i] == "--size" and i + 1 < len(sys.argv):
                image_size = sys.argv[i + 1]
                i += 2
            elif sys.argv[i] == "--input" and i + 1 < len(sys.argv):
                input_image = sys.argv[i + 1]
                i += 2
            else:
                print(f"Unknown argument: {sys.argv[i]}", file=sys.stderr)
                sys.exit(1)

        generate(prompt, output=output, aspect_ratio=aspect_ratio,
                 image_size=image_size, input_image=input_image)
    else:
        print(f"Unknown command: {cmd}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()