|
| 1 | +import { Command, InvalidArgumentError } from "commander"; |
| 2 | +import type { components } from "@shotstack/schemas"; |
| 3 | +import { createClient, type Client } from "../http/client.ts"; |
| 4 | +import { requireApiKey } from "../http/auth.ts"; |
| 5 | +import { resolveEnv, ENV_NAMES } from "../http/env.ts"; |
| 6 | +import { emit, type OutputFormat, parseOutputFormat } from "../output.ts"; |
| 7 | +import { withRecording, commandArgv, type CommandResult } from "../recorder.ts"; |
| 8 | +import { fetchModel, formatAvailability, type GenerationModel } from "./models.ts"; |
| 9 | + |
| 10 | +export type Generation = components["schemas"]["GenerationResponse"]; |
| 11 | + |
| 12 | +export interface GenerateInput { |
| 13 | + model: string; |
| 14 | + prompt: string; |
| 15 | + options?: Record<string, unknown>; |
| 16 | + length?: number; |
| 17 | +} |
| 18 | + |
| 19 | +const TERMINAL_STATES = new Set(["done", "failed"]); |
| 20 | +const POLL_INTERVAL_MS = 3000; |
| 21 | + |
| 22 | +export const generateCommand = new Command("generate") |
| 23 | + .description("Generate one image, video or audio asset from a prompt") |
| 24 | + .argument("<model>", "Model id from `shotstack models`") |
| 25 | + .argument("<prompt>", "Text prompt; for a speech model, the words to speak") |
| 26 | + .option("--options <json>", "Model options as a JSON object (see `shotstack models <model>`)", parseOptionsJson) |
| 27 | + .option("--length <seconds>", "Length of the clip the asset fills; models that generate to a duration use it", parseSeconds) |
| 28 | + .option("--watch", "Poll until the asset is ready, then print its URL") |
| 29 | + .option(`--env <name>`, `Environment: ${ENV_NAMES.join(" | ")}`) |
| 30 | + .option("--output <format>", "Output format: text | json", "text") |
| 31 | + .action( |
| 32 | + async ( |
| 33 | + model: string, |
| 34 | + prompt: string, |
| 35 | + options: { options?: Record<string, unknown>; length?: number; watch?: boolean; env?: string; output: string }, |
| 36 | + ) => { |
| 37 | + await withRecording("generate", commandArgv("generate"), async () => { |
| 38 | + const format = parseOutputFormat(options.output); |
| 39 | + const env = resolveEnv(options.env); |
| 40 | + const apiKey = requireApiKey(env.name); |
| 41 | + const client = createClient({ apiKey, env }); |
| 42 | + const input = { model, prompt, options: options.options, length: options.length }; |
| 43 | + return runGenerate(client, input, format, options.watch === true); |
| 44 | + }); |
| 45 | + }, |
| 46 | + ); |
| 47 | + |
| 48 | +export async function runGenerate( |
| 49 | + client: Client, |
| 50 | + input: GenerateInput, |
| 51 | + format: OutputFormat, |
| 52 | + watch: boolean, |
| 53 | + intervalMs: number = POLL_INTERVAL_MS, |
| 54 | +): Promise<CommandResult> { |
| 55 | + // The model's type is needed for the asset, and an unavailable model is refused before anything is billed. |
| 56 | + const model = await fetchModel(client, input.model); |
| 57 | + if (model.available === false) { |
| 58 | + const refusal = { model: model.model, available: false, unavailableReason: model.unavailableReason }; |
| 59 | + if (format === "json") console.log(JSON.stringify(refusal)); |
| 60 | + else console.error(`✗ ${model.model} is ${formatAvailability(model)}`); |
| 61 | + return { response: refusal, exitCode: 1 }; |
| 62 | + } |
| 63 | + |
| 64 | + // No Idempotency-Key: without one, identical requests share a job and its cached result. |
| 65 | + const submitted = await client.post<Generation>("/generate", buildRequest(model, input)); |
| 66 | + emit(format, submitted, formatHuman(submitted)); |
| 67 | + |
| 68 | + let final = submitted; |
| 69 | + if (!TERMINAL_STATES.has(submitted.status)) { |
| 70 | + if (watch) final = await pollGeneration(client, submitted.id, format, intervalMs); |
| 71 | + else if (format !== "json") { |
| 72 | + console.error("Still generating. Run the same command with --watch to wait; the same request returns this job instead of starting another."); |
| 73 | + } |
| 74 | + } |
| 75 | + return { renderId: final.id, response: final, exitCode: final.status === "failed" ? 1 : 0 }; |
| 76 | +} |
| 77 | + |
| 78 | +export function buildRequest(model: GenerationModel, input: GenerateInput) { |
| 79 | + return { |
| 80 | + asset: { |
| 81 | + type: model.type, |
| 82 | + prompt: input.prompt, |
| 83 | + model: model.model, |
| 84 | + ...(input.options && { options: input.options }), |
| 85 | + }, |
| 86 | + ...(input.length !== undefined && { length: input.length }), |
| 87 | + }; |
| 88 | +} |
| 89 | + |
| 90 | +export async function pollGeneration( |
| 91 | + client: Client, |
| 92 | + id: string, |
| 93 | + format: OutputFormat, |
| 94 | + intervalMs: number = POLL_INTERVAL_MS, |
| 95 | +): Promise<Generation> { |
| 96 | + while (true) { |
| 97 | + await sleep(intervalMs); |
| 98 | + const generation = await client.get<Generation>(`/generate/${encodeURIComponent(id)}`); |
| 99 | + emit(format, generation, formatHuman(generation)); |
| 100 | + if (TERMINAL_STATES.has(generation.status)) return generation; |
| 101 | + } |
| 102 | +} |
| 103 | + |
| 104 | +export function parseOptionsJson(value: string): Record<string, unknown> { |
| 105 | + let parsed: unknown; |
| 106 | + try { |
| 107 | + parsed = JSON.parse(value); |
| 108 | + } catch { |
| 109 | + throw new InvalidArgumentError("Not valid JSON."); |
| 110 | + } |
| 111 | + if (parsed === null || typeof parsed !== "object" || Array.isArray(parsed)) { |
| 112 | + throw new InvalidArgumentError("Expected a JSON object."); |
| 113 | + } |
| 114 | + return parsed as Record<string, unknown>; |
| 115 | +} |
| 116 | + |
| 117 | +// Checked here because NaN serialises to null, which the API reads as no length and bills the model's default. |
| 118 | +export function parseSeconds(value: string): number { |
| 119 | + const seconds = Number(value); |
| 120 | + if (!Number.isFinite(seconds) || seconds <= 0) throw new InvalidArgumentError("Expected a positive number of seconds."); |
| 121 | + return seconds; |
| 122 | +} |
| 123 | + |
| 124 | +function formatHuman(g: Generation): string { |
| 125 | + if (g.status === "done") return `done ${g.url ?? ""}`.trimEnd(); |
| 126 | + if (g.status === "failed") return `failed error: ${g.error ?? "unknown error"}`; |
| 127 | + return `${g.status} ${g.id}`; |
| 128 | +} |
| 129 | + |
| 130 | +function sleep(ms: number): Promise<void> { |
| 131 | + return new Promise((resolve) => setTimeout(resolve, ms)); |
| 132 | +} |
0 commit comments