@astrofoundry/pi-astro 0.5.1 → 0.6.0

package/extensions/gemini-image/index.ts

@@ -0,0 +1,313 @@
+import { readFileSync, writeFileSync } from "node:fs";
+import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
+import type { AgentToolResult } from "@mariozechner/pi-agent-core";
+import { GoogleGenAI } from "@google/genai";
+import { StringEnum } from "@mariozechner/pi-ai";
+import { Type } from "typebox";
+import { persistApiKey, resolveExistingApiKey } from "./credentials.ts";
+import {
+    ALL_MODELS,
+    DEFAULT_MODEL,
+    GEMINI_ASPECT_RATIOS,
+    GEMINI_IMAGE_SIZES,
+    IMAGEN_ASPECT_RATIOS,
+    IMAGEN_IMAGE_SIZES,
+    PERSON_GENERATION_VALUES,
+    isGeminiNative,
+    isImagen,
+    type ImageModel,
+} from "./models.ts";
+import { estimateCost, formatUsd, hasPricing } from "./pricing.ts";
+
+const params = Type.Object({
+    prompt: Type.String({ description: "Text prompt describing the image to generate." }),
+    model: Type.Optional(
+        StringEnum(ALL_MODELS as unknown as readonly [string, ...string[]], {
+            description: `Model. Default: ${DEFAULT_MODEL}.`,
+        }),
+    ),
+    aspect_ratio: Type.Optional(Type.String({ description: "Aspect ratio; validated per model family." })),
+    image_size: Type.Optional(
+        StringEnum(["1K", "2K", "4K"] as const, { description: "Output resolution. 4K is Gemini-native only." }),
+    ),
+    number_of_images: Type.Optional(
+        Type.Integer({ minimum: 1, maximum: 4, description: "Imagen only. 1-4." }),
+    ),
+    negative_prompt: Type.Optional(Type.String({ description: "Imagen only." })),
+    person_generation: Type.Optional(
+        StringEnum(PERSON_GENERATION_VALUES as unknown as readonly [string, ...string[]], {
+            description: "People policy: dont_allow | allow_adult | allow_all.",
+        }),
+    ),
+    seed: Type.Optional(Type.Integer({ description: "Imagen only. Reproducibility seed." })),
+    input_images: Type.Optional(
+        Type.Array(Type.String(), {
+            description: "Gemini native only. Absolute paths to images for editing.",
+        }),
+    ),
+    save_to: Type.Optional(Type.String({ description: "Optional directory to write generated PNG(s)." })),
+    skip_confirm: Type.Optional(Type.Boolean({ description: "Skip the cost confirmation prompt." })),
+});
+
+type Params = {
+    prompt: string;
+    model?: string;
+    aspect_ratio?: string;
+    image_size?: "1K" | "2K" | "4K";
+    number_of_images?: number;
+    negative_prompt?: string;
+    person_generation?: (typeof PERSON_GENERATION_VALUES)[number];
+    seed?: number;
+    input_images?: string[];
+    save_to?: string;
+    skip_confirm?: boolean;
+};
+
+async function ensureApiKey(ctx: ExtensionContext): Promise<string> {
+    const existing = resolveExistingApiKey();
+    if (existing) return existing;
+
+    if (!ctx.hasUI) {
+        throw new Error("Gemini API key is not configured and no interactive UI is available to ask for one.");
+    }
+
+    const pasted = await ctx.ui.input("Paste your Gemini API key:");
+    if (!pasted || !pasted.trim()) {
+        throw new Error("No API key provided.");
+    }
+    const key = pasted.trim();
+
+    const save = await ctx.ui.confirm("Save for future use?", "The key will be kept locally for later calls.");
+    if (save) {
+        persistApiKey(key);
+    }
+    return key;
+}
+
+function validateForModel(model: ImageModel, p: Params): void {
+    if (isImagen(model)) {
+        if (p.aspect_ratio && !(IMAGEN_ASPECT_RATIOS as readonly string[]).includes(p.aspect_ratio)) {
+            throw new Error(
+                `aspect_ratio "${p.aspect_ratio}" is not supported for Imagen. Valid: ${IMAGEN_ASPECT_RATIOS.join(", ")}.`,
+            );
+        }
+        if (p.image_size && !(IMAGEN_IMAGE_SIZES as readonly string[]).includes(p.image_size)) {
+            throw new Error(
+                `image_size "${p.image_size}" is not supported for Imagen. Valid: ${IMAGEN_IMAGE_SIZES.join(", ")}.`,
+            );
+        }
+        if (p.input_images && p.input_images.length > 0) {
+            throw new Error("Imagen does not support input_images (text-to-image only). Use a Gemini-native model.");
+        }
+    } else if (isGeminiNative(model)) {
+        if (p.aspect_ratio && !(GEMINI_ASPECT_RATIOS as readonly string[]).includes(p.aspect_ratio)) {
+            throw new Error(
+                `aspect_ratio "${p.aspect_ratio}" is not supported for Gemini native. Valid: ${GEMINI_ASPECT_RATIOS.join(", ")}.`,
+            );
+        }
+        if (p.image_size && !(GEMINI_IMAGE_SIZES as readonly string[]).includes(p.image_size)) {
+            throw new Error(
+                `image_size "${p.image_size}" is not supported for Gemini native. Valid: ${GEMINI_IMAGE_SIZES.join(", ")}.`,
+            );
+        }
+        if (p.number_of_images !== undefined && p.number_of_images !== 1) {
+            throw new Error("Gemini native produces a single image per call; number_of_images only applies to Imagen.");
+        }
+        if (p.negative_prompt) {
+            throw new Error("negative_prompt is Imagen-only.");
+        }
+        if (p.seed !== undefined) {
+            throw new Error("seed is Imagen-only.");
+        }
+    }
+}
+
+interface DecodedImage {
+    data: string; // base64
+    mimeType: string;
+}
+
+async function callGeminiNative(
+    ai: GoogleGenAI,
+    model: string,
+    p: Params,
+    signal: AbortSignal | undefined,
+): Promise<DecodedImage[]> {
+    const imageConfig: Record<string, string> = {};
+    if (p.aspect_ratio) imageConfig.aspectRatio = p.aspect_ratio;
+    if (p.image_size) imageConfig.imageSize = p.image_size;
+    if (p.person_generation) imageConfig.personGeneration = p.person_generation;
+
+    const contents: Array<Record<string, unknown>> = [{ text: p.prompt }];
+    for (const path of p.input_images ?? []) {
+        const buffer = readFileSync(path);
+        const mimeType = path.toLowerCase().endsWith(".jpg") || path.toLowerCase().endsWith(".jpeg")
+            ? "image/jpeg"
+            : path.toLowerCase().endsWith(".webp")
+            ? "image/webp"
+            : "image/png";
+        contents.push({
+            inlineData: {
+                mimeType,
+                data: buffer.toString("base64"),
+            },
+        });
+    }
+
+    const response = await ai.models.generateContent({
+        model,
+        contents,
+        config: {
+            responseModalities: ["Text", "Image"],
+            ...(Object.keys(imageConfig).length > 0 ? { imageConfig } : {}),
+            ...(signal ? { abortSignal: signal } : {}),
+        },
+    });
+
+    const images: DecodedImage[] = [];
+    for (const candidate of response.candidates ?? []) {
+        for (const part of candidate.content?.parts ?? []) {
+            if (part.inlineData?.data && part.inlineData.mimeType) {
+                images.push({ data: part.inlineData.data, mimeType: part.inlineData.mimeType });
+            }
+        }
+    }
+    return images;
+}
+
+async function callImagen(
+    ai: GoogleGenAI,
+    model: string,
+    p: Params,
+    signal: AbortSignal | undefined,
+): Promise<DecodedImage[]> {
+    const config: Record<string, unknown> = {};
+    if (p.aspect_ratio) config.aspectRatio = p.aspect_ratio;
+    if (p.image_size) config.imageSize = p.image_size;
+    if (p.number_of_images !== undefined) config.numberOfImages = p.number_of_images;
+    if (p.negative_prompt) config.negativePrompt = p.negative_prompt;
+    if (p.person_generation) config.personGeneration = p.person_generation;
+    if (p.seed !== undefined) config.seed = p.seed;
+    if (signal) config.abortSignal = signal;
+
+    const response = await ai.models.generateImages({
+        model,
+        prompt: p.prompt,
+        ...(Object.keys(config).length > 0 ? { config } : {}),
+    });
+
+    const images: DecodedImage[] = [];
+    for (const generated of response.generatedImages ?? []) {
+        if (generated.image?.imageBytes) {
+            images.push({
+                data: generated.image.imageBytes,
+                mimeType: generated.image.mimeType ?? "image/png",
+            });
+        }
+    }
+    return images;
+}
+
+function saveImages(images: DecodedImage[], dir: string, prompt: string): string[] {
+    const slug = prompt
+        .toLowerCase()
+        .replace(/[^a-z0-9]+/g, "-")
+        .replace(/^-+|-+$/g, "")
+        .slice(0, 40) || "image";
+    const ts = new Date().toISOString().replace(/[:.]/g, "-");
+    const paths: string[] = [];
+    for (let i = 0; i < images.length; i++) {
+        const ext = images[i].mimeType === "image/jpeg" ? "jpg" : "png";
+        const base = images.length > 1 ? `${slug}-${ts}-${i + 1}.${ext}` : `${slug}-${ts}.${ext}`;
+        const full = `${dir.replace(/\/$/, "")}/${base}`;
+        writeFileSync(full, Buffer.from(images[i].data, "base64"));
+        paths.push(full);
+    }
+    return paths;
+}
+
+export default function geminiImageExtension(pi: ExtensionAPI): void {
+    pi.registerTool({
+        name: "gemini_image",
+        label: "Gemini Image",
+        description:
+            "Generate or edit images via Google's Gemini (native) and Imagen models. Supports aspect ratios, resolutions, image editing, negative prompts, and seeds.",
+        promptGuidelines: [
+            "Use gemini-2.5-flash-image (default) for general-purpose generation.",
+            "Use gemini-3-pro-image-preview for studio-quality 4K, complex layouts, or precise text rendering.",
+            "Use imagen-4.0-*-generate-001 models for batch text-to-image (1-4 per call) or when negative_prompt / seed is needed.",
+            "Image editing (input_images) is only supported by Gemini native models.",
+            "4K output is only supported by Gemini native models.",
+        ],
+        parameters: params,
+        async execute(_toolCallId, input, signal, _onUpdate, ctx) {
+            const p = input as Params;
+            const model = (p.model as ImageModel | undefined) ?? DEFAULT_MODEL;
+
+            if (!(ALL_MODELS as readonly string[]).includes(model)) {
+                throw new Error(`Unknown model "${model}". Valid: ${ALL_MODELS.join(", ")}.`);
+            }
+
+            validateForModel(model, p);
+
+            const apiKey = await ensureApiKey(ctx);
+
+            const numImages = isImagen(model) ? p.number_of_images ?? 1 : 1;
+
+            if (!p.skip_confirm && hasPricing(model) && ctx.hasUI) {
+                const est = estimateCost({
+                    model,
+                    prompt: p.prompt,
+                    numberOfImages: numImages,
+                    imageSize: p.image_size,
+                });
+                const proceed = await ctx.ui.confirm(
+                    "Confirm image generation",
+                    `Model: ${model}\nImages: ${numImages}${p.image_size ? `\nSize: ${p.image_size}` : ""}\nEstimated: ~${formatUsd(est.estimatedUsd)} (${est.breakdown})`,
+                );
+                if (!proceed) {
+                    throw new Error("User cancelled image generation.");
+                }
+            }
+
+            const ai = new GoogleGenAI({ apiKey });
+
+            const images = isImagen(model)
+                ? await callImagen(ai, model, p, signal)
+                : await callGeminiNative(ai, model, p, signal);
+
+            if (images.length === 0) {
+                throw new Error("Model returned no images.");
+            }
+
+            const savedPaths = p.save_to ? saveImages(images, p.save_to, p.prompt) : [];
+
+            const summaryLines = [
+                `Generated ${images.length} image${images.length > 1 ? "s" : ""} with ${model}.`,
+                `Prompt: "${p.prompt}"`,
+            ];
+            if (p.aspect_ratio) summaryLines.push(`Aspect ratio: ${p.aspect_ratio}`);
+            if (p.image_size) summaryLines.push(`Size: ${p.image_size}`);
+            if (p.negative_prompt) summaryLines.push(`Negative: "${p.negative_prompt}"`);
+            if (p.person_generation) summaryLines.push(`People: ${p.person_generation}`);
+            if (p.seed !== undefined) summaryLines.push(`Seed: ${p.seed}`);
+            if (savedPaths.length > 0) summaryLines.push(`Saved to:\n  ${savedPaths.join("\n  ")}`);
+
+            const content: AgentToolResult<unknown>["content"] = images.map((img) => ({
+                type: "image",
+                data: img.data,
+                mimeType: img.mimeType,
+            }));
+            content.push({ type: "text", text: summaryLines.join("\n") });
+
+            return {
+                content,
+                details: {
+                    model,
+                    images: images.length,
+                    savedPaths,
+                },
+            };
+        },
+    });
+}

package/extensions/gemini-image/models.test.ts

@@ -0,0 +1,45 @@
+import { describe, expect, it } from "vitest";
+import {
+    ALL_MODELS,
+    DEFAULT_MODEL,
+    GEMINI_ASPECT_RATIOS,
+    GEMINI_NATIVE_MODELS,
+    IMAGEN_ASPECT_RATIOS,
+    IMAGEN_MODELS,
+    isGeminiNative,
+    isImagen,
+} from "./models.ts";
+
+describe("models", () => {
+    it("has 3 Gemini native + 3 Imagen = 6 models", () => {
+        expect(GEMINI_NATIVE_MODELS).toHaveLength(3);
+        expect(IMAGEN_MODELS).toHaveLength(3);
+        expect(ALL_MODELS).toHaveLength(6);
+    });
+
+    it("default is gemini-2.5-flash-image", () => {
+        expect(DEFAULT_MODEL).toBe("gemini-2.5-flash-image");
+    });
+
+    it("Gemini aspect ratios include 4:5 and 5:4 (missing from the junior's impl)", () => {
+        expect(GEMINI_ASPECT_RATIOS).toContain("4:5");
+        expect(GEMINI_ASPECT_RATIOS).toContain("5:4");
+    });
+
+    it("Imagen has 5 aspect ratios (no 2:3 / 3:2 / 4:5 / 5:4 / 21:9)", () => {
+        expect(IMAGEN_ASPECT_RATIOS).toHaveLength(5);
+        expect(IMAGEN_ASPECT_RATIOS).not.toContain("2:3");
+    });
+
+    it("isGeminiNative narrows correctly", () => {
+        expect(isGeminiNative("gemini-2.5-flash-image")).toBe(true);
+        expect(isGeminiNative("imagen-4.0-generate-001")).toBe(false);
+        expect(isGeminiNative("unknown")).toBe(false);
+    });
+
+    it("isImagen narrows correctly", () => {
+        expect(isImagen("imagen-4.0-ultra-generate-001")).toBe(true);
+        expect(isImagen("gemini-3-pro-image-preview")).toBe(false);
+        expect(isImagen("unknown")).toBe(false);
+    });
+});

package/extensions/gemini-image/models.ts

@@ -0,0 +1,50 @@
+export const GEMINI_NATIVE_MODELS = [
+    "gemini-2.5-flash-image",
+    "gemini-3.1-flash-image-preview",
+    "gemini-3-pro-image-preview",
+] as const;
+
+export const IMAGEN_MODELS = [
+    "imagen-4.0-fast-generate-001",
+    "imagen-4.0-generate-001",
+    "imagen-4.0-ultra-generate-001",
+] as const;
+
+export type GeminiNativeModel = (typeof GEMINI_NATIVE_MODELS)[number];
+export type ImagenModel = (typeof IMAGEN_MODELS)[number];
+export type ImageModel = GeminiNativeModel | ImagenModel;
+
+export const ALL_MODELS: readonly ImageModel[] = [
+    ...GEMINI_NATIVE_MODELS,
+    ...IMAGEN_MODELS,
+];
+
+export const DEFAULT_MODEL: ImageModel = "gemini-2.5-flash-image";
+
+export function isGeminiNative(model: string): model is GeminiNativeModel {
+    return (GEMINI_NATIVE_MODELS as readonly string[]).includes(model);
+}
+
+export function isImagen(model: string): model is ImagenModel {
+    return (IMAGEN_MODELS as readonly string[]).includes(model);
+}
+
+export const GEMINI_ASPECT_RATIOS = [
+    "1:1",
+    "2:3",
+    "3:2",
+    "3:4",
+    "4:3",
+    "4:5",
+    "5:4",
+    "9:16",
+    "16:9",
+    "21:9",
+] as const;
+
+export const IMAGEN_ASPECT_RATIOS = ["1:1", "3:4", "4:3", "9:16", "16:9"] as const;
+
+export const GEMINI_IMAGE_SIZES = ["1K", "2K", "4K"] as const;
+export const IMAGEN_IMAGE_SIZES = ["1K", "2K"] as const;
+
+export const PERSON_GENERATION_VALUES = ["dont_allow", "allow_adult", "allow_all"] as const;

package/extensions/gemini-image/pricing.test.ts

@@ -0,0 +1,95 @@
+import { describe, expect, it } from "vitest";
+import { estimateCost, formatUsd, hasPricing } from "./pricing.ts";
+
+describe("estimateCost", () => {
+    it("Imagen 4 Fast: $0.02 × N images", () => {
+        const e = estimateCost({
+            model: "imagen-4.0-fast-generate-001",
+            prompt: "cat",
+            numberOfImages: 3,
+        });
+        expect(e.estimatedUsd).toBeCloseTo(0.06, 5);
+        expect(e.breakdown).toContain("0.02");
+    });
+
+    it("Imagen 4 Standard single image", () => {
+        const e = estimateCost({
+            model: "imagen-4.0-generate-001",
+            prompt: "x",
+            numberOfImages: 1,
+        });
+        expect(e.estimatedUsd).toBeCloseTo(0.04, 5);
+    });
+
+    it("Imagen 4 Ultra", () => {
+        const e = estimateCost({
+            model: "imagen-4.0-ultra-generate-001",
+            prompt: "x",
+            numberOfImages: 4,
+        });
+        expect(e.estimatedUsd).toBeCloseTo(0.24, 5);
+    });
+
+    it("Gemini 3 Pro 4K is more expensive than 2K", () => {
+        const at4k = estimateCost({
+            model: "gemini-3-pro-image-preview",
+            prompt: "small",
+            numberOfImages: 1,
+            imageSize: "4K",
+        });
+        const at2k = estimateCost({
+            model: "gemini-3-pro-image-preview",
+            prompt: "small",
+            numberOfImages: 1,
+            imageSize: "2K",
+        });
+        expect(at4k.estimatedUsd).toBeGreaterThan(at2k.estimatedUsd);
+    });
+
+    it("Gemini 3 Pro prompt tokens contribute to cost (long prompt > short)", () => {
+        const longPrompt = "x".repeat(10_000);
+        const shortEst = estimateCost({
+            model: "gemini-3-pro-image-preview",
+            prompt: "x",
+            numberOfImages: 1,
+        });
+        const longEst = estimateCost({
+            model: "gemini-3-pro-image-preview",
+            prompt: longPrompt,
+            numberOfImages: 1,
+        });
+        expect(longEst.estimatedUsd).toBeGreaterThan(shortEst.estimatedUsd);
+    });
+
+    it("Flash estimate scales linearly with numberOfImages", () => {
+        const one = estimateCost({
+            model: "gemini-2.5-flash-image",
+            prompt: "x",
+            numberOfImages: 1,
+        });
+        const four = estimateCost({
+            model: "gemini-2.5-flash-image",
+            prompt: "x",
+            numberOfImages: 4,
+        });
+        // Output dominates; roughly 4x
+        expect(four.estimatedUsd).toBeGreaterThan(3.5 * one.estimatedUsd * 0.99);
+    });
+});
+
+describe("formatUsd", () => {
+    it("renders sub-cent as <$0.01", () => {
+        expect(formatUsd(0.002)).toBe("<$0.01");
+    });
+
+    it("renders 4 decimals for normal amounts", () => {
+        expect(formatUsd(0.1234)).toBe("$0.1234");
+    });
+});
+
+describe("hasPricing", () => {
+    it("returns true for all known models", () => {
+        expect(hasPricing("gemini-2.5-flash-image")).toBe(true);
+        expect(hasPricing("imagen-4.0-fast-generate-001")).toBe(true);
+    });
+});

package/extensions/gemini-image/pricing.ts

@@ -0,0 +1,102 @@
+/**
+ * Gemini image-generation price table.
+ *
+ * Last verified: 2026-04-24 via ai.google.dev/gemini-api/docs/pricing.
+ * Prices drift — edit when re-checked.
+ */
+
+import type { ImageModel } from "./models.ts";
+
+interface ImagenPrice {
+    kind: "flat";
+    perImage: number; // USD
+}
+
+interface GeminiNativePrice {
+    kind: "token";
+    // $ per 1M input tokens
+    inputPer1M: number;
+    // $ per output image at 1K/2K resolution
+    outputPerImage1Kor2K: number;
+    // $ per output image at 4K resolution
+    outputPerImage4K: number;
+    // Approx output tokens per image (for reference; not used in estimate)
+    tokensPerImage: number;
+}
+
+type ModelPrice = ImagenPrice | GeminiNativePrice;
+
+const PRICES: Record<ImageModel, ModelPrice> = {
+    "imagen-4.0-fast-generate-001": { kind: "flat", perImage: 0.02 },
+    "imagen-4.0-generate-001": { kind: "flat", perImage: 0.04 },
+    "imagen-4.0-ultra-generate-001": { kind: "flat", perImage: 0.06 },
+    "gemini-2.5-flash-image": {
+        kind: "token",
+        inputPer1M: 0.3,
+        outputPerImage1Kor2K: 0.039,
+        outputPerImage4K: 0.039,
+        tokensPerImage: 1290,
+    },
+    "gemini-3.1-flash-image-preview": {
+        kind: "token",
+        inputPer1M: 0.3,
+        outputPerImage1Kor2K: 0.039,
+        outputPerImage4K: 0.078,
+        tokensPerImage: 1290,
+    },
+    "gemini-3-pro-image-preview": {
+        kind: "token",
+        inputPer1M: 2.0,
+        outputPerImage1Kor2K: 0.134,
+        outputPerImage4K: 0.24,
+        tokensPerImage: 1290,
+    },
+};
+
+export interface EstimateInput {
+    model: ImageModel;
+    prompt: string;
+    numberOfImages: number;
+    imageSize?: "1K" | "2K" | "4K";
+    inputImagesBytes?: number;
+}
+
+export interface EstimateResult {
+    estimatedUsd: number;
+    breakdown: string;
+}
+
+const CHARS_PER_TOKEN_APPROX = 4;
+
+export function estimateCost(input: EstimateInput): EstimateResult {
+    const price = PRICES[input.model];
+    if (price.kind === "flat") {
+        const usd = price.perImage * input.numberOfImages;
+        return {
+            estimatedUsd: usd,
+            breakdown: `${input.numberOfImages} × $${price.perImage.toFixed(2)}/image`,
+        };
+    }
+
+    const inputTokens = Math.ceil(input.prompt.length / CHARS_PER_TOKEN_APPROX);
+    const inputCost = (inputTokens / 1_000_000) * price.inputPer1M;
+    const perImage =
+        input.imageSize === "4K" ? price.outputPerImage4K : price.outputPerImage1Kor2K;
+    const outputCost = perImage * input.numberOfImages;
+    const total = inputCost + outputCost;
+    return {
+        estimatedUsd: total,
+        breakdown:
+            `${input.numberOfImages} × $${perImage.toFixed(3)} (${input.imageSize ?? "1K/2K"}) + ` +
+            `~${inputTokens} input tokens`,
+    };
+}
+
+export function formatUsd(usd: number): string {
+    if (usd < 0.01) return `<$0.01`;
+    return `$${usd.toFixed(4)}`;
+}
+
+export function hasPricing(model: ImageModel): boolean {
+    return model in PRICES;
+}