pi-asr/extensions/asr.ts

/**
 * pi ASR extension — speech-to-text through the NVIDIA Speech NIM realtime WebSocket API.
 *
 * Tools:
 *   asr_transcribe   transcribe an audio/video file (wav, mp3, m4a, aac, ogg, flac, opus, mp4, mov, mkv)
 *   asr_health       pre-flight: health/live + health/ready + the model the NIM actually advertises
 * Command:
 *   /asr <file> [--lang xx-YY] [--words] [--speed N] [--start S] [--out transcript.txt]
 *
 * Config via env (built-in defaults otherwise):
 *   ASR_NIM_URL   https://ai.iaminsane.online/openai/streaming-asr
 *   ASR_API_KEY   key sent as `Authorization: Bearer ...` on every HTTPS/WS call
 */

import { existsSync, statSync } from "node:fs";
import { mkdir, writeFile } from "node:fs/promises";
import { dirname, isAbsolute, relative, resolve } from "node:path";
import type { AgentToolResult, ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
import { Text } from "@mariozechner/pi-tui";
import { Type } from "@sinclair/typebox";
import {
	DEFAULT_BASE_URL, RealtimeAsrClient, hasFfmpeg, loadPcm16,
	type TranscriptResult, type WordInfo,
} from "../src/client.js";

interface TranscribeParams {
	path: string;
	language?: string;
	model?: string;
	words?: boolean;
	punctuation?: boolean;
	start?: number;
	speed?: number;
	pace?: boolean;
	out?: string;
}

interface AsrDetails {
	path: string;
	text: string;
	seconds: number;
	wallMs: number;
	model: string | null;
	sessionId: string;
	wordCount: number;
	segments: { transcript: string; isLast: boolean }[];
	words?: WordInfo[];
	outPath?: string;
	partial?: string;
}

const wordCount = (t: string): number => (t.trim() ? t.trim().split(/\s+/).length : 0);
const rel = (p: string, cwd: string): string => {
	const r = relative(cwd, p);
	return r && !r.startsWith("..") && !isAbsolute(r) ? r : p;
};

export default function asrExtension(pi: ExtensionAPI) {
	async function transcribeFile(
		rawPath: string,
		params: TranscribeParams,
		cwd: string,
		signal: AbortSignal | undefined,
		onDelta?: (t: string) => void,
	): Promise<{ result: TranscriptResult; wallMs: number; absPath: string }> {
		const absPath = resolve(cwd, rawPath.replace(/^@+/, ""));
		if (!existsSync(absPath)) throw new Error(`no such file: ${absPath}`);
		if (statSync(absPath).isDirectory()) throw new Error(`${absPath} is a directory, expected an audio file`);

		const client = new RealtimeAsrClient({
			language: params.language ?? "en-US",
			model: params.model || null,
			punctuation: params.punctuation ?? true,
			wordOffsets: params.words ?? false,
		});

		const t0 = Date.now();
		const pcm = await loadPcm16(absPath, {
			sampleRate: 16000, channels: 1,
			start: params.start ?? 0, speed: params.speed ?? 1,
		});
		if (!pcm.length) throw new Error(`decoded ${absPath} to zero bytes of audio (corrupt or empty track?)`);

		const result = await client.transcribe(pcm, { pace: params.pace ?? true, frameMs: 100, onDelta, signal });
		if (signal?.aborted) throw new Error("transcription cancelled");
		if (result.failed.length) {
			throw new Error(`server reported a transcription failure: ${JSON.stringify(result.failed[0]).slice(0, 300)}`);
		}
		return { result, wallMs: Date.now() - t0, absPath };
	}

	async function runTranscribe(
		params: TranscribeParams,
		ctx: ExtensionContext,
		onUpdate?: (r: AgentToolResult<AsrDetails>) => void,
		signal?: AbortSignal,
	): Promise<AgentToolResult<AsrDetails>> {
		if (!params.path?.trim()) throw new Error("`path` is required");
		if (signal?.aborted) throw new Error("transcription cancelled");
		let lastPartial = "";
		const started = Date.now();

		const { result, wallMs, absPath } = await transcribeFile(params.path.trim(), params, ctx.cwd, signal, (t) => {
			if (!t || t === lastPartial) return;
			lastPartial = t;
			onUpdate?.({
				content: [{ type: "text", text: t }],
				details: {
					path: rel(absPath, ctx.cwd), text: t, seconds: 0, wallMs: Date.now() - started,
					model: null, sessionId: "", wordCount: wordCount(t), segments: [], partial: t,
				},
			});
		});

		const outPath = params.out ? resolve(ctx.cwd, params.out.replace(/^@+/, "")) : undefined;
		if (outPath) { await mkdir(dirname(outPath), { recursive: true }); await writeFile(outPath, `${result.text}\n`, "utf-8"); }

		const details: AsrDetails = {
			path: rel(absPath, ctx.cwd),
			text: result.text,
			seconds: result.seconds,
			wallMs,
			model: result.model,
			sessionId: result.sessionId,
			wordCount: wordCount(result.text),
			segments: result.segments,
			...(params.words ? { words: result.words } : {}),
			...(outPath ? { outPath: rel(outPath, ctx.cwd) } : {}),
		};

		const summary = `${result.text || "(empty transcript)"}\n\n— ${details.wordCount} words, ${result.seconds}s audio, ${wallMs}ms wall, model=${result.model ?? "?"}${outPath ? `, wrote ${details.outPath}` : ""}`;
		const timings = params.words && result.words.length
			? `\n\nword timings:\n${result.words.map(w => `  ${(+w.start_time).toFixed(2).padStart(7)} → ${(+w.end_time).toFixed(2).padStart(7)}  ${w.word}${w.confidence != null ? ` (${(+w.confidence).toFixed(2)})` : ""}`).join("\n")}`
			: "";
		return { content: [{ type: "text", text: summary + timings }], details };
	}

	/* ---------------------------------------------------------------- tools -- */
	pi.registerTool({
		name: "asr_transcribe",
		label: "ASR transcribe",
		description: "Speech-to-text a local audio or video file via the NVIDIA Speech NIM realtime streaming endpoint (Parakeet/Nemotron streaming ASR). Accepts wav, mp3, m4a, aac, ogg, opus, flac, mp4, mov, mkv - anything ffmpeg can decode; non-16kHz/mono input is resampled automatically. Returns the final transcript, optionally with per-word timings.",
		promptSnippet: "Transcribe an audio or video file to text (speech-to-text) with the streaming ASR server",
		promptGuidelines: [
			"Use asr_transcribe for any speech-to-text request (voice notes, meeting recordings, podcast clips, screen recordings) instead of hand-rolling ffmpeg/python.",
			"asr_transcribe is real-time bound: budget roughly the duration of the audio. Pass speed=2 to halve wall time when slight accuracy loss is fine, or start=N to skip leading silence.",
			"If asr_transcribe returns an empty transcript, run asr_health to confirm the server is up and the language/model are right.",
		],
		parameters: Type.Object({
			path: Type.String({ description: "Audio/video file, relative to cwd or absolute (a leading @ is stripped)" }),
			language: Type.Optional(Type.String({ description: "BCP-47 language tag, default en-US" })),
			model: Type.Optional(Type.String({ description: "Override the ASR model name (default: whatever the server advertises)" })),
			words: Type.Optional(Type.Boolean({ description: "Include per-word timings + confidence (default false)" })),
			punctuation: Type.Optional(Type.Boolean({ description: "Server-side automatic punctuation (default true)" })),
			start: Type.Optional(Type.Number({ description: "Skip the first N seconds" })),
			speed: Type.Optional(Type.Number({ description: "Playback rate 0.5..4 (ffmpeg atempo); >1 is faster, slightly less accurate" })),
			pace: Type.Optional(Type.Boolean({ description: "Stream at 1x realtime pacing (default true)" })),
			out: Type.Optional(Type.String({ description: "Also write the plain transcript to this path" })),
		}),

		async execute(_id, params: TranscribeParams, signal, onUpdate, ctx): Promise<AgentToolResult<AsrDetails>> {
			return runTranscribe(params, ctx, onUpdate as any, signal);
		},

		renderCall(args: any, theme: any) {
			const bits = [
				theme.fg("toolTitle", theme.bold("asr_transcribe ")),
				theme.fg("muted", String(args?.path ?? "")),
			];
			if (args?.language && args.language !== "en-US") bits.push(theme.fg("dim", ` ${args.language}`));
			if (args?.speed && args.speed !== 1) bits.push(theme.fg("dim", ` ${args.speed}x`));
			if (args?.words) bits.push(theme.fg("dim", " +words"));
			if (args?.out) bits.push(theme.fg("dim", ` → ${args.out}`));
			return new Text(bits.join(""), 0, 0);
		},

		renderResult(result: any, options: any, theme: any) {
			const d: AsrDetails | undefined = result?.details;
			if (!d) return new Text("", 0, 0);
			if (options?.isPartial && d.partial) return new Text(theme.fg("warning", `… ${d.partial}`), 0, 0);
			const head = theme.fg("muted", `${d.path}  ${d.seconds}s audio → ${d.wallMs}ms  ${d.model ?? ""}`.trimEnd());
			const body = d.text
				? theme.fg("success", d.text)
				: theme.fg("warning", "(empty transcript — check --lang / model, or run asr_health)");
			let out = `${head}\n${body}`;
			if (options?.expanded && d.words?.length) {
				out += "\n" + d.words.slice(0, 50).map(w => theme.fg("dim", `  ${(+w.start_time).toFixed(2).padStart(7)} → ${(+w.end_time).toFixed(2).padStart(7)}  ${w.word}`)).join("\n");
			} else if (d.wordCount) {
				out += " " + theme.fg("dim", `(${d.wordCount} words)`);
			}
			if (options?.expanded && d.outPath) out += "\n" + theme.fg("dim", `  wrote ${d.outPath}`);
			return new Text(out, 0, 0);
		},
	});

	pi.registerTool({
		name: "asr_health",
		label: "ASR health",
		description:
			"Check the ASR server before/while debugging: GET health/live + health/ready, and mint a realtime session to " +
			"report the model, language and sample rate the server will actually use. Also verifies ffmpeg availability " +
			"and reports the API key being used (masked).",
		promptSnippet: "Check ASR server health and the currently advertised model/sample-rate config",
		parameters: Type.Object({}),

		async execute(_id, _params, _signal, _onUpdate, ctx): Promise<AgentToolResult<any>> {
			const client = new RealtimeAsrClient();
			const [live, ready, minted, ffmpeg] = await Promise.all([
				client.health("live", 10000).catch((e: any) => ({ error: e?.message ?? String(e) })),
				client.health("ready", 10000).catch((e: any) => ({ error: e?.message ?? String(e) })),
				client.mint(10000).catch((e: any) => ({ error: e?.message ?? String(e) })),
				hasFfmpeg(),
			]);
			const ok = live?.status === "live" && ready?.status === "ready" && !minted?.error;
			const key = client.cfg.apiKey;
			const secs = minted?.client_secret ? Math.max(0, (minted.client_secret.expires_at ?? 0) - Math.floor(Date.now() / 1000)) : 0;
			const text = [
				`baseUrl    ${process.env.ASR_NIM_URL ?? DEFAULT_BASE_URL}`,
				`apiKey     ${key ? `${key.slice(0, 5)}…${key.slice(-3)} (${key.length} chars)` : "NOT SET"}`,
				`health/live ${JSON.stringify(live)}`,
				`health/ready ${JSON.stringify(ready)}`,
				minted?.error
					? `mint        ERROR ${minted.error}`
					: [
						`mint        ${minted.id}`,
						`  model     ${minted.input_audio_transcription?.model}`,
						`  language  ${minted.input_audio_transcription?.language}`,
						`  audio     ${minted.input_audio_format}, ${minted.input_audio_params?.sample_rate_hz}Hz x${minted.input_audio_params?.num_channels}`,
						`  ws token  single-use, expires in ${secs}s (mint immediately before each connect)`,
					].join("\n"),
				`ffmpeg     ${ffmpeg ? "available (any container/codec accepted)" : "MISSING — only 16kHz mono wav files will work"}`,
				`note       the realtime decoder runs at ~1x realtime; a 10min file takes ~10min (speed=2 halves it)`,
			].join("\n");
			if (!ok) throw new Error(`ASR server is not usable:\n${text}`);
			return { content: [{ type: "text", text }], details: { ok, live, ready, minted, ffmpeg } };
		},

		renderResult(result: any, _options: any, theme: any) {
			const d: any = result?.details;
			if (!d) return new Text("", 0, 0);
			const m = d.minted?.input_audio_transcription;
			const p = d.minted?.input_audio_params;
			return new Text([
				d.ok ? theme.fg("success", "✓ ASR server live") : theme.fg("error", "✗ ASR server not usable"),
				theme.fg("muted", `  model ${m?.model ?? d.minted?.error ?? "?"}`),
				theme.fg("dim", `  ${p?.sample_rate_hz ?? "?"}Hz x${p?.num_channels ?? "?"} ${d.minted?.input_audio_format ?? ""} · ${m?.language ?? "?"} · ffmpeg ${d.ffmpeg ? "ok" : "MISSING"}`),
			].join("\n"), 0, 0);
		},
	});

	/* --------------------------------------------------------------- /asr --- */
	pi.registerCommand("asr", {
		description: "Transcribe an audio/video file: /asr <file> [--lang xx-YY] [--words] [--speed N] [--start S] [--out file.txt]",
		handler: async (args: string, ctx: ExtensionContext) => {
			const toks = args.trim().split(/\s+/).filter(Boolean);
			if (!toks.length) {
				ctx.ui.notify("usage: /asr <file> [--lang xx-YY] [--words] [--speed N] [--start S] [--out file.txt]", "info");
				return;
			}
			const p: TranscribeParams = { path: "" };
			for (let i = 0; i < toks.length; i++) {
				const t = toks[i];
				if (t === "--lang") p.language = toks[++i];
				else if (t === "--words") p.words = true;
				else if (t === "--speed") p.speed = Number.parseFloat(toks[++i]) || 1;
				else if (t === "--start") p.start = Number.parseFloat(toks[++i]) || 0;
				else if (t === "--out") p.out = toks[++i];
				else if (!p.path) p.path = t;
			}
			if (!p.path) { ctx.ui.notify("no file given", "error"); return; }

			let last = "";
			ctx.ui.setStatus("asr", `transcribing ${p.path} …`);
			try {
				const r = await runTranscribe(p, ctx);
				const d = r.details;
				last = d.text || "(empty transcript)";
				ctx.ui.notify(`${d.path}: ${d.wordCount} words, ${d.seconds}s audio in ${d.wallMs}ms`, "info");
				if (ctx.hasUI) {
					await ctx.ui.custom<void>((_tui, theme, _kb, done) => {
						const body = new Text([
							theme.fg("toolTitle", theme.bold(` transcribed ${d.path} `)),
							theme.fg("muted", `${d.seconds}s audio → ${d.wallMs}ms · ${d.model ?? "?"} · ${d.wordCount} words`),
							"",
							theme.fg(d.text ? "success" : "warning", d.text || "(empty transcript)"),
							...(d.words?.length ? [theme.fg("dim", `\n${d.words.slice(0, 12).map(w => `${(+w.start_time).toFixed(2).padStart(7)} → ${(+w.end_time).toFixed(2).padStart(7)}  ${w.word}`).join("\n")}`)] : []),
							"",
							theme.fg("dim", " enter/esc to close"),
						].join("\n"), 1, 1);
						body.onKey = (key: string) => { if (key === "return" || key === "escape" || key === " ") { done(undefined); return true; } return true; };
						return body;
					});
				} else {
					console.log(last);
				}
			} catch (e: any) {
				ctx.ui.notify(`ASR failed: ${e?.message ?? e}`, "error");
			} finally {
				ctx.ui.setStatus("asr", undefined);
			}
		},
	});
}