@pi-unipi/image 2.2.0

package/src/settings.ts

@@ -0,0 +1,149 @@
+/**
+ * @pi-unipi/image — Settings
+ *
+ * Config lives at `~/.unipi/config/image/config.json`. Every read is
+ * try/catch-to-defaults so a corrupt file can never break the tools.
+ *
+ * The config directory is overridable via `UNIPI_IMAGE_CONFIG_DIR` so tests
+ * do not have to reach into the real home directory.
+ */
+
+import * as fs from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+
+/** Default system prompt for image recognition. */
+export const DEFAULT_RECOGNIZE_SYSTEM_PROMPT =
+  "You are a precise image analyst assisting a software engineer. " +
+  "Describe what is actually visible — never speculate about what is not shown. " +
+  "For screenshots, transcribe visible text, UI structure, and any errors verbatim. " +
+  "For diagrams, describe the components and their relationships. " +
+  "For photographs, describe the subject, setting, and notable detail. " +
+  "Be specific and concise; lead with the single most important observation.";
+
+/** Default image-generation model. */
+export const DEFAULT_GENERATE_MODEL = "openrouter/google/gemini-3-pro-image";
+
+export interface GenerateSettings {
+  /** Whether the image_generate tool is registered. */
+  enabled: boolean;
+  /** Model as "provider/model-id". */
+  model: string;
+  /** Directory for saved images. `~` is expanded. */
+  outputDir: string;
+  /** Whether to also write generated images to disk. */
+  saveToDisk: boolean;
+}
+
+export interface RecognizeSettings {
+  /** Whether the image_recognize tool is registered. */
+  enabled: boolean;
+  /** Model as "provider/model-id". Empty = use the session's current model. */
+  model: string;
+  /** System prompt sent with every recognition request. */
+  systemPrompt: string;
+}
+
+export interface ImageConfig {
+  generate: GenerateSettings;
+  recognize: RecognizeSettings;
+}
+
+export const DEFAULT_CONFIG: ImageConfig = {
+  generate: {
+    enabled: true,
+    model: DEFAULT_GENERATE_MODEL,
+    outputDir: "~/.unipi/images",
+    saveToDisk: true,
+  },
+  recognize: {
+    enabled: true,
+    model: "",
+    systemPrompt: DEFAULT_RECOGNIZE_SYSTEM_PROMPT,
+  },
+};
+
+/** Resolve the config directory, honouring the test override. */
+export function getConfigDir(): string {
+  const override = process.env.UNIPI_IMAGE_CONFIG_DIR;
+  if (override && override.trim().length > 0) return override;
+  return path.join(os.homedir(), ".unipi", "config", "image");
+}
+
+function getConfigPath(): string {
+  return path.join(getConfigDir(), "config.json");
+}
+
+/** Expand a leading `~` to the user's home directory. */
+export function expandHome(target: string): string {
+  if (target === "~") return os.homedir();
+  if (target.startsWith("~/") || target.startsWith("~\\")) {
+    return path.join(os.homedir(), target.slice(2));
+  }
+  return target;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return value !== null && typeof value === "object" && !Array.isArray(value);
+}
+
+/** Merge a loaded section over its defaults, ignoring wrong-typed fields. */
+function mergeSection<T extends object>(defaults: T, loaded: unknown): T {
+  if (!isRecord(loaded)) return { ...defaults };
+
+  const merged: T = { ...defaults };
+  for (const key of Object.keys(defaults) as Array<keyof T & string>) {
+    const value = loaded[key];
+    if (value === undefined || value === null) continue;
+    // Only accept a value whose type matches the default's.
+    if (typeof value === typeof defaults[key]) {
+      merged[key] = value as T[keyof T & string];
+    }
+  }
+  return merged;
+}
+
+/** Load config from disk, falling back to defaults on any problem. */
+export function loadConfig(): ImageConfig {
+  try {
+    const raw = fs.readFileSync(getConfigPath(), "utf-8");
+    const parsed: unknown = JSON.parse(raw);
+    if (!isRecord(parsed)) return structuredClone(DEFAULT_CONFIG);
+
+    return {
+      generate: mergeSection(DEFAULT_CONFIG.generate, parsed.generate),
+      recognize: mergeSection(DEFAULT_CONFIG.recognize, parsed.recognize),
+    };
+  } catch {
+    return structuredClone(DEFAULT_CONFIG);
+  }
+}
+
+/** Persist config. Returns false instead of throwing when the write fails. */
+export function saveConfig(config: ImageConfig): boolean {
+  try {
+    const dir = getConfigDir();
+    fs.mkdirSync(dir, { recursive: true });
+    fs.writeFileSync(getConfigPath(), `${JSON.stringify(config, null, 2)}\n`, "utf-8");
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+/** Apply a partial update, merging one level deep. */
+export function updateConfig(partial: Partial<ImageConfig>): ImageConfig {
+  const current = loadConfig();
+  const next: ImageConfig = {
+    generate: { ...current.generate, ...partial.generate },
+    recognize: { ...current.recognize, ...partial.recognize },
+  };
+  saveConfig(next);
+  return next;
+}
+
+/** Resolved absolute output directory for generated images. */
+export function getOutputDir(config: ImageConfig = loadConfig()): string {
+  const dir = config.generate.outputDir?.trim();
+  return expandHome(dir && dir.length > 0 ? dir : DEFAULT_CONFIG.generate.outputDir);
+}

package/src/tools.ts

@@ -0,0 +1,296 @@
+/**
+ * @pi-unipi/image — Agent tool registration
+ *
+ * Registers `image_generate` and `image_recognize`. Each tool is registered
+ * only when enabled in config, so a user who wants just one does not have the
+ * other consuming context in the system prompt.
+ */
+
+import { Type } from "typebox";
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { IMAGE_TOOLS } from "@pi-unipi/core";
+
+import { generateImage } from "./generate.js";
+import { loadImage } from "./image-source.js";
+import {
+  formatModelRef,
+  listImageGenModels,
+  resolveImageGenModel,
+  resolveVisionModel,
+  splitModelRef,
+  type ChatModelRegistry,
+  type VisionModel,
+} from "./models.js";
+import { recognizeImage } from "./recognize.js";
+import { getOutputDir, loadConfig } from "./settings.js";
+
+/** Error shape shared by both tools. */
+function errorResult(message: string) {
+  return {
+    content: [{ type: "text" as const, text: message }],
+    isError: true,
+    details: {},
+  };
+}
+
+function messageOf(error: unknown): string {
+  return error instanceof Error ? error.message : String(error);
+}
+
+/** Read the chat model registry off the extension context. */
+function getRegistry(ctx: ExtensionContext): ChatModelRegistry | undefined {
+  return (ctx as unknown as { modelRegistry?: ChatModelRegistry }).modelRegistry;
+}
+
+/** Resolve an API key for a provider through pi's auth storage. */
+async function resolveApiKey(
+  registry: ChatModelRegistry | undefined,
+  provider: string,
+): Promise<string | undefined> {
+  try {
+    const key = await registry?.getApiKeyForProvider?.(provider);
+    if (key) return key;
+  } catch {
+    // Fall through to the environment.
+  }
+
+  const envName = `${provider.toUpperCase().replace(/[^A-Z0-9]/g, "_")}_API_KEY`;
+  return process.env[envName] || undefined;
+}
+
+/** Look up the full pi-ai model object for a vision model. */
+function findChatModel(
+  registry: ChatModelRegistry,
+  model: VisionModel,
+): { baseUrl?: string; api?: string } | undefined {
+  try {
+    return registry.find(model.provider, model.id) as
+      | { baseUrl?: string; api?: string }
+      | undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+export function registerImageTools(pi: ExtensionAPI): void {
+  const config = loadConfig();
+
+  if (config.generate.enabled) registerGenerateTool(pi);
+  if (config.recognize.enabled) registerRecognizeTool(pi);
+}
+
+function registerGenerateTool(pi: ExtensionAPI): void {
+  pi.registerTool({
+    name: IMAGE_TOOLS.GENERATE,
+    label: "Generate Image",
+    description:
+      "Generate an image from a text prompt using an image model. " +
+      "The image is returned inline and, when enabled, saved to disk.",
+    promptSnippet: "Generate an image from a text prompt.",
+    promptGuidelines: [
+      "Use image_generate to create images from a text description.",
+      "Write a detailed prompt — subject, style, composition and lighting all help.",
+      "Omit model to use the one configured in /unipi:image-settings.",
+      "Generated images cost money per call; do not regenerate without being asked.",
+    ],
+    parameters: Type.Object({
+      prompt: Type.String({
+        description: "Description of the image to generate. Be specific.",
+      }),
+      model: Type.Optional(
+        Type.String({
+          description:
+            'Image model override, e.g. "flux.2-pro" or ' +
+            '"openrouter/google/gemini-3-pro-image". Omit to use the configured default.',
+        }),
+      ),
+    }),
+    async execute(_toolCallId, params, signal, _onUpdate, ctx) {
+      try {
+        const config = loadConfig();
+        const models = await listImageGenModels();
+
+        const requested = params.model?.trim() || config.generate.model;
+        const resolved = resolveImageGenModel(requested, models);
+        if (typeof resolved === "string") return errorResult(resolved);
+
+        // pi-ai resolves image auth from its own credential store; only fall
+        // back to pi's chat-provider key when that comes up empty.
+        const registry = getRegistry(ctx);
+        const fallbackKey = await resolveApiKey(registry, resolved.provider);
+
+        const result = await generateImage({
+          prompt: params.prompt,
+          model: resolved,
+          ...(fallbackKey ? { apiKey: fallbackKey } : {}),
+          signal,
+          outputDir: config.generate.saveToDisk ? getOutputDir(config) : undefined,
+        });
+
+        const saved = result.images
+          .map((image) => image.path)
+          .filter((path): path is string => Boolean(path));
+
+        const summary = [
+          `Generated ${result.images.length} image${result.images.length === 1 ? "" : "s"} ` +
+            `with ${formatModelRef(resolved)}.`,
+          saved.length > 0 ? `Saved to:\n${saved.map((p) => `  ${p}`).join("\n")}` : "",
+          config.generate.saveToDisk && saved.length === 0
+            ? "Could not write to the output directory — returning the image inline only."
+            : "",
+          result.text,
+        ]
+          .filter(Boolean)
+          .join("\n");
+
+        return {
+          content: [
+            { type: "text" as const, text: summary },
+            ...result.images.map((image) => ({
+              type: "image" as const,
+              data: image.data,
+              mimeType: image.mimeType,
+            })),
+          ],
+          details: {
+            model: formatModelRef(resolved),
+            count: result.images.length,
+            paths: saved,
+          },
+        };
+      } catch (error) {
+        return errorResult(`Image generation failed: ${messageOf(error)}`);
+      }
+    },
+  });
+}
+
+function registerRecognizeTool(pi: ExtensionAPI): void {
+  pi.registerTool({
+    name: IMAGE_TOOLS.RECOGNIZE,
+    label: "Recognize Image",
+    description:
+      "Analyze an image and answer questions about it using a vision model. " +
+      "Accepts a local file path, a data: URL, or base64 image data.",
+    promptSnippet: "Analyze an image and answer questions about it.",
+    promptGuidelines: [
+      "Use image_recognize to read screenshots, diagrams, mockups and photos.",
+      "Pass a local file path whenever possible — it is cheaper than inlining base64.",
+      "Ask a specific question in `prompt` to focus the analysis.",
+      "Remote URLs are not fetched; download the image first.",
+    ],
+    parameters: Type.Object({
+      image: Type.String({
+        description:
+          "Local file path, data: URL, or base64 image data. " +
+          "Supported types: PNG, JPEG, GIF, WebP.",
+      }),
+      prompt: Type.Optional(
+        Type.String({
+          description:
+            "What to ask about the image. Defaults to a general description.",
+        }),
+      ),
+      model: Type.Optional(
+        Type.String({
+          description:
+            "Vision model override. Must accept image input. " +
+            "Omit to use the configured default or the current session model.",
+        }),
+      ),
+      systemPrompt: Type.Optional(
+        Type.String({
+          description:
+            "Override the configured system prompt for this call only.",
+        }),
+      ),
+    }),
+    async execute(_toolCallId, params, signal, _onUpdate, ctx) {
+      try {
+        const config = loadConfig();
+
+        const registry = getRegistry(ctx);
+        if (!registry) {
+          return errorResult(
+            "Model registry unavailable — image_recognize needs an active session.",
+          );
+        }
+
+        // Precedence: per-call override, configured model, session model.
+        const requested =
+          params.model?.trim() ||
+          config.recognize.model.trim() ||
+          currentSessionModel(ctx);
+
+        if (!requested) {
+          return errorResult(
+            "No vision model configured.\n" +
+              "→ Choose one with /unipi:image-settings, or pass `model`.",
+          );
+        }
+
+        const resolved = resolveVisionModel(requested, registry);
+        if (typeof resolved === "string") return errorResult(resolved);
+
+        const image = loadImage(params.image, ctx.cwd ?? process.cwd());
+
+        const chatModel = findChatModel(registry, resolved);
+        const baseUrl = chatModel?.baseUrl;
+        if (!baseUrl) {
+          return errorResult(
+            `Could not determine the API endpoint for ${formatModelRef(resolved)}.`,
+          );
+        }
+
+        const apiKey = await resolveApiKey(registry, resolved.provider);
+        if (!apiKey) {
+          return errorResult(
+            `No API key for provider "${resolved.provider}".\n` +
+              "→ Sign in with /login, or set the provider's API key environment variable.",
+          );
+        }
+
+        const result = await recognizeImage({
+          image,
+          prompt: params.prompt?.trim() || "Describe this image in detail.",
+          systemPrompt: params.systemPrompt?.trim() || config.recognize.systemPrompt,
+          apiKey,
+          baseUrl,
+          api: chatModel?.api ?? "openai-completions",
+          modelId: resolved.id,
+          signal,
+        });
+
+        const origin = image.path ? ` (${image.path})` : "";
+
+        return {
+          content: [
+            {
+              type: "text" as const,
+              text: `${result.text}\n\n— analyzed with ${formatModelRef(resolved)}${origin}`,
+            },
+          ],
+          details: {
+            model: formatModelRef(resolved),
+            mimeType: image.mimeType,
+            source: image.source,
+            ...(image.path ? { path: image.path } : {}),
+          },
+        };
+      } catch (error) {
+        return errorResult(`Image recognition failed: ${messageOf(error)}`);
+      }
+    },
+  });
+}
+
+/** The session's current model as "provider/id", when discoverable. */
+function currentSessionModel(ctx: ExtensionContext): string {
+  const model = (ctx as unknown as { model?: { provider?: string; id?: string } }).model;
+  if (model?.provider && model?.id) {
+    return formatModelRef({ provider: model.provider, id: model.id });
+  }
+  return "";
+}
+
+export { splitModelRef };