@vellumai/assistant 0.6.4 → 0.6.5

package/src/cli/commands/image-generation.ts

@@ -0,0 +1,300 @@
+import { existsSync, mkdirSync, writeFileSync } from "node:fs";
+import os from "node:os";
+import { join } from "node:path";
+
+import { Command } from "commander";
+
+import { getConfig } from "../../config/loader.js";
+import {
+  generateImage,
+  type ImageGenCredentials,
+  mapGeminiError,
+} from "../../media/gemini-image-service.js";
+import {
+  buildManagedBaseUrl,
+  resolveManagedProxyContext,
+} from "../../providers/managed-proxy/context.js";
+import { getProviderKeyAsync } from "../../security/secure-keys.js";
+import { log } from "../logger.js";
+
+// ---------------------------------------------------------------------------
+// MIME type → file extension mapping
+// ---------------------------------------------------------------------------
+
+function extensionForMime(mimeType: string): string {
+  switch (mimeType) {
+    case "image/png":
+      return "png";
+    case "image/jpeg":
+    case "image/jpg":
+      return "jpg";
+    case "image/webp":
+      return "webp";
+    case "image/gif":
+      return "gif";
+    default:
+      return "png";
+  }
+}
+
+// ---------------------------------------------------------------------------
+// MIME type from file extension (for source images)
+// ---------------------------------------------------------------------------
+
+function mimeForExtension(filePath: string): string {
+  const ext = filePath.split(".").pop()?.toLowerCase();
+  switch (ext) {
+    case "png":
+      return "image/png";
+    case "jpg":
+    case "jpeg":
+      return "image/jpeg";
+    case "webp":
+      return "image/webp";
+    case "gif":
+      return "image/gif";
+    default:
+      return "application/octet-stream";
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Command registration
+// ---------------------------------------------------------------------------
+
+export function registerImageGenerationCommand(program: Command): void {
+  const imageGen = program
+    .command("image-generation")
+    .description("AI image generation and editing");
+
+  imageGen.addHelpText(
+    "after",
+    `
+Modes:
+  managed    — Uses platform-managed credentials (requires login to Vellum).
+  your-own   — Uses your own Gemini API key configured in settings.
+
+Supported models:
+  gemini-3.1-flash-image-preview (default)
+  gemini-3-pro-image-preview
+
+Examples:
+  $ assistant image-generation generate --prompt "A sunset over the ocean"
+  $ assistant image-generation generate --prompt "Remove background" --mode edit --source photo.png
+  $ assistant image-generation generate --prompt "Logo design" --variants 3 --output-dir ./output
+  $ assistant image-generation generate --prompt "A cat" --json`,
+  );
+
+  const generate = imageGen
+    .command("generate")
+    .description("Generate or edit images using AI")
+    .requiredOption(
+      "--prompt <text>",
+      "Description of the image to generate or edits to apply",
+    )
+    .option("--mode <mode>", "generate (default) or edit", "generate")
+    .option(
+      "--source <path...>",
+      "Source image file path for edit mode (repeatable)",
+    )
+    .option("--model <model-id>", "Model override")
+    .option(
+      "--variants <n>",
+      "Number of variants (1-4, default 1)",
+      (v: string) => parseInt(v, 10),
+      1,
+    )
+    .option("--output-dir <dir>", "Directory to save images")
+    .option("--json", "Output structured JSON");
+
+  generate.addHelpText(
+    "after",
+    `
+Notes:
+  Edit mode (--mode edit) requires at least one --source image file.
+  Output files are named image-1.png, image-2.png, etc. (extension matches MIME type).
+  Default output directory is the system temp directory.
+
+Examples:
+  $ assistant image-generation generate --prompt "A mountain landscape at dawn"
+  $ assistant image-generation generate --prompt "Make it darker" --mode edit --source input.png
+  $ assistant image-generation generate --prompt "Logo variations" --variants 4 --output-dir ./logos
+  $ assistant image-generation generate --prompt "A robot" --model gemini-3-pro-image-preview --json`,
+  );
+
+  generate.action(async (opts) => {
+    const jsonOutput = opts.json === true;
+    const prompt: string = opts.prompt;
+    const mode: "generate" | "edit" =
+      opts.mode === "edit" ? "edit" : "generate";
+    const sourcePaths: string[] | undefined = opts.source;
+    const modelOverride: string | undefined = opts.model;
+    const rawVariants = opts.variants ?? 1;
+    const variants: number = Number.isNaN(rawVariants)
+      ? 1
+      : Math.max(1, Math.min(rawVariants, 4));
+    const outputDir: string = opts.outputDir ?? os.tmpdir();
+
+    // --- Validate edit mode requires --source ---
+    if (mode === "edit" && (!sourcePaths || sourcePaths.length === 0)) {
+      const msg = "Edit mode requires at least one --source image file.";
+      if (jsonOutput) {
+        process.stdout.write(JSON.stringify({ ok: false, error: msg }) + "\n");
+      } else {
+        log.error(msg);
+      }
+      process.exitCode = 1;
+      return;
+    }
+
+    // --- Resolve credentials ---
+    const config = getConfig();
+    const imageGenMode = config.services["image-generation"].mode;
+
+    let credentials: ImageGenCredentials | undefined;
+
+    if (imageGenMode === "managed") {
+      const managedBaseUrl = await buildManagedBaseUrl("gemini");
+      if (managedBaseUrl) {
+        const ctx = await resolveManagedProxyContext();
+        credentials = {
+          type: "managed-proxy",
+          assistantApiKey: ctx.assistantApiKey,
+          baseUrl: managedBaseUrl,
+        };
+      }
+    } else {
+      const apiKey = await getProviderKeyAsync("gemini");
+      if (apiKey) {
+        credentials = { type: "direct", apiKey };
+      }
+    }
+
+    if (!credentials) {
+      const hint =
+        imageGenMode === "managed"
+          ? "Managed proxy is not available. Please log in to Vellum or switch to your-own mode:\n  Run 'assistant auth login' to authenticate, or set services.image-generation.mode to 'your-own' in config."
+          : "No Gemini API key configured. Add your Gemini API key:\n  Run 'assistant keys set gemini' or configure it in Settings > Models & Services.";
+
+      if (jsonOutput) {
+        process.stdout.write(JSON.stringify({ ok: false, error: hint }) + "\n");
+      } else {
+        log.error(hint);
+      }
+      process.exitCode = 1;
+      return;
+    }
+
+    // --- Read source images for edit mode ---
+    let sourceImages:
+      | Array<{ mimeType: string; dataBase64: string }>
+      | undefined;
+
+    if (mode === "edit" && sourcePaths && sourcePaths.length > 0) {
+      const errors: string[] = [];
+      const validImages: Array<{ mimeType: string; dataBase64: string }> = [];
+
+      for (const filePath of sourcePaths) {
+        if (!existsSync(filePath)) {
+          errors.push(`File not found: ${filePath}`);
+          continue;
+        }
+        try {
+          const file = Bun.file(filePath);
+          const buffer = Buffer.from(await file.arrayBuffer());
+          const mimeType =
+            file.type !== "application/octet-stream"
+              ? file.type
+              : mimeForExtension(filePath);
+          validImages.push({
+            mimeType,
+            dataBase64: buffer.toString("base64"),
+          });
+        } catch (err) {
+          errors.push(`Could not read ${filePath}: ${(err as Error).message}`);
+        }
+      }
+
+      if (validImages.length === 0) {
+        const errorMsg = `No source images could be read.\n${errors.join("\n")}`;
+        if (jsonOutput) {
+          process.stdout.write(
+            JSON.stringify({ ok: false, error: errorMsg }) + "\n",
+          );
+        } else {
+          log.error(errorMsg);
+        }
+        process.exitCode = 1;
+        return;
+      }
+      sourceImages = validImages;
+    }
+
+    // --- Resolve model ---
+    const model = modelOverride ?? config.services["image-generation"].model;
+
+    // --- Generate image ---
+    try {
+      const result = await generateImage(credentials, {
+        prompt,
+        mode,
+        sourceImages,
+        model,
+        variants,
+      });
+
+      // --- Ensure output directory exists ---
+      if (!existsSync(outputDir)) {
+        mkdirSync(outputDir, { recursive: true });
+      }
+
+      // --- Write images to disk ---
+      const imageOutputs: Array<{
+        path: string;
+        mimeType: string;
+        sizeBytes: number;
+      }> = [];
+
+      for (let i = 0; i < result.images.length; i++) {
+        const img = result.images[i];
+        const ext = extensionForMime(img.mimeType);
+        const fileName = `image-${i + 1}.${ext}`;
+        const filePath = join(outputDir, fileName);
+        const buffer = Buffer.from(img.dataBase64, "base64");
+        writeFileSync(filePath, buffer);
+        imageOutputs.push({
+          path: filePath,
+          mimeType: img.mimeType,
+          sizeBytes: buffer.length,
+        });
+      }
+
+      // --- Output ---
+      if (jsonOutput) {
+        const output: Record<string, unknown> = {
+          ok: true,
+          images: imageOutputs,
+          model: result.resolvedModel,
+        };
+        if (result.text) {
+          output.text = result.text;
+        }
+        process.stdout.write(JSON.stringify(output) + "\n");
+      } else {
+        for (const img of imageOutputs) {
+          process.stdout.write(img.path + "\n");
+        }
+      }
+    } catch (error) {
+      const errorMsg = mapGeminiError(error);
+      if (jsonOutput) {
+        process.stdout.write(
+          JSON.stringify({ ok: false, error: errorMsg }) + "\n",
+        );
+      } else {
+        log.error(errorMsg);
+      }
+      process.exitCode = 1;
+    }
+  });
+}

package/src/cli/commands/inference.ts

@@ -0,0 +1,200 @@
+import { readFileSync } from "node:fs";
+
+import type { Command } from "commander";
+
+import {
+  extractAllText,
+  getConfiguredProvider,
+  userMessage,
+} from "../../providers/provider-send-message.js";
+import { log } from "../logger.js";
+
+/**
+ * Attach the `send` subcommand to the given command group (`inference` or
+ * `llm`). Both groups share the same implementation.
+ */
+function attachSendSubcommand(group: Command): void {
+  group
+    .command("send")
+    .description("Send a message to the configured LLM and print the response")
+    .option("--system-prompt <text>", "System prompt for the model")
+    .option("--model <model-id>", "Model override")
+    .option("--max-tokens <n>", "Max response tokens", undefined)
+    .option("--json", "Output structured JSON")
+    .argument("[message...]", "User message (joined with spaces)")
+    .addHelpText(
+      "after",
+      `
+Behavioral notes:
+  - If no message argument is provided, reads from stdin.
+  - If --model is omitted, uses the configured default model.
+  - Requires a configured LLM provider (see 'assistant config set').
+
+Examples:
+  $ assistant inference send "What is 2+2?"
+  $ echo "Summarize this" | assistant inference send
+  $ assistant llm send --system-prompt "You are a poet" "Write a haiku"
+  $ assistant inference send --model claude-sonnet-4-20250514 --json "Hello"`,
+    )
+    .action(
+      async (
+        messageParts: string[],
+        opts: {
+          systemPrompt?: string;
+          model?: string;
+          maxTokens?: string;
+          json?: boolean;
+        },
+      ) => {
+        const { systemPrompt, model, json: jsonOutput } = opts;
+        const maxTokens = opts.maxTokens
+          ? parseInt(opts.maxTokens, 10)
+          : undefined;
+
+        if (
+          opts.maxTokens !== undefined &&
+          (!Number.isFinite(maxTokens) || maxTokens! < 1)
+        ) {
+          const msg = "Invalid --max-tokens value. Must be a positive integer.";
+          if (jsonOutput) {
+            process.stdout.write(
+              JSON.stringify({ ok: false, error: msg }) + "\n",
+            );
+          } else {
+            log.error(msg);
+          }
+          process.exitCode = 1;
+          return;
+        }
+
+        // Determine user message: positional args or stdin.
+        let messageText = messageParts.length > 0 ? messageParts.join(" ") : "";
+
+        if (!messageText && !process.stdin.isTTY) {
+          try {
+            messageText = readFileSync("/dev/stdin", "utf-8").trim();
+          } catch {
+            // stdin not available or empty
+          }
+        }
+
+        if (!messageText) {
+          const msg =
+            "No message provided. Pass a message as an argument or pipe via stdin.";
+          if (jsonOutput) {
+            process.stdout.write(
+              JSON.stringify({ ok: false, error: msg }) + "\n",
+            );
+          } else {
+            log.error(msg);
+          }
+          process.exitCode = 1;
+          return;
+        }
+
+        // Resolve provider.
+        const provider = await getConfiguredProvider("inference");
+        if (!provider) {
+          const msg =
+            "No LLM provider is configured. Run 'assistant config set llm.default.provider <provider>' to set one up.";
+          if (jsonOutput) {
+            process.stdout.write(
+              JSON.stringify({ ok: false, error: msg }) + "\n",
+            );
+          } else {
+            log.error(msg);
+          }
+          process.exitCode = 1;
+          return;
+        }
+
+        try {
+          const response = await provider.sendMessage(
+            [userMessage(messageText)],
+            undefined,
+            systemPrompt,
+            {
+              config: {
+                callSite: "inference",
+                max_tokens: maxTokens,
+                model,
+              },
+            },
+          );
+
+          const text = extractAllText(response);
+
+          if (jsonOutput) {
+            process.stdout.write(
+              JSON.stringify({
+                ok: true,
+                response: text,
+                model: response.model,
+                usage: {
+                  inputTokens: response.usage.inputTokens,
+                  outputTokens: response.usage.outputTokens,
+                },
+              }) + "\n",
+            );
+          } else {
+            process.stdout.write(text + "\n");
+          }
+        } catch (err) {
+          const msg =
+            err instanceof Error ? err.message : "Unknown error occurred";
+          if (jsonOutput) {
+            process.stdout.write(
+              JSON.stringify({ ok: false, error: msg }) + "\n",
+            );
+          } else {
+            log.error(msg);
+          }
+          process.exitCode = 1;
+        }
+      },
+    );
+}
+
+/**
+ * Register `inference` and `llm` command groups on the top-level program.
+ * Both expose the same subcommands — `llm` is an alias for `inference`.
+ */
+export function registerInferenceCommand(program: Command): void {
+  const inference = program
+    .command("inference")
+    .description("LLM inference operations");
+
+  inference.addHelpText(
+    "after",
+    `
+The inference command group sends requests to your configured LLM provider.
+The provider is resolved from your assistant config (llm.default.provider).
+
+Examples:
+  $ assistant inference send "What is the capital of France?"
+  $ echo "Explain quantum computing" | assistant inference send
+  $ assistant llm send --system-prompt "Be concise" "What is TCP?"
+  $ assistant inference send --model claude-sonnet-4-20250514 --json "Hello"`,
+  );
+
+  attachSendSubcommand(inference);
+
+  const llm = program
+    .command("llm")
+    .description("LLM inference operations (alias for 'inference')");
+
+  llm.addHelpText(
+    "after",
+    `
+The llm command group is an alias for 'inference'. It sends requests to your
+configured LLM provider, resolved from your assistant config (llm.default.provider).
+
+Examples:
+  $ assistant llm send "What is the capital of France?"
+  $ echo "Explain quantum computing" | assistant llm send
+  $ assistant llm send --system-prompt "Be concise" "What is TCP?"
+  $ assistant llm send --model claude-sonnet-4-20250514 --json "Hello"`,
+  );
+
+  attachSendSubcommand(llm);
+}